perf: coalesce small shuffle-read batches before native sort

[peterxcli]

Validation update: end-to-end Spark A/B did not reproduce the native-only result. The common 128-writer / 128-reducer Sort case regressed by about 11%, while larger Sort and Final HashAggregate stress cases were neutral or below a 3% wall-clock gain. No implementation PR is planned from the current approach; see the full results. Future work needs a memory-accounted/adaptive mechanism and a stable stage-level win.

What is the problem the feature request solves?

ShuffleScanExec exposes one record batch per shuffle block. With many map tasks and sparse reduce partitions, a native sort can therefore receive thousands of tiny batches, making per-batch sorting overhead dominate useful work.

A release-mode native benchmark modeled that reduce-side shape with 1 million (Int64, Float64, Utf8) rows and an 8,192-row target. It used one warmup and seven alternating raw/coalesced samples, included coalescing time, and verified identical row count, ordering, and value digests. The experiment was repeated independently with similar results; the second run was:

Input rows/batch	Raw sort	Coalesced sort	Speedup
16	1836.9 ms	58.1 ms	31.6x
64	300.2 ms	48.7 ms	6.17x
256	106.5 ms	46.8 ms	2.28x
1,024	69.4 ms	46.6 ms	1.49x
8,192	43.6 ms	43.2 ms	neutral

This gives a focused Sort result for the broader investigation in #2187 and addresses the shuffle-specific concern raised in #495.

Describe the potential solution

Coalesce small batches immediately before native SortExec when its stage reads native shuffle data. Reuse DataFusion's existing CoalesceBatchesExec, target the configured native batch size, and honor datafusion.execution.coalesce_batches. Register the internal coalescer with the Sort plan so its elapsed time remains visible in the Sort metric. No new user-facing configuration is needed.

Keep this first slice Sort-only. A matching experiment on final aggregation showed a small regression for 1,024-row inputs in both runs, so aggregates, joins, and windows need separate evidence before receiving the same treatment.

Additional context

This is a focused implementation slice of #2187. It concerns reduce-side operator input, rather than shuffle-write block construction (#5002), shuffle-read JNI/IPC micro-costs (#5198), or broadcast collection (#3703).

[peterxcli]

End-to-end release A/B changes the conclusion from the native-only microbenchmark: the current blanket read-side coalescing patch should not be published.

Setup: macOS arm64 release build, Spark 4.1.3, AQE off, native shuffle/direct read, the same patched binary with datafusion.execution.coalesce_batches toggled, one warmup per mode, alternating measured runs, and exact result checks.

Consumer / shape	Coalescing off	Coalescing on	Result
Sort, 1M rows, 128 writers / 128 reducers (3 runs), wall	647 ms	722 ms	11.6% slower
Sort, same shape, native operator	248 ms	276 ms	11.3% slower
Sort, 1M rows, 1024 writers / 1 reducer (7 runs), wall	492 ms	476 ms	effectively neutral
Sort, same shape, native operator	54 ms	53 ms	effectively neutral
Final grouped SUM/COUNT, 8,388,608 rows, 512 writers / 8 reducers (7 runs), wall	116 ms	116 ms	neutral
Final aggregate, same shape, native operator	28 ms	25 ms	10.7% operator-only gain
Final grouped SUM/COUNT, 8,388,608 rows, 2048 writers / 8 reducers (11 runs), wall	371 ms	363 ms	2.2% faster
Final aggregate, same shape, native operator	102 ms	99 ms	2.9% faster

The fragmentation is real: each Comet shuffle frame becomes one native RecordBatch, and writer-side coalescing cannot combine tails from different map tasks. But a read-side wrapper does not remove JNI pulls, IPC decode/decompression, dictionary unpacking, or schema conversion.

For Sort, DataFusion already buffers and bulk-concatenates inputs when the partition is below its default 1 MiB in-place threshold. The extra wrapper duplicates that work and regresses the common small-reducer case. Final grouped HashAggregate has no equivalent concat shortcut, so it can save per-batch dispatch, but even a 2048-writer stress case produced only a 2.2% wall improvement.

There is also a production risk: DataFusion's generic CoalesceBatchesExec is deprecated and its Arrow BatchCoalescer builders are not tracked by a DataFusion MemoryReservation. Applying it to variable-width aggregate state could hide significant memory and perturb spilling.

Decision: I am not opening the implementation PR from this experiment. A future patch should first demonstrate a stable stage-level win (not just operator time) on a real workload and use a memory-accounted/adaptive mechanism that avoids the small-Sort regression.

[peterxcli]

take