Skip to content

perf: serialize Python input directly from Comet Arrow vectors - #5368

Draft
sunchao wants to merge 2 commits into
apache:mainfrom
sunchao:dev/chao/codex/comet-4383-direct-python-arrow-ipc
Draft

perf: serialize Python input directly from Comet Arrow vectors#5368
sunchao wants to merge 2 commits into
apache:mainfrom
sunchao:dev/chao/codex/comet-4383-direct-python-arrow-ipc

Conversation

@sunchao

@sunchao sunchao commented Aug 15, 2026

Copy link
Copy Markdown
Member

Why are the changes needed?

Closes #4383.

Comet already produces Arrow-backed columnar batches during native execution. A Python mapInArrow or mapInPandas worker 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:

def identity(batches):
    yield from batches

result = events.mapInArrow(identity, events.schema)

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:

Before

Comet's existing Arrow buffers
        |
        | allocate destination vectors and copy every column buffer
        v
writer-owned Arrow vector tree
        |
        | serialize Arrow IPC bytes
        v
Python worker 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.

After

Comet's existing Arrow buffers
        |
        | describe those same buffers as an Arrow IPC record batch
        v
Arrow IPC writer
        |
        | serialize Arrow IPC bytes
        v
Python worker pipe

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 \
  test

The 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:

37 rows -> 7 + 7 + 7 + 7 + 7 + 2

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.py passes all 178 Linux/macOS suite-registration checks. Maven ScalaStyle and Spotless checks pass, and ruff check spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py passes. 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 mapInArrow and 1.27x for mapInPandas when 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Drop the per-batch Comet→Spark buffer copy in CometColumnarPythonInput

1 participant