[Python] Allow reusing compute kernel state across repeated function calls via pc.bind_function() / FunctionExecutor (GH-49214)

python/pyarrow/_compute.pxd

@@ -38,6 +38,14 @@ cdef class FunctionOptions(_Weakrefable):
     cdef inline shared_ptr[CFunctionOptions] unwrap(self)
 
 
+cdef class FunctionExecutor(_Weakrefable):
+    cdef:
+        shared_ptr[CFunctionExecutor] sp_executor
+        CFunctionExecutor* base_executor
+
+    cdef void init(self, const shared_ptr[CFunctionExecutor]& sp_executor) except *
+
+
 cdef class _SortOptions(FunctionOptions):
     pass
 

python/pyarrow/_compute.pyx

@@ -611,6 +611,155 @@ def call_function(name, args, options=None, memory_pool=None, length=None):
                      length=length)
 
 
+def bind_function(name, args, options=None, memory_pool=None, registry=None):
+    """
+    Bind a function to its input types and options, returning a
+    reusable :class:`FunctionExecutor`.
+
+    The returned executor preconfigures the function's kernel (including any
+    one-time state such as hash tables for `is_in` / `index_in`) and can be
+    called repeatedly on data of the same type without the overhead of
+    repeated bind & kernel initialization.
+
+    This is useful when the same function with the same options needs to be
+    applied to many different input arrays, e.g. in a processing loop.
+
+    Parameters
+    ----------
+    name : str
+        The name of the function to bind (e.g. "is_in", "index_in", ...).
+    args : iterable
+        A prototype of the arguments to this function.  Accepted types are
+        the same as for :func:`call_function`: Array, ChunkedArray, Scalar,
+        RecordBatch, Table, or Python scalars.  Only the types matter for
+        binding — values are not copied into the executor except where the
+        options capture a `value_set`.
+    options : FunctionOptions, optional
+        Options instance for the bound function.  If the options carry a
+        `value_set` (as is the case for :class:`SetLookupOptions`), its
+        hash table is built once during binding and reused for every call
+        to :meth:`FunctionExecutor.execute`.
+    memory_pool : pyarrow.MemoryPool, optional
+        Memory pool used during binding (notably for allocating any
+        precomputed state such as a hash table).
+    registry : FunctionRegistry, optional
+        The registry to look up the function in.  If not passed, the global
+        registry (as returned by :func:`function_registry`) is used.
+
+    Returns
+    -------
+    FunctionExecutor
+        A reusable executor.  Call :meth:`FunctionExecutor.execute` on it to
+        run the bound computation on fresh data.
+
+    Examples
+    --------
+    >>> import pyarrow as pa
+    >>> import pyarrow.compute as pc
+    >>> value_set = pa.array([1, 2, 3])
+    >>> executor = pc.bind_function("is_in", [pa.array([1])],
+    ...                               options=pc.SetLookupOptions(value_set))
+    >>> for arr in [pa.array([1, 2, 4]), pa.array([2, 3, 5])]:
+    ...     print(executor.execute([arr]))
+    ...
+    [ True  True False]
+    [ True  True False]
+    """
+    cdef:
+        vector[CDatum] c_args
+        const CFunctionOptions* c_options = NULL
+        CFunctionRegistry* c_registry = NULL
+        CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
+        CExecContext c_exec_ctx = CExecContext(pool)
+        shared_ptr[CFunctionExecutor] sp_executor
+        c_string c_name = tobytes(name)
+        FunctionOptions c_options_obj
+        FunctionRegistry c_registry_obj
+
+    _pack_compute_args(args, &c_args)
+
+    if options is not None:
+        c_options_obj = <FunctionOptions> options
+        c_options = c_options_obj.get_options()
+
+    if registry is not None:
+        c_registry_obj = <FunctionRegistry> registry
+        c_registry = c_registry_obj.registry
+
+    with nogil:
+        sp_executor = GetResultValue(
+            GetFunctionExecutor(c_name, c_args, c_options, c_registry)
+        )
+
+    check_status(sp_executor.get().Init(c_options, &c_exec_ctx))
+
+    cdef FunctionExecutor executor = FunctionExecutor.__new__(FunctionExecutor)
+    executor.init(sp_executor)
+    return executor
+
+
+cdef class FunctionExecutor(_Weakrefable):
+    """
+    A pre-bound compute function executor.
+
+    Obtain an instance using :func:`bind_function`.  An executor holds
+    precomputed state (for example, a hash table over the ``value_set`` for
+    ``is_in`` / ``index_in``) and can be run many times on new input data
+    of the same shape and types.
+
+    Calling conventions are the same as :meth:`Function.call`.
+    """
+
+    def __init__(self):
+        raise TypeError(
+            "Do not call FunctionExecutor's constructor directly. "
+            "Use pyarrow.compute.bind_function() instead."
+        )
+
+    cdef void init(self, const shared_ptr[CFunctionExecutor]& sp_executor) except *:
+        self.sp_executor = sp_executor
+        self.base_executor = sp_executor.get()
+
+    def execute(self, args, memory_pool=None, length=None):
+        """
+        Execute the bound function on new arguments.
+
+        Parameters
+        ----------
+        args : iterable
+            Arguments to pass to the function.  Types must match the ones
+            used when :func:`bind_function` was called.
+        memory_pool : pyarrow.MemoryPool, optional
+            Memory pool to use for allocations during execution.  If not
+            passed, allocates from the default memory pool.
+        length : int, optional
+            Batch size for execution, for nullary functions.  If not passed,
+            inferred from the data.
+
+        Returns
+        -------
+        Any
+            Result of the bound function, wrapped as a PyArrow / Python
+            object — same types as :func:`call_function`.
+        """
+        cdef:
+            vector[CDatum] c_args
+            CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
+            CExecContext c_exec_ctx = CExecContext(pool)
+            CDatum result
+
+        _pack_compute_args(args, &c_args)
+
+        cdef int64_t c_length = length if length is not None else -1
+
+        with nogil:
+            result = GetResultValue(
+                self.base_executor.Execute(c_args, c_length)
+            )
+
+        return wrap_datum(result)
+
+
 cdef class FunctionOptions(_Weakrefable):
     __slots__ = ()  # avoid mistakingly creating attributes
 

python/pyarrow/includes/libarrow.pxd

@@ -2900,6 +2900,16 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
 
     cdef c_string ToString(DatumType kind)
 
+    cdef cppclass CFunctionExecutor "arrow::compute::FunctionExecutor":
+        CStatus Init(const CFunctionOptions* options,
+                     CExecContext* exec_ctx)
+        CResult[CDatum] Execute(const vector[CDatum]& args,
+                                 int64_t length)
+
+    CResult[shared_ptr[CFunctionExecutor]] GetFunctionExecutor(
+        const c_string& func_name, const vector[CDatum]& args,
+        const CFunctionOptions* options, CFunctionRegistry* func_registry)
+
 
 cdef extern from * namespace "arrow::compute":
     # inlined from compute/function_internal.h to avoid exposing

python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py

@@ -0,0 +1,81 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+"""Benchmark: bind_function + FunctionExecutor vs repeated is_in calls.
+
+Reproduces the scenario from issue #49214: a large value_set queried
+against many small input arrays, as happens in Ray Data checkpointing
+(100M IDs × 20 blocks).
+"""
+
+import time
+import pyarrow as pa
+import pyarrow.compute as pc
+import numpy as np
+
+
+def benchmark_repeated_is_in(value_set_size, num_arrays, array_size, seed=0):
+    """Baseline: call pc.is_in() repeatedly, rebuilding MemoTable each time."""
+    rng = np.random.default_rng(seed)
+    value_set = pa.array(rng.choice(1_000_000, size=value_set_size).tolist())
+    arrays = [pa.array(rng.choice(1_000_000, size=array_size).tolist())
+              for _ in range(num_arrays)]
+
+    for arr in arrays[:2]:
+        pc.is_in(arr, value_set=value_set)
+
+    start = time.perf_counter()
+    for arr in arrays:
+        pc.is_in(arr, value_set=value_set)
+    return time.perf_counter() - start
+
+
+def benchmark_executor(value_set_size, num_arrays, array_size, seed=0):
+    """New path: bind once, execute many - MemoTable built once."""
+    rng = np.random.default_rng(seed)
+    value_set = pa.array(rng.choice(1_000_000, size=value_set_size).tolist())
+    arrays = [pa.array(rng.choice(1_000_000, size=array_size).tolist())
+              for _ in range(num_arrays)]
+
+    executor = pc.bind_function(
+        "is_in", [arrays[0]],
+        options=pc.SetLookupOptions(value_set))
+
+    for arr in arrays[:2]:
+        executor.execute([arr])
+
+    start = time.perf_counter()
+    for arr in arrays:
+        executor.execute([arr])
+    return time.perf_counter() - start
+
+
+def main():
+    scenarios = [
+        (10_000, 50, 1_000,  "small"),
+        (100_000, 200, 5_000, "medium"),
+        (500_000, 500, 10_000, "large"),
+    ]
+    print(f"{'scenario':>10}  {'repeated':>12}  {'executor':>12}  {'speedup':>10}")
+    print("-" * 55)
+    for vs_size, n_arr, arr_size, label in scenarios:
+        baseline = benchmark_repeated_is_in(vs_size, n_arr, arr_size)
+        try:
+            improved = benchmark_executor(vs_size, n_arr, arr_size)
+            speedup = baseline / improved if improved > 0 else float("inf")
+            print(f"{label:>10}  {baseline*1000:>10.1f} ms  "
+                  f"{improved*1000:>10.1f} ms  {speedup:>8.2f}x")
+        except AttributeError:
+            print(f"{label:>10}  {baseline*1000:>10.1f} ms  "
+                  f"{'N/A':>12}  {'(bind_function not yet built)':>10}")
+
+
+if __name__ == "__main__":
+    main()

python/pyarrow/tests/test_compute.py

@@ -4518,3 +4518,128 @@ def test_winsorize():
     result = pc.winsorize(
         arr, options=pc.WinsorizeOptions(lower_limit=0.1, upper_limit=0.8))
     assert result.to_pylist() == [8, 4, 8, 8, 5, 3, 7, 2, 2, 6]
+
+
+
+# ---------------------------------------------------------------------------
+# Tests for bind_function / FunctionExecutor  (GH-49214)
+# ---------------------------------------------------------------------------
+
+def test_bind_function_basic():
+    """bind_function returns a FunctionExecutor that can execute."""
+    value_set = pa.array([1, 2, 3, 5])
+    arr = pa.array([1, 2, 4, 5, 6])
+
+    executor = pc.bind_function(
+        "is_in", [arr], options=pc.SetLookupOptions(value_set))
+    assert isinstance(executor, pc.FunctionExecutor)
+
+    result = executor.execute([arr])
+    expected = pc.is_in(arr, value_set=value_set)
+    assert result.equals(expected)
+
+
+def test_bind_function_reuse_is_in():
+    """Same value_set, multiple input arrays — executor reuses MemoTable."""
+    value_set = pa.array(range(1000))
+    arrays = [pa.array(range(i, i + 500)) for i in range(10)]
+
+    executor = pc.bind_function(
+        "is_in", [arrays[0]],
+        options=pc.SetLookupOptions(value_set))
+
+    for arr in arrays:
+        result = executor.execute([arr])
+        expected = pc.is_in(arr, value_set=value_set)
+        assert result.equals(expected)
+
+
+def test_bind_function_reuse_index_in():
+    """Same value_set, multiple input arrays for index_in."""
+    value_set = pa.array(["a", "b", "c"])
+    arrays = [pa.array(["a", "b", "d"]), pa.array(["c", "a", "e"])]
+
+    executor = pc.bind_function(
+        "index_in", [arrays[0]],
+        options=pc.SetLookupOptions(value_set))
+
+    for arr in arrays:
+        result = executor.execute([arr])
+        expected = pc.index_in(arr, value_set=value_set)
+        assert result.equals(expected)
+
+
+def test_bind_function_no_options():
+    """bind_function works for functions without options (e.g. add)."""
+    a = pa.array([1, 2, 3])
+    b = pa.array([4, 5, 6])
+
+    executor = pc.bind_function("add", [a, b])
+    result = executor.execute([a, b])
+    expected = pc.add(a, b)
+    assert result.equals(expected)
+
+
+def test_bind_function_different_options():
+    """Different options produce different executors."""
+    vs1 = pa.array([1, 2])
+    vs2 = pa.array([2, 3])
+    arr = pa.array([1, 2, 3])
+
+    exec1 = pc.bind_function(
+        "is_in", [arr], options=pc.SetLookupOptions(vs1))
+    exec2 = pc.bind_function(
+        "is_in", [arr], options=pc.SetLookupOptions(vs2))
+
+    assert exec1.execute([arr]).to_pylist() == [True, True, False]
+    assert exec2.execute([arr]).to_pylist() == [False, True, True]
+
+
+def test_bind_function_no_direct_construction():
+    """FunctionExecutor cannot be constructed directly."""
+    with pytest.raises(TypeError, match="Do not call FunctionExecutor"):
+        pc.FunctionExecutor()
+
+
+def test_bind_function_with_custom_registry():
+    """bind_function accepts a custom FunctionRegistry."""
+    registry = pc.function_registry()
+    value_set = pa.array([1, 2, 3])
+    arr = pa.array([1, 2])
+
+    executor = pc.bind_function(
+        "is_in", [arr], options=pc.SetLookupOptions(value_set),
+        registry=registry)
+    result = executor.execute([arr])
+    assert result.to_pylist() == [True, True]
+
+
+def test_bind_function_wrong_arg_type():
+    """Arguments whose types don't match the bound types raise on execute."""
+    value_set = pa.array([1, 2, 3])
+    int_arr = pa.array([1, 2])
+    str_arr = pa.array(["a", "b"])
+
+    executor = pc.bind_function(
+        "is_in", [int_arr], options=pc.SetLookupOptions(value_set))
+
+    # Different types — should raise
+    with pytest.raises((pa.ArrowInvalid, pa.ArrowNotImplementedError)):
+        executor.execute([str_arr])
+
+
+def test_bind_function_chunks():
+    """bind_function + executor work on ChunkedArrays too."""
+    value_set = pa.array(range(100))
+    chunked = pa.chunked_array([pa.array(range(0, 50)),
+                                pa.array(range(50, 100))])
+
+    executor = pc.bind_function(
+        "is_in", [chunked],
+        options=pc.SetLookupOptions(value_set))
+
+    # Use a fresh ChunkedArray of the same type as input
+    chunked2 = pa.chunked_array([pa.array([1, 2, 3]), pa.array([4, 5])])
+    result = executor.execute([chunked2])
+    expected = pc.is_in(chunked2, value_set=value_set)
+    assert result.equals(expected)