perf: serialize Python input directly from Comet Arrow vectors - #5368
Draft
sunchao wants to merge 2 commits into
Draft
perf: serialize Python input directly from Comet Arrow vectors#5368sunchao wants to merge 2 commits into
sunchao wants to merge 2 commits into
Conversation
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
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.
Why are the changes needed?
Closes #4383.
Comet already produces Arrow-backed columnar batches during native execution. A Python
mapInArrowormapInPandasworker also consumes Arrow batches, so the data reaching the Python boundary is already in the right physical representation. However, the existing Spark 4.x Comet runner still makes a complete intermediate copy before sending each batch to the worker.For example, even an identity UDF pays this cost:
For every input batch, the current runner allocates writer-owned destination vectors and copies each column's validity, offset, and value buffers into them. Nested arrays, structs, and maps repeat the same work recursively. The Arrow IPC writer then reads those copied buffers and writes their contents to the Python worker's pipe:
That first copy does not transform the values, change the schema, or help the worker interpret the data. It only adds per-batch allocation, CPU work, and memory-bandwidth pressure. The cost becomes more visible with wide schemas, variable-width values, nested columns, or many batches per partition.
As an illustration, an 8,192-row batch containing roughly 16 MiB of Arrow column buffers currently requires roughly another 16 MiB of destination buffers and an extra 16 MiB memory copy before those same bytes are written to the pipe. After this change, the only new Arrow data buffer needed for that batch is the wrapping struct's validity bitmap: 8,192 bits, or 1 KiB. The 16 MiB example is illustrative; the exact savings depend on the batch layout.
The IPC write itself remains necessary because the Python worker is a separate process. This change removes the avoidable copy before that write.
What changes were proposed in this PR?
The runner now treats Comet's existing Arrow vectors as the data source for the outgoing IPC record batch, instead of first materializing equivalent writer-owned vectors. The stream keeps its existing schema header and Python-worker contract, but the root used to advertise that schema becomes schema-only. When a batch arrives, the runner constructs the record-batch view directly over the original column buffers and hands it to the same Arrow stream writer.
Python expects the input columns beneath one non-null struct, so the outgoing record batch adds the corresponding struct field node and a small all-valid bitmap before the existing column buffers. Arrow's normal vector-unloading logic preserves the depth-first layout of nested structs, lists, and maps, along with null counts and buffer metadata. The worker therefore receives the same logical schema, column values, and stream framing as before; the intermediate column-buffer copy simply disappears.
The native-backed vectors remain owned by their original Comet batch throughout this process. Their buffers are retained while the IPC message is written synchronously and released afterward, including when serialization fails. No buffers are transferred between Arrow allocators, no borrowed vectors are closed, and no shared-root allocator changes are required. This keeps the change separate from the allocator work discussed in #4294.
The optimization remains on the existing opt-in Spark 4.x execution path. Spark 3.5 continues to use its current fallback, and configurations requiring incompatible large-variable-width Arrow layouts still fall back to Spark. Regression coverage, CI suite registration, user documentation, and benchmark setup are updated to reflect the new data path.
How was this PR tested?
The focused JVM suite exercises direct serialization across separate allocators, verifies that borrowed buffer reference counts return to their original values, checks nested lists/structs/maps and null fields, covers empty and zero-column batches, and injects a write failure to verify cleanup:
JAVA_HOME=/opt/homebrew/opt/openjdk@17/libexec/openjdk.jdk/Contents/Home \ mvn -B -Pspark-4.0 -Pscala-2.13 \ -DwildcardSuites=org.apache.spark.sql.execution.python.CometArrowPythonRunnerSuite \ testThe Spark 4.0 reactor completed successfully with 5 focused Arrow tests and 23 additional JVM unit tests passing. The same focused Arrow suite also passed on Spark 4.1, and the Spark 3.5 profile compiled successfully.
Real PySpark 4.0 workers were exercised on both the accelerated and fallback paths. One regression sends 37 rows containing nested structs, arrays, maps, and nulls through an identity-style worker in six source batches:
The test checks both the returned nested values and the batch boundaries observed by Python, covering source-vector turnover rather than only a single-batch round trip.
python3 dev/ci/check-suites.pypasses all 178 Linux/macOS suite-registration checks. Maven ScalaStyle and Spotless checks pass, andruff check spark/src/test/resources/pyspark/benchmark_pyarrow_udf.pypasses. The documented benchmark build was also checked with:make -n release PROFILES='-Pspark-4.0 -Pscala-2.13'For additional context, the local wide-schema end-to-end benchmark measured approximately 1.24x for
mapInArrowand 1.27x formapInPandaswhen comparing the optimized Comet path with vanilla Spark. These numbers include Python worker and IPC overhead; they are workload-dependent and do not isolate the incremental benefit of removing this copy.