apache / apache/datafusion-comet
perf: coalesce small shuffle-read batches before native sort
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
> **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](https://github.com/apache/datafusion-comet/issues/5686#issuecomment-5543566525). 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).
Contributor guide
Assessment
This issue has not been assessed yet.