[jonasdedden]

Rationale for this change

There is no direct way to check whether each list in a list array contains a value. This PR closes #33295.

What changes are included in this PR?

A binary scalar function list_contains(lists, value) returning a boolean per list, for list, large list, list view, large list view and fixed-size list inputs.

value can be a scalar, or an array of the same length, in which case each list is searched for the value at the same index.

• List array: each list is searched for its value directly, without materializing comparison bitmaps or repeated values. Short lists (fewer than 32 values on average) are compared in full without data-dependent branches; longer lists stop after the first 16-value chunk holding a match. With null child values, long lists are first searched without validity checks, and only lists found are searched again with them.
• Scalar list with an array of values: the values are looked up with is_in instead of broadcasting the list.
• Comparison semantics follow equal: values and list values are cast to the types equal would compare them as (so implicit casts apply, e.g. an int8 list and an int64 value), and equal's type errors are raised. float16 is compared as float32 and binary views as large binary.

Semantics:

• Null lists give null.
• Null list values don't match a non-null value.
• A null value matches lists that hold a null, as in Polars. Its type must still be comparable with the list value type, like a non-null value.
• A NaN value (including float16) matches NaN list values, as in is_in and Polars (equal says NaN != NaN).
• -0.0 matches 0.0, as with equal.

Are these changes tested?

Yes, in scalar_nested_test.cc and test_compute.py. Cases include sliced, chunked and out-of-order list-view inputs, arrays of values and scalar lists with arrays of values for 12 value types, short and long lists, null values whose slots hold the searched value, null and NaN values (including float16), -0.0, implicit casts, and type and time zone errors (also for null values).

Are there any user-facing changes?

Yes, a new compute function, pyarrow.compute.list_contains in Python.

Was AI used for this PR?

PR code and description written by:

• Human
• AI

Reviewed before submission by:

• Human
• AI
• Not reviewed

• GitHub Issue: [C++] Add a "list_contains" kernel #33295

[jonasdedden]

@pitrou @zanmato1984 and/or @AlenkaF

[jonasdedden]

972bc4c speeds up the per-list scan: short lists are now 2.4-4.1x faster, long lists are unchanged. The scan used to build an OptionalBinaryBitBlockCounter per list. Now it folds validity into the match bitmap once, reads a single unaligned word for lists spanning <= 64 bits, and looks up the offsets/sizes/width once per batch.

List length	Match rate	equal alone	list_contains before (95d6b87)	list_contains after (972bc4c)	Speedup
1	1e-5	18 ms	346 ms	108 ms	3.2×
3	1e-5	19 ms	140 ms	48 ms	2.9×
3	30%	19 ms	196 ms	48 ms	4.1×
10	30%	19 ms	68 ms	28 ms	2.4×
100	30%	18 ms	26 ms	26 ms	1.0×
1,000	1e-5	19 ms	21 ms	20 ms	1.0×
1,000	30%	18 ms	19 ms	19 ms	1.0×

Setup and script

50M int64 values in one list<int64> array, int64 item, best of 7, best of two runs per build (runs differed by ≤ 6 ms). Release build, 8 cores, system allocator; same PyArrow, only libarrow_compute swapped between the two commits. equal alone is the comparison over the same values, a lower bound for the kernel.

"""Time pc.list_contains against the `equal` pass it is built on. Run once per build."""

import json
import sys
import time

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc

N_VALUES = 50_000_000
CASES = [(1, 1e-5), (3, 1e-5), (3, 0.3), (10, 0.3), (100, 0.3), (1000, 1e-5), (1000, 0.3)]


def best_ms(func, repeats=7):
    times = []
    for _ in range(repeats):
        start = time.perf_counter()
        func()
        times.append(time.perf_counter() - start)
    return min(times) * 1e3


results = {}
for list_length, match_rate in CASES:
    rng = np.random.default_rng(42)
    values = rng.integers(2, 1 << 40, N_VALUES)
    values[rng.random(N_VALUES) < match_rate] = 1
    offsets = np.arange(0, N_VALUES + 1, list_length, dtype=np.int32)
    lists = pa.ListArray.from_arrays(offsets, pa.array(values))
    flat = lists.values
    item = pa.scalar(1, pa.int64())
    results[f"{list_length},{match_rate}"] = {
        "list_contains": best_ms(lambda: pc.list_contains(lists, item)),
        "equal": best_ms(lambda: pc.equal(flat, item)),
    }
json.dump(results, sys.stdout)

[jonasdedden]

@pitrou @zanmato1984 and/or @AlenkaF any remarks? :) I think this solution hopefully is quite performant and would close a 4 year old feature request.

[zanmato1984]

Thanks for working on this! I’d like to understand the design trade-offs for these three input forms:

1. List array + scalar query — e.g. [[2, 7, 8], [2, 3, 4]], searching for 2.
The current implementation compares all child values before checking each list’s match bitmap. Was this chosen to favor efficient batched comparison over a per-list search that could stop comparing once a match is found? Have you evaluated that trade-off for different list lengths and match positions?

2. List array + query array — e.g. [[1, 2], [3, 4]], with queries [2, 4].
Here, the query values are expanded using take before comparison. Is the extra allocation a deliberate trade-off to reuse equal and its type-promotion behavior, rather than introducing a dedicated kernel that searches each list using its corresponding query without materializing the repeated values?

3. Scalar list + query array — e.g. one list scalar [1, 2, 3, 4], queried with [2, 4, 9].
Have you considered a dedicated set-lookup path, potentially reusing is_in where the comparison semantics align, instead of broadcasting the list? It seems worth exploring particularly for a long list paired with many queries, although building the lookup table may not pay off for small inputs.

I haven’t benchmarked these alternatives; I’m mainly trying to understand whether they were evaluated already or would make sense as follow-up optimizations.

[jonasdedden]

@zanmato1984 thanks so much for the investigation! There indeed was a lot of performance optimization left to do, and I think now things are way faster and don't have extreme memory overhead in edge cases.

1. List array + scalar query: batched equal + bitmap scan was the simplest way to reuse equal's type promotion; the only early exit was in the bitmap scan. Each list is now searched directly. Short lists (with < 32 values) are compared in full without data-dependent branches, which keeps them stable when matches are frequent (early-exit loop lead to mispredicts in benchmarks). Long lists stop after the first 16-value chunk holding a match.

2. List array + query array: the take was also a way to reuse equal. The same per-list search now reads queries[i] directly, so the repeated values (and the indices to build them) are gone.

3. Scalar list + query array: now uses is_in instead of broadcasting. It aligns with equal once both sides are cast to equal's common type, with one difference: is_in tells -0.0 from 0.0, so floats get 0.0 added first.

Promotion still follows equal: values and list values are cast to the types equal would compare them as, and a length-1 equal call on nulls raises equal's exact type and time zone errors.

Some benchmarks with 50M int64 values in one list<int64> array:

List length	Match position	Scalar query: before -> after	Query array: before -> after
3	first	49 ms, 10 MB -> 34 ms, 4 MB	218 ms, 804 MB -> 34 ms, 4 MB
3	last	50 ms, 10 MB -> 33 ms, 4 MB	218 ms, 804 MB -> 35 ms, 4 MB
3	none	49 ms, 10 MB -> 34 ms, 4 MB	219 ms, 804 MB -> 35 ms, 4 MB
3	random (30%)	49 ms, 10 MB -> 36 ms, 4 MB	217 ms, 804 MB -> 33 ms, 4 MB
10	none	28 ms, 8 MB -> 19 ms, 1 MB	196 ms, 801 MB -> 21 ms, 1 MB
10	random (30%)	28 ms, 8 MB -> 20 ms, 1 MB	198 ms, 801 MB -> 21 ms, 1 MB
100	first	26 ms, 6 MB -> 9 ms, 0 MB	202 ms, 800 MB -> 9 ms, 0 MB
100	none	34 ms, 6 MB -> 19 ms, 0 MB	210 ms, 800 MB -> 20 ms, 0 MB
1,000	first	20 ms, 6 MB -> 1 ms, 0 MB	184 ms, 800 MB -> 1 ms, 0 MB
1,000	middle	20 ms, 6 MB -> 15 ms, 0 MB	183 ms, 800 MB -> 15 ms, 0 MB
1,000	none	21 ms, 6 MB -> 17 ms, 0 MB	186 ms, 800 MB -> 18 ms, 0 MB
1,000	random (30%)	20 ms, 6 MB -> 1 ms, 0 MB	184 ms, 800 MB -> 1 ms, 0 MB

With 10% null child values:

List length	Match position	Scalar query: before -> after	Query array: before -> after
3	none	53 ms -> 51 ms	219 ms -> 51 ms
3	random (30%)	52 ms -> 50 ms	219 ms -> 50 ms
1,000	none	22 ms -> 17 ms	185 ms -> 18 ms
1,000	random (30%)	20 ms -> 3 ms	184 ms -> 3 ms

Scalar list + query array (about half the queries match):

List length	Queries	Before (broadcast)	After (is_in)
10	1,000,000	156 ms, 245 MB	9.2 ms, 1 MB
1,000	1,000	4.1 ms, 24 MB	<0.1 ms, 0 MB
1,000	100,000	428 ms, 2,400 MB	1.4 ms, 0 MB
100,000	1,000	427 ms, 2,400 MB	0.9 ms, 6 MB

The lookup table paid off even at the smallest size I tried; broadcasting copies the list once per query, so its memory grows with list length times queries.

Benchmark scripts

List array, scalar query vs query array, by match position (python bench_positions.py):

import json
import subprocess
import sys
import time

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc

N_VALUES = 50_000_000


def make_input(length, where):
    rng = np.random.default_rng(42)
    values = rng.integers(2, 1 << 40, N_VALUES)
    starts = np.arange(0, N_VALUES // length * length, length)
    if where == "first":
        values[starts] = 1
    elif where == "middle":
        values[starts + length // 2] = 1
    elif where == "last":
        values[starts + length - 1] = 1
    elif where == "random":
        values[rng.random(N_VALUES) < 0.3] = 1
    offsets = np.arange(0, N_VALUES + 1, length, dtype=np.int32)
    return pa.ListArray.from_arrays(offsets, pa.array(values))


def run_one(length, where, mode):
    lists = make_input(length, where)
    query = (pa.scalar(1, pa.int64()) if mode == "scalar"
             else pa.array(np.full(len(lists), 1, dtype=np.int64)))
    pool = pa.default_memory_pool()
    baseline = pool.bytes_allocated()
    times = []
    for _ in range(5):
        start = time.perf_counter()
        pc.list_contains(lists, query)
        times.append(time.perf_counter() - start)
    return {"ms": min(times) * 1e3, "mb": (pool.max_memory() - baseline) / 1e6}


if __name__ == "__main__":
    if len(sys.argv) == 4:
        print(json.dumps(run_one(int(sys.argv[1]), sys.argv[2], sys.argv[3])))
        sys.exit()
    for length in (3, 10, 100, 1000):
        for where in ("first", "middle", "last", "none", "random"):
            cells = []
            for mode in ("scalar", "array"):
                out = json.loads(subprocess.run(
                    [sys.executable, __file__, str(length), where, mode],
                    capture_output=True, text=True, check=True).stdout)
                cells.append(f"{out['ms']:.0f} ms, {out['mb']:.0f} MB")
            print(f"| {length} | {where} | " + " | ".join(cells) + " |", flush=True)

With 10% null child values (uses make_input from above):

import time

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc

from bench_positions import make_input

for length in (3, 1000):
    for where in ("none", "random"):
        clean = make_input(length, where)
        rng = np.random.default_rng(7)
        values = clean.values.to_numpy()
        values = pa.array(values, mask=rng.random(len(values)) < 0.1)
        lists = pa.ListArray.from_arrays(clean.offsets, values)
        queries = pa.array(np.full(len(lists), 1, dtype=np.int64))
        cells = []
        for q in (pa.scalar(1, pa.int64()), queries):
            ts = []
            for _ in range(5):
                t = time.perf_counter()
                pc.list_contains(lists, q)
                ts.append(time.perf_counter() - t)
            cells.append(f"{min(ts) * 1e3:.0f} ms")
        print(f"| {length} | {where} | " + " | ".join(cells) + " |", flush=True)

Scalar list + query array:

import json
import subprocess
import sys
import time

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc

SIZES = [(10, 1_000_000), (1_000, 1_000), (1_000, 100_000), (100_000, 1_000)]


def run_one(list_length, n_queries):
    rng = np.random.default_rng(42)
    values = rng.permutation(2 * list_length)[:list_length].astype(np.int64)
    lst = pa.scalar(values, pa.list_(pa.int64()))
    queries = pa.array(rng.integers(0, 2 * list_length, n_queries))
    pool = pa.default_memory_pool()
    baseline = pool.bytes_allocated()
    times = []
    for _ in range(5):
        start = time.perf_counter()
        pc.list_contains(lst, queries)
        times.append(time.perf_counter() - start)
    return {"ms": min(times) * 1e3, "mb": (pool.max_memory() - baseline) / 1e6}


if __name__ == "__main__":
    if len(sys.argv) == 3:
        print(json.dumps(run_one(int(sys.argv[1]), int(sys.argv[2]))))
        sys.exit()
    for list_length, n_queries in SIZES:
        out = json.loads(subprocess.run(
            [sys.executable, __file__, str(list_length), str(n_queries)],
            capture_output=True, text=True, check=True).stdout)
        print(f"| {list_length:,} | {n_queries:,} | {out['ms']:,.1f} ms, {out['mb']:,.0f} MB |")

[dmitry-chirkov-dremio]

NaN matching differs between the two input forms. A scalar list containing double bits 0x7ff8000000000001, queried with 0x7ff8000000000002, returns false. The same list as an array returns true.

Please normalize NaNs before is_in and add a regression test with different NaN payloads.

[jonasdedden]

Fixed in 6fc4996: before is_in, NaNs in both the list and the queries are replaced by one canonical NaN (next to the existing -0.0 → 0.0 step). Added a regression test with payloads 0x7ff8000000000001 / 0x7ff8000000000002 for list arrays with scalar and array queries, and for a scalar list. Cost: +0.4–1.4 ms for 1M float queries.

[dmitry-chirkov-dremio]

fixed_size_binary(0) is valid, and its non-null values compare equal. This branch rejects it for list arrays, while the scalar-list path works.

Please handle zero-width binary values and add tests for both input forms.

[jonasdedden]

Fixed in 0110702: width 0 now matches every value (without memcmp, which could otherwise get null data pointers). Tests cover list arrays with scalar and array queries and a scalar list.

[dmitry-chirkov-dremio]

For list views, this includes the whole child array, so the later cast can fail on values outside every view. For example, child values [1, 9007199254740993], offset [0], size [1], and query float64(1) raise a conversion error, although the visible list is just [1].

Please avoid converting unused child values and add a regression test.

[jonasdedden]

Fixed in d9dd833. Without a cast, nothing changed (the child is searched in place, nothing is converted). With a cast, only the span from the first to the last referenced value is cast, as a zero-copy slice, so your example no longer converts the unused value. If converting the span fails, which can only happen in a gap between views, just the referenced values are gathered and cast. Regression tests cover the value outside every view and one in a gap, with scalar and array queries.

I measured the alternatives first. Flattening the views was up to ~100× slower for short out-of-order views (it concatenates one slice per view), and always gathering made casted views 2–4× slower with twice the memory. With this approach, casted views take the same time and memory as before (e.g. 131 ms / 404 MB for 16.7M 3-element views), and views without a cast got 3–13% faster.

[dmitry-chirkov-dremio]

How about same float16-to-float32 promotion here as in ComparableTypes?

Tried locally - for list containing [[1], [null]], float32 queries [null, null] fail, but [null, 1] succeeds. A float32 null scalar query also fails.

[zanmato1984]

I can reproduce this as well. Two additional cases are affected:

• A scalar list<float16> containing [1, null], with float32 queries [null, null], fails instead of returning [true, true].
• In the reverse direction, list<float32> containing [[1], [null]], with float16 queries [null, null], fails instead of returning [false, true].

Could the null-query path apply the same float16 normalization as the ordinary comparison path?

[jonasdedden]

Fixed in cb5b082: the null-query type check now goes through the same type resolution as other queries, including the float16 → float32 promotion. All the reported cases are covered by tests: list<float16> with float32 [null, null] and a float32 null scalar, a scalar list<float16> with float32 nulls, and list<float32> with float16 nulls.

[zanmato1984]

Thanks for the updates. Two type-validation inconsistencies still need correction: the typed-null float16 issue discussed in the existing thread, and the null scalar-list shortcut below. The shared-validator suggestion is non-blocking and can be a follow-up.

[zanmato1984]

Could we validate the declared list value type against the query type before this early return?

A null list<int32> scalar with string queries ["a", "b"] currently returns [null, null], whereas the equivalent list<int32> array [null, null] rejects the incompatible types. The same inconsistency occurs with typed-null string queries.

Type acceptance should not depend on whether the list scalar is valid.

[jonasdedden]

Fixed in b946a72: the types are now checked against the list type's declared value type before the null-list shortcut, so a null list<int32> scalar with string queries (null or not) raises like the equivalent array. Added tests for both.

[zanmato1984]

Non-blocking design question: could the execution-time comparison type checks, such as the timestamp timezone compatibility check, be extracted into an internal helper shared by the comparison kernels and list_contains, while retaining DispatchBest for type selection?

Currently, CheckComparable constructs two length-one null arrays and executes equal solely to obtain its validation status. A shared validator would avoid that dummy execution and make the validation contract explicit, without introducing a public API or changing equal’s supported types.

This could also be a follow-up if it would broaden the scope here.

[jonasdedden]

Done in 54532f4, since it was small: the time zone check moved into an internal ValidateComparison helper in codegen_internal, which both the comparison kernels (CompareTimestamps::Exec) and list_contains call. list_contains now resolves types with DispatchBest plus ValidateComparison, so the dummy equal execution is gone. Errors are unchanged (all comparison tests pass), and small calls got 1.3–3.3 µs cheaper. I left is_in's own time zone check alone, as it raises TypeError rather than Invalid.

[dmitry-chirkov-dremio]

What should this return when the list values and search value have different fixed-size binary widths?

import pyarrow as pa
import pyarrow.compute as pc


def test_list_contains_different_binary_widths():
    lists = pa.array([[b"a"]], type=pa.list_(pa.binary(1)))
    value = pa.scalar(b"ab", type=pa.binary(2))

    assert pc.equal(lists.values, value).to_pylist() == [False]
    assert pc.list_contains(lists, value).to_pylist() == [False]