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/_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/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 296148f52816..bf5134770094 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 @@ -3148,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 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)