jonasdedden commented on issue #47118: URL: https://github.com/apache/arrow/issues/47118#issuecomment-5875294308
I worked on an alternative to what @amoeba mentioned in [this](https://github.com/apache/arrow/issues/47118#issuecomment-3075893244) comment, which seems to be up to an order of magnitude faster in some cases, although vastly more complex. ```Python def list_contains(array: ChunkedArrayAny, item: NonNestedLiteral) -> ChunkedArrayAny: """Whether each list holds `item`, or a null element if `item` is None. A running count of matches over the flattened values grows within a list iff the list holds a match. That keeps this linear, where a group-by per list would sort. """ # Per chunk, as combining them can overflow 32-bit list offsets. chunks = [_list_contains(arr, item) for arr in array.chunks] return pa.chunked_array(chunks, pa.bool_()) def _list_contains(arr: ArrayAny, item: NonNestedLiteral) -> ArrayAny: # Each list's bounds in the flattened values. Null lists get null bounds, which # `take` below turns into a null result. lengths = pc.list_value_length(arr) ends = pc.cumulative_sum(lengths, skip_nulls=True) starts = pc.subtract(ends, lengths) # Like Polars, and unlike `pc.equal`, NaN matches NaN. values = pc.list_flatten(arr) 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, lit(item)) # `running[i]` counts the matches before flattened position `i`. Prepending the # leading 0 to the bit-packed `hits`, and counting in the width of the offsets, # keeps the large intermediates to a single one. hits = pa.concat_arrays([pa.array([False]), hits.fill_null(False)]) running = pc.cumulative_sum(hits.cast(lengths.type)) return pc.greater(running.take(ends), running.take(starts)) ``` | List length, how often the item matches | Ours | Comment approach (completed) | |---|---|---| | 3 elements, rare (1e-5) | **183 ms**, 625 MB extra | 380 ms, **413 MB** extra | | 3 elements, common (30%) | **197 ms**, **625 MB** extra | 1,519 ms, 944 MB extra | | 1,000 elements, rare | **114 ms**, 426 MB extra | 183 ms, **413 MB** extra | | 1,000 elements, common | **124 ms**, **426 MB** extra | 523 ms, 940 MB extra | -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
