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

[andrewthad]

Describe the usage question you have. Please include as many useful details as possible.

I'm trying to use pyarrow to query arrow files that have a column named tags of type list<string>. The most common query is "does the list contain this string?" but I'm not able to figure out how to write this query. As an example:

>>> table = p.Table.from_arrays([p.array([["foo"], [], ["foo","bar"], ["bar"], ["bar"]], p.list_(p.string()))], names=['tags'])
>>> table
pyarrow.Table
tags: list<item: string>
  child 0, item: string
----
tags: [[["foo"],[],["foo","bar"],["bar"],["bar"]]]

From here, I would expect that I should be able to do something like table.filter(pc.list_has_element(pc.field('tags')), 'foo') but there is not a function list_has_element, and there doesn't seem to be anything like it. I'm wonderful if there is something that I'm missing, some part of the API that is used to lift other compute functions over lists.

Related Issues:

I see in #45167 that it's not currently possible to check that lists are equal, so maybe lists are a less developed part of the library.

Component(s)

Python

[amoeba]

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"]))

[andrewthad]

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.

[andrewthad]

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.

[jonasdedden]

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.