[sunchao]

Which issue does this PR close?

Closes #7257.

Rationale for this change

A broadcast build fed by native shuffle can receive thousands of tiny batches. Each batch crosses Arrow FFI and opens a separate serialized stream.

The shared serializer allocation and codec improvement is now a separate PR, #7310. This PR excludes those serializer files. Please land #7310 first; the measurements below evaluate the remaining batching benefit with that improved serializer applied.

What changes are included in this PR?

Add experimental, default-off spark.comet.exec.broadcastExchange.producerCoalesce.enabled. Direct local Comet native shuffle inputs group small decoded INT, BIGINT, BOOLEAN, STRING and BINARY batches before Arrow export and serialization. AQE partition ranges, row order and shared shuffle dependencies are preserved. Large, unsupported and singleton inputs pass through; execution stays native.

Groups are bounded by rows, retained/copy estimates and input count. Conservative bitmap padding remains in the copy estimate, including for aliased Boolean columns. Reader buffers remain outside execution-pool accounting; the contributor memory guide documents their limits and lifetime. Driver broadcast assembly is unchanged.

Review cleanup shares direct-buffer validation, test frames and stream fixtures, centralizes row validation, removes constructor-only catches and overlapping/non-observing tests, and shortens documentation. Task completion and runtime-error cleanup remain. Shared benchmark mode selection lives in the inert base companion to avoid initializing another Spark session.

The mode remains default-off. The fresh wide-input control regresses, and complete-query timings are order-sensitive; broader workload benefits and concurrent-task memory measurements are still needed before default enablement.

How are these changes tested?

• Passed 15 focused Rust tests: 13 coalescer and two JNI append tests.
• Passed 152 JVM tests across native-shuffle, decoder, Celeborn-reader, input-RDD and block-reader suites. Six existing Spark 4-only cases were canceled by version guards. Coverage retains exact [3,3,8,2] → [6,8,2] output order, nullable supported types with AQE on/off, shared dependencies and cleanup.
• Passed strict Spark 3.5/Delta production and test compilation, formatting, CI suite/config checks, benchmark-runner checks and git diff --check.
• The separate perf: reduce per-batch broadcast serializer allocations #7310 serializer passed all 16 UtilsSuite tests against the same dependency classpath. Its new metadata/chunk test fails against main at the intended multiple-chunk assertion.
• Five independent source reviews have no remaining actionable P1/P2 findings. Hosted CI is separate and still running; local Clippy was not rerun.

Current performance measurements

Measured 48e1024bfc plus #7310's serializer (2660acf816) on four pinned AMD EPYC Milan CPUs, JDK 17.0.21, Spark 3.5.9 and a 4 GiB JVM heap. The #7310 comparison overrides only compiled Utils classes; benchmark classes, native library, inputs and settings are otherwise identical. Class origins and native-load traces were verified, and source/artifact hashes stayed unchanged across all eight successful JVM runs.

Each serializer ran in fresh JVMs in OFF/ON/OFF and ON/OFF/ON order, with warmup and at least five iterations/two measured seconds per arm. Values below are logged elapsed milliseconds divided by iteration count; ranges are ranges of arm means, not confidence intervals.

The producer benchmark runs native decode → Arrow FFI → broadcast serialization over 8,192 rows in prepared, uncompressed in-memory frames. Transport and driver assembly are excluded. Ordered values/nulls, row counts, streams and bytes are checked outside timing. The repeated 6 KiB string is a compression/byte-budget control, not representative JSON.

Input rows per frame	#7310 serializer, batching OFF (ms)	#7310 serializer, batching ON (ms)	Streams OFF → ON
1, INT + BIGINT	164.3–187.1	15.83–17.17	8,192 → 8
8, INT + BIGINT	22.21–23.34	2.257–2.331	1,024 → 1
8,192, full-batch control	0.211–0.214	0.215–0.217	1 → 1
64, plus a 6 KiB STRING	31.98–32.08	35.89–60.79	128 → 7

Tiny-batch producer time improves by roughly 10x against #7310. Full-batch overhead is about 0.004–0.006 ms. The wide-string control is slower with batching in both orders, with substantial order variation; no wide-row speedup or general performance gain is claimed. One-row serialized output falls from 2,295,384 to 64,784 bytes. Decoder-only means for that input were 52.15 ms OFF / 15.16 ms ON, including native decoding, JNI/FFI, iteration and cleanup rather than isolating FFI alone.

CometBroadcastLocalShuffleBenchmark measures a complete query: 65,536 build rows in 128 tasks → 64 local native shuffle partitions → broadcast hash join with 1,048,576 probe rows → count/sum collection. Every iteration regenerates the inputs and includes shuffle, broadcast assembly/distribution, join and aggregation, without cached input/stages. All four query runs returned [1048576,88351366032], verified native plans through aggregation, and reduced broadcast input streams from 8,192 to 64.

With #7310 applied:

Order	First arm (ms)	Middle arm (ms)	Last arm (ms)
OFF / ON / OFF	856.4	630.4	694.4
ON / OFF / ON	769.6	764.2	535.4

The OFF range is 694.4–856.4 ms and ON is 535.4–769.6 ms. These overlap; the initial reverse-order ON arm is effectively neutral versus OFF. The faster later ON arms do not establish a uniform 25–30% query improvement.

The query uses explicit broadcast, AQE off, local[4], and the benchmark framework's on-heap mode with an unbounded native pool. These are controlled local measurements, not TPC-H, AQE conversion, general workload or production-memory evidence. Both benchmark classes accept --on-first; the producer also accepts --decode-only.

Local native builds use the current source in an isolated snapshot with an alternate Cargo.lock because the registry lacks a submitted AWS dependency version. Arrow 59.3.1 and DataFusion 55.2.0 are unchanged, but 43 package versions differ; the submitted Cargo.lock is untouched. These local results do not validate that lockfile. The release library uses LTO off and 16 codegen units.

[sunchao]

cc @andygrove @viirya @comphead

[andygrove]

Thanks @sunchao. The tiny-batch numbers are impressive, and I like that the PR is upfront about what the benchmarks do and don't show.

I read this with FFI, shuffle and memory in mind and did not find a correctness problem. The decoded batch owns its data, so the direct buffer can be reused. Native handle access is under the same monitor as close(). The flush-plus-passthrough case that leaves two outputs ready is handled. Cleanup runs on the close, error and task-completion paths. The new config is off by default, so I see no behavior change against the latest release unless it is turned on.

I also checked it against #7193, which caps the driver-side coalesceBroadcastBatches at 1 GiB per string or binary vector. The two merge cleanly and are independent. This PR only groups small inputs on the producer side, so the total string bytes reaching the driver are unchanged and #7193 is still needed. Its per-input split and its fallback for a single oversized batch work the same with larger or fewer inputs.

I have two questions.

1. Why is spark.comet.exec.broadcastExchange.producerCoalesce.enabled disabled by default? The description points at the wide-string control, where one ON arm was slower and the variance was high. Is that the whole reason, or is there something else, such as the memory held outside the pool? Related to that, what would you want to see before flipping the default? For example, whole-query numbers for the local-shuffle path, since the TPC-H results are from the Celeborn implementation. I'd like to know the plan so we don't leave a large win behind an opt-in flag with no path to enabling it.

2. The Rust tests cover a flush followed by a passthrough, where two outputs are ready. Does a Scala-level test exercise that sequence through NativeBatchDecoderIterator.fetchCoalesced, with small batches followed by one at least half the row target? The row order across the two outputs is what the broadcast result depends on, so I'd like to see it pinned at the JVM boundary as well.

[viirya]

Thanks for the detailed write-up and benchmark. Before going further on the native side, I'd like to understand where the OFF-path cost actually goes. This adds ~1.3k lines (five JNI entry points, a native handle lifecycle, a second decoder mode, an RDD flag) for an experimental, default-off feature, so I'd like the necessity question settled first.

Utils.serializeBatches creates a ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) and a fresh codec stream for every batch. In Spark, the first write allocates a full zeroed 1 MiB heap chunk, and toChunkedByteBuffer then copies the used part into a right-sized buffer. For 8,192 one-row batches that's roughly 8 GiB of zeroed allocation, which is about the size of the 1.16 s OFF time in the first benchmark row. Could you break the OFF time down into serialization setup vs. FFI export? If most of it is in serializeBatches, a JVM-only change would get most of the win for every broadcast input, including columnar shuffle, Celeborn, and all types. That could mean a size-aware initial chunk, or one stream per partition instead of one per batch.

Related: the whole-query numbers come from the internal native-Celeborn implementation, and this PR excludes Celeborn. Could you add an end-to-end measurement of this code on local native shuffle, e.g. a TPC-H query where AQE turns an SMJ into a BHJ?

Smaller points:

• Repeated caveat: the "8 MiB is not a total memory bound" caveat appears in the config doc, three docs pages and the Rust docs, and the user-facing docs list internal constants users can't tune. One explanation in the contributor guide, plus a short user-facing description, should be enough.
• Test placement: the new local-decoder tests live in CometCelebornShuffleReaderSuite, presumably for TrackingInputStream/intFrame. Moving those helpers next to NativeBatchDecoderIteratorLifecycleChecks, or into a small shared test util, would keep the suites aligned with what they test.

[comphead]

nice nice, checking

[comphead]

Thanks for the detailed write-up. I read this against the existing thread and kept to points that are not already raised there, mostly about reusing what the shuffle crate already has and where the tests live.

[sunchao]

@andygrove, addressed both questions in b97daacb45:

1. The mode remains default-off for both reasons: wide-row results do not establish a consistent additional gain, and the reader retains native buffers outside the execution pool. I added a concrete enablement bar to the description: broader whole-query benefit over the improved JVM serializer, neutral large/wide controls, and concurrent-task memory measurements. The new complete local-shuffle query does show a residual benefit (930–1,114 ms OFF versus 644–823 ms ON across both orders), but it uses a controlled explicit broadcast, AQE off, and the benchmark's unbounded on-heap execution mode. It is useful current-path evidence, not enough by itself to flip the default.
2. Added the JVM boundary case with distinct rows in frames [3,3,8,2] and target 16. It asserts output sizes [6,8,2] and every row in order, exercising the pending-group flush and the simultaneously ready passthrough before the final group. These local tests now live in NativeBatchDecoderIteratorSuite, registered in both CI matrices.

The driver-side size cap remains independent. Local validation passed 17 focused Rust tests and 169 JVM tests, with six existing version-gated cancellations; strict compilation, formatting and Clippy also passed. Hosted CI is running on the updated head.

[sunchao]

@viirya, your serializer-allocation diagnosis was right. I added the shared JVM improvement as a separate commit, 070de7f1bf: output chunks use Arrow data/dictionary sizes, clamped to 4 KiB–1 MiB, with no new config. The remaining review changes and benchmarks are in b97daacb45.

I reran the producer benchmark with the original serializer restored as an isolated Utils override, sharing all other current Java/native artifacts, and with the new serializer. Fresh JVMs ran both OFF/ON/OFF and ON/OFF/ON. For 8,192 one-row inputs, ranges of arm means were:

Serializer / batching	Producer time
Original serializer / OFF	1,135–1,257 ms
Current serializer / OFF	232–314 ms
Current serializer / ON	18.1–19.0 ms

The current decoder-only means were 51.3 ms OFF / 18.0 ms ON. Those include decoding, JNI/FFI and cleanup, rather than isolating FFI itself. The shared allocation fix removes much of the original cost; native grouping still reduces the remaining per-batch work and streams from 8,192 to 8. The description now compares against that improved baseline and replaces the superseded measurements.

I also added a complete local native-shuffle → broadcast join → aggregation benchmark. With the improved serializer it measured 930–1,114 ms OFF / 644–823 ms ON, with identical results and 8,192 → 64 driver input streams. Both case orders are reported. This is an explicit-broadcast local[4] query with AQE off and an unbounded native pool in benchmark on-heap mode; it is not TPC-H, an AQE conversion measurement, or production memory evidence. Wide-string producer results overlap across orders, so no repeatable wide-row benefit is claimed and native batching remains opt-in. The local alternate-lockfile limitation is still stated in the description.

The documentation details are consolidated in the contributor memory guide, and the local decoder tests moved into their own registered suite with a shared framing fixture. The inline replies cover the accounting, metric, cleanup, trace and test changes. Local validation passed 17 focused Rust tests and 169 JVM tests (six existing version-gated cancellations), including the new multi-chunk binary serialization test; hosted CI is running.

[comphead]

Thanks @sunchao this looks good to me, let me run simplification round

[comphead]

Thanks for working through the first round. All five of my earlier points are addressed at b97daacb45. This pass looks for code that could be shared or dropped, tests that overlap, and comments that could be shorter. None of it changes behavior.

Two smaller points that don't need their own threads:

• The CometNativeShuffleInputRDDSuite test mostly checks that withNativeBatchCoalescing passes its constructor arguments through. The AQE arm of the CometNativeShuffleSuite test already exercises CoalescedPartitionSpec end to end, so the unit test only adds the skew specs.
• In memory_management.md the last two lines of the new paragraph are wrapped short, which looks like a reflow leftover. "Measure the setting on your workload before enabling it, particularly for wide rows" appears both in the CometConf doc, which generates the configs page, and in tuning/shuffle.md.

[comphead]

Is the bitmap padding accounting needed? concat_batches builds one null buffer per output column, so the 64-byte padding is allocated once per group, not once per input. Charging padded_bitmap_bytes, plus the extra Boolean padding, for every input overcounts the copy. The doc comment allows overcounting, and the 8 MiB budget is far below the i32 offset limit. Using get_slice_memory_size() plus metadata as the copy estimate would remove bitmap_bytes, padded_bitmap_bytes, the subtraction and the Boolean branch. BatchBytes could then become three locals in eligible_bytes.

[sunchao]

I kept the conservative bitmap allowance. You are right that charging padding per input can overcount, but removing it does not preserve the output-copy bound: get_slice_memory_size() reports logical bitmap bytes, while concatenation rounds Boolean values and null buffers to 64-byte blocks. With aliased Boolean columns, retained source buffers are shared but each output column needs its own bitmaps, so logical bytes plus metadata need not cover those allocations. I added a short comment explaining why the allowance stays; BatchBytes still keeps that retained-versus-copied accounting together.

Updated in 48e1024bfc.

[viirya]

Thanks for the quick and thorough turnaround. The numbers against the improved serializer answer my main question: batching still gives about 12x on the producer and about 25–30% on the full local query, and it also cuts the decode/FFI and driver-side batch counts, which a JVM-only change wouldn't reach. You were also right about Celeborn: only useLocalShuffle dependencies get a local handle, and materialization excludes them, so the takeover I described can't happen. The replies to my earlier threads work for me.

One structural request: could 070de7f1bf go in its own PR? It changes serialization for every broadcast, regardless of the new flag, and it stands on its own. We squash-merge, so landing it here would fold a default-on change into a commit titled as an experimental opt-in feature, which makes later bisects and reverts harder. As a separate PR it can land first, and this PR's baseline becomes plain main.

With that split and the simplification points from @comphead's latest pass addressed, I'm happy to see this go in.

[sunchao]

@comphead, the simplification pass is addressed in 48e1024bfc. The inline replies cover the shared JNI buffer validation, centralized row check, direct constructor setup, shared fixtures/benchmark helpers, and consolidated tests. I also removed the forwarding-only RDD test and the duplicated/reflowed documentation from the review body.

I kept conservative bitmap padding for the output-copy bound and explained the aliased-Boolean case in that thread, which remains open for discussion. The other addressed threads are resolved.

Validation passed 15 focused Rust tests and 168 JVM tests across both changes (six existing version-gated cancellations), plus strict Spark 3.5/Delta compilation and formatting/config checks. Five independent source-review passes have no remaining actionable P1/P2 findings. One benchmark initialization issue found during review was fixed before any timings were taken; mode selection now lives in the inert base companion.

Both PR descriptions now contain fresh measurements. The wide-string control is slower with native batching, and the complete-query results are order-sensitive, so those limitations are explicit and the flag remains default-off. Hosted CI is still separate from the local validation.

[sunchao]

@viirya, split the shared serializer work into #7310 (2660acf816) and removed Utils.scala/UtilsSuite.scala from this PR's diff. It includes the lazy per-iterator codec reuse and a regression that actually exercises underestimated small chunks. #7310 can land first, independently of the opt-in batching change.

I refreshed both serializer baselines and both batching orders using the same Java/native artifacts apart from the Utils override. With #7310 applied, the one-row producer case measures 164.3–187.1 ms OFF versus 15.83–17.17 ms ON, with 8,192 → 8 streams. The complete-query means are OFF/ON/OFF = 856.4/630.4/694.4 ms and ON/OFF/ON = 769.6/764.2/535.4 ms. The initial reverse-order ON arm is effectively neutral; I removed the earlier blanket query-speedup characterization. The wide-string control also regresses, so there is no general or wide-row speedup claim.

All eight fresh JVM runs returned the expected values/counts, with unchanged source and artifact hashes. Full methodology, controls and local alternate-lock limitation are in the descriptions. #7310 separately reports the shared serializer improvement with batching OFF in both arms. The cleanup from @comphead's pass is pushed in 48e1024bfc, with the bitmap-padding rationale left for discussion.

[comphead]

it would be great to test, there is some use case waiting for this fix.