[Python] How to Test Lists for Existence of an Element #47118

Activity

1. andrewthad

   added

   Type: usageIssue is a user questionIssue is a user question

   on Jul 15, 2025
2. github-actions

   added

   Component: Python

   on Jul 15, 2025
3. amoeba

   changed the title [-]How to Test Lists for Existence of an Element[/-] [+][Python] How to Test Lists for Existence of an Element[/+] on Jul 15, 2025
4. amoeba commented on Jul 15, 2025

   

   amoeba

   on Jul 15, 2025

   Member

   More actions

   Hi @andrewthad, thanks for the question.

   I think #33295 and #32045 are related and that this operation isn't supported.

   Until that's made to work (PRs welcome), you could work around it by flattening first, though it's not efficient,

   pc.is_in(pc.list_flatten(tbl.column("tags")), value_set=pa.array(["foo"]))
5. andrewthad commented on Jul 16, 2025

   

   andrewthad

   on Jul 16, 2025

   Author

   More actions

   Thanks. I had not come across #33295 in my searching. I'll close this issue since the question was answered by your response, and #33295 is the right place to track any work that gets done on this.

   👍React with 👍1Reacted by Bryce Mecum
6. andrewthad commented on Jul 16, 2025

   

   andrewthad

   on Jul 16, 2025

   Author

   More actions

   I spent an hour trying to make headway on this but didn't get anywhere. If anyone tries to do this in the future, here is one difficulty to be aware of. The arrow library is currently able to express type signatures where the output type refers to something from an input type. For example, list_element has a type like List a -> IntLike -> a, and this type can be expressed because there is a function ListValuesType that can be used as the OutputType, and its meaning is "take the first input type, expect it to be a list, and then get the element type from it". But there is no equivalent for an input type. So it is not currently possible to express a type like List a -> a -> Bool, which is needed for list_contains.
7. andrewthad

   closed this as completedon Jul 16, 2025
8. jonasdedden commented on Sep 28, 2026

   

   jonasdedden

   on Sep 28, 2026

   Last edited by jonasdedden

   Contributor

   More actions

   I worked on an alternative to what @amoeba mentioned in this comment. It is up to an order of magnitude faster, and needs next to no extra memory. A native kernel is still faster again: list_contains is implemented in #51611 (GH-33295).

   import math from typing import Any
   import pyarrow as pa
   import pyarrow.compute as pc
   
   BLOCK_VALUES = 1 << 21
   
   def list_contains(array: pa.ChunkedArray, item: Any) -> pa.ChunkedArray:
      """Whether each list holds `item`, or a null element if `item` is None."""
      # Per chunk, as combining them can overflow 32-bit list offsets.
      return pa.chunked_array([_list_contains(arr, item) for arr in array.chunks], pa.bool_())
   
   def _list_contains(arr: pa.ListArray, item: Any) -> pa.BooleanArray:
      # Zero-copy slices of whole lists, of about 2M values each, keep the intermediates
      # small enough to stay in cache.
      n_lists = len(arr)
      n_values = arr.offsets[-1].as_py() - arr.offsets[0].as_py()
      step = max(1, n_lists * BLOCK_VALUES // max(1, n_values))
      blocks = [_list_contains_block(arr.slice(i, step), item) for i in range(0, n_lists, step)]
      return pa.concat_arrays(blocks) if blocks else pa.array([], pa.bool_())
   
   def _list_contains_block(arr: pa.ListArray, item: Any) -> pa.BooleanArray:
      offsets = arr.offsets     first = offsets[0].as_py()
      values = arr.values.slice(first, offsets[-1].as_py() - first)
      if first:
          # A plain `int` would widen int32 offsets to int64, which is slower to `take`.
          offsets = pc.subtract(offsets, pa.scalar(first, offsets.type))
      ends, starts = offsets.slice(1), offsets.slice(0, len(arr))
   
      # Like Polars, and unlike `pc.equal`, NaN matches NaN.
      if item is None:
          hits = pc.is_null(values)
      elif isinstance(item, float) and math.isnan(item):
          hits = pc.is_nan(values)
      else:
          hits = pc.equal(values, pa.scalar(item))
   
      # `running[i]` counts the matches before flattened position `i`, modulo 2^k. That is
      # exact within lists shorter than 2^k, so counting in the narrowest type fitting the
      # longest list keeps the only large intermediates small.
      longest = pc.max(pc.subtract(ends, starts)).as_py() or 0
      count_type = (
          pa.uint8() if longest < 1 << 8
          else pa.uint16() if longest < 1 << 16
          else pa.uint32() if longest < 1 << 32
          else pa.uint64()
      )
      hits = pa.concat_arrays([pa.array([False]), hits.fill_null(False)])
      running = pc.cumulative_sum(hits.cast(count_type))
      in_list = pc.subtract(running.take(ends), running.take(starts))
      found = pc.not_equal(in_list, pa.scalar(0, count_type))
      return found if arr.null_count == 0 else pc.if_else(arr.is_valid(), found, None)

   List length, how often the item matches	Comment approach	Previous cumsum (edit history of this comment)	This	Native list_contains (#51611)
   3 elements, rare (1e-5)	429 ms, 546 MB	219 ms, 613 MB	140 ms, 9 MB	51 ms, 10 MB
   3 elements, common (30%)	2,108 ms, 1,274 MB	217 ms, 613 MB	144 ms, 9 MB	51 ms, 10 MB
   1,000 elements, rare	202 ms, 413 MB	135 ms, 413 MB	167 ms, 9 MB	20 ms, 6 MB
   1,000 elements, common	520 ms, 535 MB	136 ms, 413 MB	151 ms, 9 MB	20 ms, 6 MB

   Setup

   • Data: 50M int64 values in one list<int64> chunk; the item is an int64.
   • Timing: best of 5 runs, a fresh process per cell. Memory is the peak above the input (Arrow pool plus NumPy).
   • Build: release build of GH-33295: [C++][Python] Add list_contains compute function #51611, 8 cores, system allocator. All approaches return identical results.
   • The comment approach needs a step to map is_in matches back to lists. I used list_parent_indices → filter → unique → is_in over the row numbers.
   • For 1,000-element lists the previous cumsum is slightly faster, but it needs ~45× the memory.

   👍React with 👍1Reacted by Bryce Mecum