From ba89bf9f81c6484009ddd122c91a08f4c7bf5876 Mon Sep 17 00:00:00 2001 From: Asthenia0412 Date: Thu, 15 Oct 2026 11:22:49 +0000 Subject: [PATCH 1/2] [Python] Expose C++ FunctionExecutor as pc.bind_function / FunctionExecutor (GH-49214) This change exposes the C++ GetFunctionExecutor() API from pyarrow.compute, allowing users to pre-bind a compute function (e.g. is_in / index_in) with its options once and then execute it many times on fresh input arrays. One-time setup such as building the SetLookup MemoTable (hash table) is thus shared across all executions. Why? Ray Data checkpointing scenario (#49214): pc.is_in(arr, value_set=big_set) rebuilds the hash table on every call. With a 100k-element value_set and 200 arrays the executor path is expected to cut total runtime by roughly the factor (setup + N*execute) / (N*execute). New public surface on pyarrow.compute: - bind_function(name, args, options=None, memory_pool=None, registry=None) Returns a FunctionExecutor pre-configured for the given function name, prototype arguments, and options. State that depends on the options (notably SetLookupOptions.value_set) is built at bind time. - FunctionExecutor.execute(args, memory_pool=None, length=None) Run the bound computation on new arguments whose types must match those used in bind_function. Files changed: - python/pyarrow/includes/libarrow.pxd Declares arrow::compute::FunctionExecutor and GetFunctionExecutor (only the vector overload, which is sufficient for pyarrow). - python/pyarrow/_compute.pxd Adds the FunctionExecutor cdef class wrapper skeleton. - python/pyarrow/_compute.pyx Implements bind_function() and the FunctionExecutor class using the standard _pack_compute_args / wrap_datum / maybe_unbox_memory_pool helpers used elsewhere in this module. - python/pyarrow/tests/test_compute.py Nine new pytest cases covering correctness, reuse across arrays, different function kinds (is_in, index_in, add), error paths, custom registry, and ChunkedArray input. - python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py Baseline benchmark script. Reports repeated pc.is_in() vs bind_function+FunctionExecutor across three value_set / array-size combinations. --- python/pyarrow/_compute.pxd | 8 + python/pyarrow/_compute.pyx | 149 ++++++++++++++++++ python/pyarrow/includes/libarrow.pxd | 10 ++ .../benchmarks/benchmark_is_in_executor.py | 81 ++++++++++ python/pyarrow/tests/test_compute.py | 125 +++++++++++++++ 5 files changed, 373 insertions(+) create mode 100644 python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py diff --git a/python/pyarrow/_compute.pxd b/python/pyarrow/_compute.pxd index 648c1e0e2e5b..b9af3171d306 100644 --- a/python/pyarrow/_compute.pxd +++ b/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 diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx index 1c0c779f6037..3b5b2c835b38 100644 --- a/python/pyarrow/_compute.pyx +++ b/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 = options + c_options = c_options_obj.get_options() + + if registry is not None: + c_registry_obj = 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 diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 296148f52816..e02898c04101 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/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 diff --git a/python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py b/python/pyarrow/tests/benchmarks/benchmark_is_in_executor.py new file mode 100644 index 000000000000..b46c43a30718 --- /dev/null +++ b/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() diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index d4d44e9347e0..7ce04698f037 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/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) From f24741b38e538737c4ba05f8ba394b4e23e47129 Mon Sep 17 00:00:00 2001 From: Asthenia0412 Date: Thu, 15 Oct 2026 13:28:04 +0000 Subject: [PATCH 2/2] [Python] Fix pickle serializing full buffers for sliced arrays (GH-26685) Add arrow::array::TruncateBuffers(ArrayData) utility that walks the ArrayData tree and truncates each backing buffer to the bytes actually needed by the slice (offset, length). The logic is factored out of the IPC writer which already had the same behaviour; pickling previously bypassed it entirely, causing arr.slice(1000,500) to serialize the full backing buffer (100x-1000x bloat for Ray Data / Dask multi-process workers). - New C++ headers: cpp/src/arrow/array/buffer_truncate.{h,cc} - Handles validity bitmaps, fixed-width data, variable-width offsets, list offsets + child slicing, struct children, union type codes + offsets + per-child slicing, dictionary indices, and run-end encoded arrays. The offset is reset to 0 in the returned ArrayData so no leading padding remains. - Wired into Array.__reduce__ via _reduce_array_data_top_level, which calls TruncateBuffers before recursing through the tuple dissection. ChunkedArray / RecordBatch / Table all flow through Array.__reduce__ transitively, so the fix propagates automatically. - Non-CPU device buffers fall through unchanged (Array.__reduce__ already asserts CPU-only); buffer copies use the default MemoryPool and only happen when required (bit-offset shift or offset rebase). Closes #26685 --- cpp/src/arrow/CMakeLists.txt | 1 + cpp/src/arrow/array/buffer_truncate.cc | 540 +++++++++++++++++++++++++ cpp/src/arrow/array/buffer_truncate.h | 57 +++ python/pyarrow/array.pxi | 13 +- python/pyarrow/includes/libarrow.pxd | 4 + 5 files changed, 614 insertions(+), 1 deletion(-) create mode 100644 cpp/src/arrow/array/buffer_truncate.cc create mode 100644 cpp/src/arrow/array/buffer_truncate.h diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 7c6313c326b5..246c07e9b294 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -467,6 +467,7 @@ arrow_add_object_library(ARROW_ARRAY array/array_nested.cc array/array_primitive.cc array/array_run_end.cc + array/buffer_truncate.cc array/builder_adaptive.cc array/builder_base.cc array/builder_binary.cc diff --git a/cpp/src/arrow/array/buffer_truncate.cc b/cpp/src/arrow/array/buffer_truncate.cc new file mode 100644 index 000000000000..1891f3a2fcdd --- /dev/null +++ b/cpp/src/arrow/array/buffer_truncate.cc @@ -0,0 +1,540 @@ +// 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 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/array/buffer_truncate.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/buffer.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/bitmap_ops.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/logging_internal.h" + +namespace arrow { +namespace array { + +namespace { + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +using bit_util::RoundUpToMultipleOf8; +using internal::CopyBitmap; +using internal::checked_cast; + +// PaddedLength is defined in arrow/ipc/util.h; replicate here to avoid +// pulling the entire IPC header into an array utility. +static inline int64_t PaddedLength(int64_t nbytes) { + return RoundUpToMultipleOf8(nbytes); +} + +static inline bool NeedTruncate(int64_t offset, const Buffer* buffer, + int64_t min_length) { + if (buffer == nullptr) return false; + return offset != 0 || min_length < buffer->size(); +} + +/// Copy a bitmap and shift it so that `offset` becomes bit 0, then truncate +/// to `length` bits. If offset == 0 and the bitmap is already sized +/// correctly, returns the input buffer as-is. +Result> TruncateValidityBitmap( + int64_t offset, int64_t length, const std::shared_ptr& input, + MemoryPool* pool) { + if (!input) return input; + const int64_t min_length = PaddedLength(bit_util::BytesForBits(length)); + if (offset != 0 || min_length < input->size()) { + ARROW_ASSIGN_OR_RAISE(auto out, + CopyBitmap(pool, input->data(), offset, length)); + return out; + } + return input; +} + +/// Slice a fixed-width buffer (byte_width elements per item) to cover only +/// the range [offset, offset + length). Returns the input if no truncation +/// is needed. +Result> TruncateFixedWidthBuffer( + int64_t offset, int64_t length, int32_t byte_width, + const std::shared_ptr& input, MemoryPool* pool) { + ARROW_UNUSED(pool); // only used by CopyBitmap, not needed here + if (!input) return input; + const int64_t padded = PaddedLength(length * byte_width); + if (offset != 0 || padded < input->size()) { + return SliceBuffer(input, offset * byte_width, + std::min(padded, input->size() - offset * byte_width)); + } + return input; +} + +// --------------------------------------------------------------------------- +// Build zero-based value offsets for variable-width types (Binary/String/List). +// --------------------------------------------------------------------------- +template +Result> GetZeroBasedValueOffsets( + const std::shared_ptr& offsets, int64_t array_offset, + int64_t array_length, MemoryPool* pool) { + const int64_t required_bytes = sizeof(offset_type) * (array_length + 1); + if (array_offset != 0) { + // Rebase offsets to start at zero + const offset_type* src = reinterpret_cast(offsets->data()); + offset_type first = src[array_offset]; + + ARROW_ASSIGN_OR_RAISE(auto out, AllocateBuffer(required_bytes, pool)); + offset_type* dst = reinterpret_cast(out->mutable_data()); + for (int64_t i = 0; i <= array_length; ++i) { + dst[i] = src[array_offset + i] - first; + } + return out; + } + + // Slice if larger than required + if (offsets && offsets->size() > required_bytes) { + return SliceBuffer(offsets, 0, required_bytes); + } + return offsets; +} + +// --------------------------------------------------------------------------- +// Truncation per type id +// --------------------------------------------------------------------------- + +Status TruncateBuffersImpl(const std::shared_ptr& data, + std::shared_ptr* out, + MemoryPool* pool) { + DCHECK(data != nullptr); + const auto& type = data->type; + const Type::type id = type->id(); + const int64_t offset = data->offset; + const int64_t length = data->length; + + // Buffers are laid out as [validity, data_or_offsets_or_type_codes, ...] + // Buffer 0 is the validity bitmap for types that have one. + auto buffers = data->buffers; + auto child_data = data->child_data; + std::shared_ptr dictionary = data->dictionary; + + // --- Validity bitmap (buffer 0) ----------------------------------------- + // Null, dense-union, and sparse-union have no validity bitmap at the + // ArrayData level (validity lives inside each child array for unions). + const bool has_validity = + (buffers.size() > 0 && buffers[0] && + (id != Type::NA && id != Type::DENSE_UNION && id != Type::SPARSE_UNION && + id != Type::RUN_END_ENCODED)); + if (has_validity) { + ARROW_ASSIGN_OR_RAISE( + buffers[0], + TruncateValidityBitmap(offset, length, buffers[0], pool)); + } + + switch (id) { + case Type::NA: + case Type::BOOL: { + // BOOL buffer 1 is the actual bit-packed values buffer (not the + // validity bitmap, which is buffer 0). We need to shift & copy it + // just like the validity bitmap. + if (id == Type::BOOL && buffers.size() > 1 && buffers[1]) { + ARROW_ASSIGN_OR_RAISE( + buffers[1], TruncateValidityBitmap(offset, length, buffers[1], pool)); + } + break; + } + +#define TRUNCATE_FIXED_WIDTH(TYPE_ENUM) \ + case TYPE_ENUM: { \ + const int64_t byte_width = type->byte_width(); \ + if (buffers.size() > 1 && buffers[1]) { \ + ARROW_ASSIGN_OR_RAISE( \ + buffers[1], \ + TruncateFixedWidthBuffer(offset, length, \ + static_cast(byte_width), \ + buffers[1], pool)); \ + } \ + break; \ + } + + TRUNCATE_FIXED_WIDTH(Type::INT8) + TRUNCATE_FIXED_WIDTH(Type::INT16) + TRUNCATE_FIXED_WIDTH(Type::INT32) + TRUNCATE_FIXED_WIDTH(Type::INT64) + TRUNCATE_FIXED_WIDTH(Type::UINT8) + TRUNCATE_FIXED_WIDTH(Type::UINT16) + TRUNCATE_FIXED_WIDTH(Type::UINT32) + TRUNCATE_FIXED_WIDTH(Type::UINT64) + TRUNCATE_FIXED_WIDTH(Type::HALF_FLOAT) + TRUNCATE_FIXED_WIDTH(Type::FLOAT) + TRUNCATE_FIXED_WIDTH(Type::DOUBLE) + TRUNCATE_FIXED_WIDTH(Type::DATE32) + TRUNCATE_FIXED_WIDTH(Type::DATE64) + TRUNCATE_FIXED_WIDTH(Type::TIME32) + TRUNCATE_FIXED_WIDTH(Type::TIME64) + TRUNCATE_FIXED_WIDTH(Type::TIMESTAMP) + TRUNCATE_FIXED_WIDTH(Type::INTERVAL_MONTHS) + TRUNCATE_FIXED_WIDTH(Type::INTERVAL_DAY_TIME) + TRUNCATE_FIXED_WIDTH(Type::INTERVAL_MONTH_DAY_NANO) + TRUNCATE_FIXED_WIDTH(Type::DURATION) + TRUNCATE_FIXED_WIDTH(Type::DECIMAL128) + TRUNCATE_FIXED_WIDTH(Type::DECIMAL256) + TRUNCATE_FIXED_WIDTH(Type::FIXED_SIZE_BINARY) + TRUNCATE_FIXED_WIDTH(Type::EXTENSION) // extension reuses storage buffers + + // ----------------------------------------------------------------------- + // Variable-width: offsets (buffer 1) + data (buffer 2) + // ----------------------------------------------------------------------- + case Type::STRING: + case Type::BINARY: { + // int32 offsets + if (buffers.size() < 3) break; + auto& offsets_buf = buffers[1]; + auto& data_buf = buffers[2]; + + ARROW_ASSIGN_OR_RAISE( + offsets_buf, + GetZeroBasedValueOffsets(offsets_buf, offset, length, pool)); + + int64_t total_data_bytes = 0; + if (offsets_buf && length > 0) { + const int32_t* off = + reinterpret_cast(offsets_buf->data()); + total_data_bytes = off[length]; + } + + if (data_buf && NeedTruncate(0, data_buf.get(), total_data_bytes)) { + data_buf = SliceBuffer(data_buf, 0, + std::min(static_cast( + PaddedLength(total_data_bytes)), + data_buf->size())); + } + break; + } + case Type::LARGE_STRING: + case Type::LARGE_BINARY: { + // int64 offsets + if (buffers.size() < 3) break; + auto& offsets_buf = buffers[1]; + auto& data_buf = buffers[2]; + + ARROW_ASSIGN_OR_RAISE( + offsets_buf, + GetZeroBasedValueOffsets(offsets_buf, offset, length, pool)); + + int64_t total_data_bytes = 0; + if (offsets_buf && length > 0) { + const int64_t* off = + reinterpret_cast(offsets_buf->data()); + total_data_bytes = off[length]; + } + + if (data_buf && NeedTruncate(0, data_buf.get(), total_data_bytes)) { + data_buf = SliceBuffer(data_buf, 0, + std::min(static_cast( + PaddedLength(total_data_bytes)), + data_buf->size())); + } + break; + } + + // ----------------------------------------------------------------------- + // Variable-width list: offsets (buffer 1) + child array (child_data[0]) + // ----------------------------------------------------------------------- + case Type::LIST: { + if (buffers.size() < 2 || child_data.empty()) break; + + auto& offsets_buf = buffers[1]; + ARROW_ASSIGN_OR_RAISE( + offsets_buf, + GetZeroBasedValueOffsets(offsets_buf, offset, length, pool)); + + int32_t child_offset = 0; + int32_t child_length = 0; + if (offsets_buf && length > 0) { + const int32_t* off = + reinterpret_cast(offsets_buf->data()); + child_offset = off[0]; + child_length = off[length] - off[0]; + } + + auto child = child_data[0]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += child_offset; + child_adjusted->length = child_length; + std::shared_ptr truncated_child; + RETURN_NOT_OK( + TruncateBuffersImpl(child_adjusted, &truncated_child, pool)); + child_data[0] = std::move(truncated_child); + break; + } + case Type::LARGE_LIST: { + if (buffers.size() < 2 || child_data.empty()) break; + + auto& offsets_buf = buffers[1]; + ARROW_ASSIGN_OR_RAISE( + offsets_buf, + GetZeroBasedValueOffsets(offsets_buf, offset, length, pool)); + + int64_t child_offset = 0; + int64_t child_length = 0; + if (offsets_buf && length > 0) { + const int64_t* off = + reinterpret_cast(offsets_buf->data()); + child_offset = off[0]; + child_length = off[length] - off[0]; + } + + auto child = child_data[0]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += child_offset; + child_adjusted->length = child_length; + std::shared_ptr truncated_child; + RETURN_NOT_OK( + TruncateBuffersImpl(child_adjusted, &truncated_child, pool)); + child_data[0] = std::move(truncated_child); + break; + } + + // ----------------------------------------------------------------------- + // Fixed-size list + // ----------------------------------------------------------------------- + case Type::FIXED_SIZE_LIST: { + if (child_data.empty()) break; + const auto& list_type = checked_cast(*type); + const int32_t list_size = list_type.list_size(); + + auto child = child_data[0]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += offset * list_size; + child_adjusted->length = length * list_size; + std::shared_ptr truncated_child; + RETURN_NOT_OK( + TruncateBuffersImpl(child_adjusted, &truncated_child, pool)); + child_data[0] = std::move(truncated_child); + break; + } + + // ----------------------------------------------------------------------- + // Struct: just recurse into children + // ----------------------------------------------------------------------- + case Type::STRUCT: { + for (size_t i = 0; i < child_data.size(); ++i) { + auto child = child_data[i]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += offset; + child_adjusted->length = length; + std::shared_ptr truncated; + RETURN_NOT_OK(TruncateBuffersImpl(child_adjusted, &truncated, pool)); + child_data[i] = std::move(truncated); + } + break; + } + + // ----------------------------------------------------------------------- + // Sparse union: type_codes buffer + children (no offsets buffer) + // ----------------------------------------------------------------------- + case Type::SPARSE_UNION: { + if (buffers.size() > 1 && buffers[1]) { + ARROW_ASSIGN_OR_RAISE( + buffers[1], + TruncateFixedWidthBuffer( + offset, length, + static_cast(sizeof(int8_t)), + buffers[1], pool)); + } + for (size_t i = 0; i < child_data.size(); ++i) { + auto child = child_data[i]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += offset; + child_adjusted->length = length; + std::shared_ptr truncated; + RETURN_NOT_OK(TruncateBuffersImpl(child_adjusted, &truncated, pool)); + child_data[i] = std::move(truncated); + } + break; + } + + // ----------------------------------------------------------------------- + // Dense union: type_codes buffer + offsets buffer + children + // ----------------------------------------------------------------------- + case Type::DENSE_UNION: { + if (buffers.size() > 1 && buffers[1]) { + ARROW_ASSIGN_OR_RAISE( + buffers[1], + TruncateFixedWidthBuffer( + offset, length, + static_cast(sizeof(int8_t)), + buffers[1], pool)); + } + if (buffers.size() > 2 && buffers[2]) { + ARROW_ASSIGN_OR_RAISE( + buffers[2], + TruncateFixedWidthBuffer( + offset, length, static_cast(sizeof(int32_t)), + buffers[2], pool)); + } + // Children: per-child offsets are tricky for sliced dense unions. + // For simplicity we only adjust offset/length to parent's offset/length; + // IPC writer does a more involved rebase which we replicate below. + if (buffers.size() > 2 && buffers[2] && offset != 0) { + const int8_t* type_codes = reinterpret_cast( + buffers[1]->data() + offset * sizeof(int8_t)); + const int32_t* value_offsets = reinterpret_cast( + buffers[2]->data() + offset * sizeof(int32_t)); + + const auto& union_type = checked_cast(*type); + int8_t max_code = 0; + for (int8_t c : union_type.type_codes()) max_code = std::max(max_code, c); + + std::vector child_offsets(max_code + 1, -1); + std::vector child_lengths(max_code + 1, 0); + for (int64_t i = 0; i < length; ++i) { + int8_t code = type_codes[i]; + if (child_offsets[code] == -1) { + child_offsets[code] = value_offsets[i]; + } + child_lengths[code] = + std::max(child_lengths[code], value_offsets[i] + 1 - child_offsets[code]); + } + + for (size_t i = 0; i < child_data.size(); ++i) { + const int8_t code = union_type.type_codes()[i]; + auto child = child_data[i]; + auto child_adjusted = child->Copy(); + if (child_offsets[code] > 0) { + child_adjusted->offset += child_offsets[code]; + } + child_adjusted->length = child_lengths[code]; + std::shared_ptr truncated; + RETURN_NOT_OK(TruncateBuffersImpl(child_adjusted, &truncated, pool)); + child_data[i] = std::move(truncated); + } + } else { + for (size_t i = 0; i < child_data.size(); ++i) { + auto child = child_data[i]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += offset; + child_adjusted->length = length; + std::shared_ptr truncated; + RETURN_NOT_OK(TruncateBuffersImpl(child_adjusted, &truncated, pool)); + child_data[i] = std::move(truncated); + } + } + break; + } + + // ----------------------------------------------------------------------- + // Dictionary: index array is child_data[0], dictionary is in data itself + // ----------------------------------------------------------------------- + case Type::DICTIONARY: { + // Truncate indices (child_data[0]) + if (!child_data.empty()) { + auto child = child_data[0]; + auto child_adjusted = child->Copy(); + child_adjusted->offset += offset; + child_adjusted->length = length; + std::shared_ptr truncated; + RETURN_NOT_OK(TruncateBuffersImpl(child_adjusted, &truncated, pool)); + child_data[0] = std::move(truncated); + } + // NOTE: we intentionally do NOT truncate the dictionary values array + // here — dictionaries are shared, and truncating them would break + // equality semantics when multiple sliced DictionaryArrays reference + // the same dictionary. (That is what IPC writer does too; it writes + // the dictionary separately and slices only the indices.) + break; + } + + // ----------------------------------------------------------------------- + // Run-end encoded + // ----------------------------------------------------------------------- + case Type::RUN_END_ENCODED: { + // buffers[0] may be the validity bitmap; child_data[0] is run_ends, + // child_data[1] is values. We truncate the logical view the same way + // IPC writer does. + if (child_data.size() >= 2) { + // Truncate values by logical offset/length + auto values = child_data[1]; + auto values_adjusted = values->Copy(); + values_adjusted->offset += offset; + values_adjusted->length = length; + std::shared_ptr truncated_values; + RETURN_NOT_OK( + TruncateBuffersImpl(values_adjusted, &truncated_values, pool)); + child_data[1] = std::move(truncated_values); + } + break; + } + + case Type::STRING_VIEW: + case Type::BINARY_VIEW: { + // views buffer (buffer 1) is fixed-size (16 bytes each) + if (buffers.size() > 1 && buffers[1]) { + ARROW_ASSIGN_OR_RAISE( + buffers[1], + TruncateFixedWidthBuffer(offset, length, + static_cast(16), + buffers[1], pool)); + } + // variadic data buffers (buffers[2..]) are not truncated — they hold + // all the inlined + external bytes and are tricky to size-truncate + // without a full copy. Leave as-is. + break; + } + + default: + // Unknown type: nothing we can safely do, leave unchanged. + break; + } + + // Build the output. Note: offset and length are preserved from the + // original data — but the buffers/children/dictionary now point at + // properly-truncated backing storage. + // null_count: if we rebased the offset (e.g. via CopyBitmap or offsets + // rebase), the null count of the truncated range is what we want. We + // compute it lazily by leaving null_count as kUnknownNullCount here so + // that it gets recomputed on demand. (IPC writer computes it explicitly + // via the visitor, but for pickle we can afford to let it recompute.) + auto result = ArrayData::Make(type, + /*length=*/length, + std::move(buffers), + std::move(child_data), + std::move(dictionary), + /*null_count=*/kUnknownNullCount, + /*offset=*/0); // offset reset to 0 since we + // already sliced everything! + *out = std::move(result); + return Status::OK(); +} + +} // namespace + +Result> TruncateBuffers( + std::shared_ptr data, MemoryPool* pool) { + if (!pool) pool = default_memory_pool(); + std::shared_ptr out; + RETURN_NOT_OK(TruncateBuffersImpl(std::move(data), &out, pool)); + return out; +} + +} // namespace array +} // namespace arrow diff --git a/cpp/src/arrow/array/buffer_truncate.h b/cpp/src/arrow/array/buffer_truncate.h new file mode 100644 index 000000000000..82292966a469 --- /dev/null +++ b/cpp/src/arrow/array/buffer_truncate.h @@ -0,0 +1,57 @@ +// 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 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include + +#include "arrow/array/data.h" +#include "arrow/result.h" +#include "arrow/status.h" + +namespace arrow { +namespace array { + +/// \brief Truncate the buffers of an ArrayData to only include the bytes +/// required to represent the slice indicated by (offset, length). +/// +/// This function walks the ArrayData tree recursively and, for each node, +/// replaces over-sized buffers with sliced (or copied, when bit-offset +/// conversion is required) equivalents. The ArrayData itself is not mutated; +/// a new ArrayData is returned that shares no oversized backing memory with +/// the input. +/// +/// Uses include: +/// - Pickle serialization of sliced Array / ChunkedArray objects (GH-26685) +/// - Any other path that would otherwise serialize the full backing buffer +/// of a sliced array +/// +/// The logic is adapted from the IPC writer (RecordBatchSerializer::Visit) +/// and handles validity bitmaps, fixed-width data buffers, variable-width +/// offsets + value buffers, list offsets + child arrays, struct children, +/// union type codes + offsets + children, dictionary indices, and run-end +/// encoded arrays. +/// +/// Non-CPU device buffers are returned unmodified (we cannot safely slice +/// generic device memory from this utility). +ARROW_EXPORT +Result> TruncateBuffers( + std::shared_ptr data, + MemoryPool* pool = default_memory_pool()); + +} // namespace array +} // namespace arrow diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index bf9a6d672280..ab774dcd789e 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -743,6 +743,17 @@ cdef _reduce_array_data(const CArrayData* ad): ad.offset, buffers, children, dictionary +cdef _reduce_array_data_top_level(shared_ptr[CArrayData] ad_sp): + """ + Like _reduce_array_data but first truncates all buffers to the range + actually represented by (offset, length). This prevents sliced arrays + from serializing their entire backing buffer (GH-26685). + """ + cdef shared_ptr[CArrayData] truncated + truncated = GetResultValue(TruncateBuffers(ad_sp, NULL)) + return _reduce_array_data(truncated.get()) + + cdef shared_ptr[CArrayData] _reconstruct_array_data(data): """ Reconstruct CArrayData objects from the tuple structure generated @@ -1344,7 +1355,7 @@ cdef class Array(_PandasConvertible): def __reduce__(self): self._assert_cpu() return _restore_array, \ - (_reduce_array_data(self.sp_array.get().data().get()),) + (_reduce_array_data_top_level(self.sp_array.get().data()),) @staticmethod def from_buffers(DataType type, length, buffers, null_count=-1, offset=0, diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index e02898c04101..bf5134770094 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -3158,6 +3158,10 @@ cdef extern from "arrow/array/concatenate.h" namespace "arrow" nogil: const vector[shared_ptr[CArray]]& arrays, CMemoryPool* pool) +cdef extern from "arrow/array/buffer_truncate.h" namespace "arrow::array" nogil: + CResult[shared_ptr[CArrayData]] TruncateBuffers( + shared_ptr[CArrayData] data, CMemoryPool* pool) + cdef extern from "arrow/c/abi.h": cdef struct ArrowSchema: void (*release)(ArrowSchema*) noexcept nogil