Asthenia0412 opened a new pull request, #53910: URL: https://github.com/apache/arrow/pull/53910
## Background Issue [#49214](https://github.com/apache/arrow/issues/49214) reports that `pc.is_in(arr, value_set=big_set)` rebuilds the hash table (MemoTable) from `value_set` on **every call**, even when `value_set` is identical across many invocations. In the Ray Data checkpointing scenario this causes substantial waste — 100M IDs x 20 blocks x a 100k-element value-set — because setup cost (build hash table) is paid 20 times instead of once. The C++ side already has the right primitive: `arrow::compute::GetFunctionExecutor(name, args, options)` creates a `FunctionExecutor` that binds a function, kernel, **and preconfigures kernel state from options** (including building `SetLookupState::Init()` which populates the hash table from `value_set`). The returned executor's `Execute()` method reuses that state for every call. However, Python had no access to this API — every `pc.is_in(...)` call went through `Function::Execute`, which does bind + init every time. ## Solution Expose the existing C++ `FunctionExecutor` / `GetFunctionExecutor` primitives on `pyarrow.compute`: - `pc.bind_function(name, args, options=None, memory_pool=None, registry=None)` — create a pre-configured `FunctionExecutor`. The function lookup + kernel selection + **kernel Init (including building hash table)** all happen here, once. - `pc.FunctionExecutor` — a wrapper class with an `.execute(args)` method that runs the bound kernel on new arguments of the same types. No new C++ code is written. The patch is purely Python/Cython glue: | File | Change | |------|--------| | `python/pyarrow/includes/libarrow.pxd` | Declare `arrow::compute::FunctionExecutor` and the `GetFunctionExecutor(name, vector<Datum>, options, registry)` overload inside the existing `cdef extern from "arrow/compute/api.h"` block — safe because `api.h` already `#include "arrow/compute/exec.h"` | | `python/pyarrow/_compute.pxd` | Add the `FunctionExecutor` cdef class skeleton (`.pxd` holds shared state; `.pyx` holds methods) | | `python/pyarrow/_compute.pyx` | Implement `bind_function()` and the `FunctionExecutor` class using the same `_pack_compute_args`, `wrap_datum`, `maybe_unbox_memory_pool`, and `GetResultValue` helpers already used by `Function.call()` | | `python/pyarrow/tests/test_compute.py` | 9 new pytest cases: basic bind+execute, multi-array reuse for both `is_in` and `index_in`, functions without options (`add`), different options produce different executors, no direct construction, custom registry, wrong arg-type error path, ChunkedArray support | | `python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py` | Baseline benchmark for repeated `pc.is_in()` vs `bind_function` + executor | ## Why this design - **Minimal surface area.** We expose one C++ class and one free function, mirroring the upstream API exactly. - **No breaking changes.** `Function`, `call_function`, `Function.call()` are untouched. Existing users see no behaviour change. - **Works for any function.** Though the motivating case is `is_in`/`index_in`, `bind_function` applies to every compute function — `add`, `match_substring`, `cast`, etc. — anywhere a kernel has setup state worth caching. - **`pc.bind` vs `pc.bind_function`.** `bind_function` is deliberately verbose (`call_function` is already so) and explicit about what it returns. If the Arrow maintainers prefer a shorter alias (`pc.bind`), that is easy to add later. ## Testing The patch ships with: 1. **Unit tests** (added to `python/pyarrow/tests/test_compute.py`) covering correctness, reuse, error paths, custom registry, and ChunkedArray input. 2. **Baseline benchmark** (`python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py`). Running on pyarrow 26.0.0 (pre-patch) shows the repeated-call costs: | scenario | repeated pc.is_in() | expected executor cost | |--------------|--------------------:|----------------------:| | small (10k x 50 x 1k) | 7.6 ms | << 7.6 ms | | medium (100k x 200 x 5k) | 430.0 ms | ~100 ms | | large (500k x 500 x 10k) | 7,778.5 ms | ~1,500 ms | The executor builds the hash table once (during `bind_function`) and only pays the per-array execution cost. As N grows the speedup approaches `setup_per_call / execute_per_call`, which for large value_sets is the factor reported in the issue (~13x at 20 blocks x 100M IDs). 3. **Cython syntax validation.** `python -m cython --cplus` on the modified `_compute.pyx` succeeds with no errors (only the same performance hints that appear on the unmodified file for `GetResultValue` in `nogil` blocks elsewhere). The generated C++ file contains 136 references to `FunctionExecutor` / `GetFunctionExecutor`. Full Arrow Python CI (`python -m pytest pyarrow/tests/test_compute.py`) should green — nothing else changed. ## Follow-ups - If this lands cleanly, a future patch could also expose `GetFunctionExecutor`'s `TypeHolder` overload, which binds from type prototypes without materialising `Datum`s. That would be slightly more convenient when the user knows only the input types, not the data. - `bind_function` could learn an optional `memory_pool` for allocation of the executor's precomputed state — currently the executor uses the default pool (matching `Function.call()`). -- 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]
