perf: bulk copy fixed-width columns in ArrowWriter - #5442
Conversation
|
cc @mbutrovich @sunchao whenever you have a moment, I’d really appreciate your review on this. TIA! |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Summary
Reviewed the full four-file contribution at d0fc1336a796b7a2fad9986b48a3decc4f4355ec against observed base and merge-base e975f1383cd79e818287d1807e6e534df7559311. I found no outstanding verified introduced P1/P2 issue after source tracing, independent adversarial probes, and discussion deduplication, and am approving this contribution. The CI and validation limitations below are part of this assessment.
Prior state and problem
The existing fixed-width no-null path pre-sized Arrow vectors but still dispatched a Spark getter and Arrow setter for every value. Spark's on-heap and off-heap vectors expose public primitive bulk getters, so compatible slices can avoid that per-value writer dispatch. The input is still Spark-owned storage rather than an Arrow buffer; this optimization necessarily copies data and must preserve slice offsets, validity, physical representation, and ownership.
The reader previously built each ColumnarArray view itself and chose the null-aware or no-null writer entry point. It also explicitly corrected the root row count after finishing, because per-column writes alone could not account for positive-row, zero-column batches.
Design approach
ArrowWriter.writeColumns and ArrowFieldWriter.writeColumnSlice move slice-aware dispatch into the writer, where the original ColumnVector, source offset, and length remain available. FixedWidthArrowFieldWriter recognizes eligible inputs and Arrow physical types, obtains a fresh primitive array through Spark's public getter, and copies it into Arrow-owned storage. Other inputs keep the existing ColumnarArray scalar path.
SparkColumnarArrowReader now uses this shared batch entry point while retaining its existing per-output-slice writer creation and source-batch ownership. Setting the writer row count outside the column loop preserves zero-column batches without a separate reader-side correction.
Correctness / compatibility analysis
The inspected Spark getters return arrays whose element zero already corresponds to startRow; the second copy therefore correctly uses the primitive-array base offset and destination row zero. Destination capacity is established before copying, byte-length multiplication is widened to Long, and validity writes cover the required bytes. Arrow's logical value count excludes trailing validity padding. Dictionary and null-bearing inputs are excluded from bulk dispatch, and the accepted Spark vector classes are final, so arbitrary overridden bulk getters cannot enter this path.
Fresh standalone probes compiled the unchanged writer/reader source against Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and experimental 4.2.0, with Arrow 18.3.0 and JDK 11/17. I reran all seven saved probe configurations successfully. They covered heap/off-heap inputs, 31/32/33-row boundaries, nonzero offsets, dictionary/custom/constant fallbacks, integer extrema, floating-point raw bits, temporal types, nested scalar fallbacks, zero-column and split batches, growth, reset, retained batches, and allocation-failure cleanup. The harnesses use exact schema-conversion method excerpts with support scaffolding; these are targeted source probes, not full reactor or packaged/shaded-distribution tests. C Stream ownership was source-checked, not exercised through a full native integration run.
All 12 CometArrowStreamSuite tests also passed after the ready-state cutoff in six Linux/macOS CI jobs spanning Spark 3.4 through 4.2. Those jobs checked out synthetic merge 418ee0facd3497717ae3b282726a13dc1d58e1be: its four contribution blobs match this head, but its full tree differs. This is relevant suite evidence, not a full exact-HEAD suite pass. CI still has running work and a failed macOS Spark 4.0 scans job. Its crash artifact and matching native library resolve the return address to hdfsThreadDestructor; the root cause and PR causation remain unproved. I have not treated that failure as green or claimed a baseline reproduction.
Key design decisions
The conservative eligibility gate requires an initial write, at least 32 rows, a little-endian host, a recognized non-dictionary heap/off-heap vector, and no input nulls. Consulting nullability across the entire input vector can retain scalar execution even when a selected slice has no nulls, which is safe and avoids assuming more precise null metadata.
The physical-type allowlist preserves the existing integer, floating-point, date, timestamp, and interval representations. Boolean bit packing, decimal precision handling, calendar-interval conversion, and nested/variable-width writers remain scalar. Public getters avoid reflection and private Spark storage assumptions, at the cost of a transient primitive-array allocation for each eligible column. No new shared mutable state or persistent raw address is introduced.
Implementation sketch
The reader bounds each output slice by the rows remaining in its current Spark batch, creates a writer against the stable root, and passes the batch and slice coordinates to writeColumns. Each field either follows its existing scalar writer or selects the matching primitive getter, ensures Arrow capacity, performs the memory copy, fills validity, and records its value count. finish then publishes the root and field counts.
The four added suite cases cover dispatch and dictionary fallback, supported physical types, zero-column row counts, and reader slice offsets; the existing allocation test also reaches the new required-integer path. The benchmark compares optimized and retained scalar paths on both memory placements across five sizes, retains nullable and row cases, and closes its readers, roots, and inputs. Benchmark code was inspected, but performance numbers were not independently reproduced.
Behavioral changes worth calling out
Eligible slices now take a bulk-copy path with an intermediate JVM array; this is not zero-copy. Tiny slices, dictionary-backed or null-bearing vectors, unsupported physical types, and non-little-endian hosts retain scalar dispatch. No new supported SQL type or dependency version is introduced, and slice order, decoded values, null semantics, and positive row counts for empty schemas are preserved in the checked cases.
The 32-row threshold is a dispatch policy, not a universal throughput guarantee. The benchmark's columnar schema and cases have also changed, so old and new timing labels should not be compared without accounting for that workload difference. Source-batch borrowing, output-buffer ownership, and task-local reader state remain as before.
Suggested improvements
No blocking code change emerged from this review. Optional regression hardening would extend the existing type-dispatch fixture beyond zero-initialized values and make the split-reader fixture include repeated bulk-sized slices, a non-byte-aligned length, and an undersized destination. Those cases passed the independent probes, but keeping them in the repository suite would protect future changes.
Keep the native HDFS cleanup crash visible until its cause is established, and distinguish a rerun on the applicable merge tree from validation of this exact head. Broader architecture/JDK and benchmark coverage would help assess the cutoff's portability; this review does not claim a full native suite, big-endian runtime validation, or a production performance result.
|
Merged, thanks @peterxcli ! |
2 participants
Which issue does this PR close?
Closes #5299.
Rationale for this change
The fixed-width no-null path introduced in #5046 still calls
setValueUnsafeonce per value. Spark's on-heap and off-heap column vectors expose public bulk getters, so sufficiently large fixed-width slices can avoid per-value getter/setter dispatch.This is a bulk-copy optimization, not zero-copy: Spark's public bulk getters allocate and populate a primitive array, then the writer copies that array into Arrow. Because those steps have a fixed cost, tiny slices remain faster on the scalar path.
The implementation uses the writer-level extension point described in #5317. This keeps vector ownership and slice offsets explicit without reflection or access to private
ColumnarArrayinternals.What changes are included in this PR?
ArrowWriter.writeColumnsandArrowFieldWriter.writeColumnSliceso field writers receive the sourceColumnVector, row offset, and slice length.OnHeapColumnVectorandOffHeapColumnVectorslices using Spark's public bulk getters andPlatform.copyMemory.SparkColumnarArrowReaderto use the shared batch API while preserving split-batch offsets and zero-column row counts.Benchmark results
Apple M4, macOS 26.5.2, OpenJDK 21.0.6, Spark 4.1 profile; three columns (
Int,Long, andDouble). Values are best time in nanoseconds per row.With the cutoff temporarily disabled, the bulk path exposed its fixed cost on tiny slices:
With the final 32-row cutoff enabled, the optimized entry point dispatches tiny slices to scalar and larger slices to bulk:
The 32-row cutoff is an empirical guard against the observed tiny-slice regression, not a claim that every CPU and JVM has the same crossover. Comet's default batch size is 8192 rows, where this run measured a 7.2-7.9x speedup.
How are these changes tested?
make core./mvnw test -Dtest=none -Dsuites=org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite(Spark 4.1 / JDK 21: 12 tests passed)./mvnw test -Pspark-3.4 -Dtest=none -Dsuites=org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite(Spark 3.4 / JDK 17: 12 tests passed)CometArrowWriterBenchmarkvia Maven Exec on the Spark 4.1 profile, both with the cutoff temporarily disabled and with the final dispatchgit diff --check