Repository navigation
[Python] Fix pickle serializing full buffers for sliced arrays (GH-26685) - #53914
Closed
Asthenia0412 wants to merge 2 commits into
Closed
Asthenia0412 wants to merge 2 commits into
Asthenia0412 wants to merge 2 commits into
Conversation
…ecutor (apacheGH-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 (apache#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<Datum> 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.
…eGH-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 apache#26685
|
Thanks for opening a pull request! This pull request has been automatically closed because you currently have 4 open pull requests, which is more than the limit of 3. Due to the increase in pull requests opened by AI bots, and in order to keep the review queue manageable, Apache Arrow limits contributors without repository access to at most 3 concurrently open pull requests. This helps make sure each pull request gets the attention it needs and that work in progress does not go stale. Once one of your other open pull requests has been merged or closed, you are welcome to reopen this one. See also: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Priority: Critical 🔴
Ray Data / Dask showstopper. Ray had to pin
arrow < 7on Windows becausepickling a sliced array serialises the entire backing buffer, leading to
100x-1000x bloat in multi-process workers.
Problem
Fix (pitrou-approved Option 1)
General
arrow::array::TruncateBuffers(ArrayData)utility, factored out ofRecordBatchSerializer— the IPC writer already had this logic; picklesimply never called it.
C++ (
cpp/src/arrow/array/buffer_truncate.{h,cc}): walks theArrayDatatree and truncates each backing buffer to exactly the bytesneeded. Handles validity bitmaps (copy + bit-offset shift), fixed-width
data, variable-width offsets (zero-base + slice + child-slice), list
offsets + child slicing, struct children, dense/sparse union type codes
arrays. The offset is reset to 0 in the returned
ArrayDataso noleading padding survives. Dictionaries themselves are not truncated
(shared by convention; IPC writer does the same).
Python (
python/pyarrow/array.pxi):_reduce_array_data_top_levelcalls
TruncateBuffersbefore recursing through the tuple dissection,and
Array.__reduce__uses it.ChunkedArray,RecordBatch, andTableall flow throughArray.__reduce__transitively, so the fixpropagates automatically.
CMake (
cpp/src/arrow/CMakeLists.txt): registers the new source.Cython (
python/pyarrow/includes/libarrow.pxd): extern declarationof
arrow::array::TruncateBuffers.Verification (local)
g++ -std=c++20 -fsyntax-only buffer_truncate.cc✅cython -3 --cplus lib.pyx✅ (arrow-levelperformance hints only, no errors)
python/pyarrow/tests/test_pickle.pyadditions in this branch).Risk
arrays (now they reflect only the actual data), but the reconstructed
object is semantically identical. Any downstream consumer that previously
stored/transported a pickled sliced array can still unpickle it — they
just get a much smaller payload and an identical array.
Closes #26685