NoahKusaba opened a new issue, #2578:
URL: https://github.com/apache/datafusion-ballista/issues/2578
**Is your feature request related to a problem or challenge? Please describe
what you are trying to do.**
`SortShuffleWriterExec` keeps every input batch and records a
`(batch_idx, row_idx)` pair per row. Only when it spills or finishes does it
build each output batch, gathering the rows with `interleave_record_batch`
(#1598, #1432). This keeps memory flat however many partitions there are.
On large shuffles that gather is slow. By the time it runs, the buffered
batches are out of CPU cache, and each output batch picks its rows out of as
many input batches as there are partitions, a few rows from each, so almost
every row is a cache miss.
Today (simplified from `buffer.rs` and `partitioned_batch_iterator.rs`):
```rust
// 1. As each input batch arrives: keep it, and note where its rows go.
for (p, rows) in per_partition_rows.iter().enumerate() {
for &r in rows {
indices[p].push((batch_idx, r)); // 8 bytes written per row
}
}
batches.push(batch);
// 2. At spill or end time, possibly hundreds of MB later: build each
// partition's output batches by gathering its rows, then write them.
for p in 0..num_partitions {
for chunk in indices[p].chunks(batch_size) {
// rows scattered across ~num_partitions long-evicted input batches
let out = interleave_record_batch(&batches, chunk)?;
writers[p].write(&out)?; // IPC stream, or the spill file
}
}
```
Proposed:
```rust
// 1. As each input batch arrives, while it is still in cache: copy each
// partition's rows into that partition's coalescer, which fills
// `batch_size`-row batches and sets each one aside when full.
for (p, rows) in per_partition_rows.iter().enumerate() {
let taken = take_record_batch(&batch, &UInt32Array::from(rows.clone()))?;
coalescers[p].push_batch(taken)?;
}
// `batch` is dropped here.
// 2. At spill or end time: the batches are already built; just write them.
for p in 0..num_partitions {
coalescers[p].finish_buffered_batch()?; // the last, partly filled batch
while let Some(out) = coalescers[p].next_completed_batch() {
writers[p].write(&out)?;
}
}
```
Both versions hold all the row data in memory until they spill or finish.
What differs is how it is arranged, and when rows are sorted into partitions:
| | Today | Proposed |
|---|---|---|
| Holds | Whole input batches, partitions mixed, plus an 8-byte index per
row | Each partition's rows in its own batches; input batches are dropped |
| Sorts rows into partitions | At spill or end time, by gathering from cold
memory | As each batch arrives, while it is in cache |
| Extra memory | The indices (8 bytes a row) | One partly filled batch per
partition, however few rows it has |
So the data takes about the same memory either way. The speedup comes from
copying each row while its batch is still in cache rather than long after.
Example (TPC-H SF10, one executor): shuffling 60M `lineitem` rows
(`l_orderkey, l_quantity`) into 12 partitions takes 1.75 s, while the scan
alone takes 0.31 s. The gather is the largest part of the shuffle write
(~0.48 s per input partition), followed by encoding (~0.44 s) and recording
the index pairs (~0.23 s).
**How it got this way**
- **Before #1598**, the writer split each input batch by partition as it
arrived, using DataFusion's `BatchPartitioner`. With 200 partitions an
8192-row batch became ~41-row pieces, each buffered, spilled and written as
its own IPC batch: heavy per-batch overhead, poor compression and slow
reads (#1432).
- **#1598** fixed that by adopting DataFusion Comet's approach: keep whole
batches plus row indices, and gather full `batch_size` batches only at
spill or write time. Batches came out the right size, with less memory
overhead as the partition count grows. The trade-off it didn't measure is
the cold-memory gather above.
- **#1626** then found that the gather kept every source string buffer alive
after the DataFusion 52 upgrade (~90× slower on string-heavy data), and
added compaction after each gather.
This proposal splits on arrival again, but into a `BatchCoalescer` per
partition, so output batches are still full-size (#1432's problem stays
fixed). That trades #1598's flat memory for per-partition buffers, which is
why it only switches when they fit.
**Prior art**
- **Comet's "immediate mode"** (apache/datafusion-comet#3845, #3855, both
closed in April 2026): split each batch on arrival into reusable
per-partition builders. It was opt-in
(`spark.comet.exec.shuffle.nativeMode`,
default `buffered`), with an automatic mode "based on the number of output
partitions" planned but never built. It was set aside because the
per-partition builders reserve memory before any data arrives, which does
not scale to the thousands of partitions common in Spark (Comet #887). A
sort-based follow-up (#3940) was 36% faster than buffered at 200 partitions
but used 22 GiB, against 7.4 GiB, at 800. Comet's hash shuffle still
buffers and gathers, like Ballista today.
- **Gluten ships both** (apache/gluten#6475): a hash shuffle writer that
scatters rows on arrival into pre-allocated, reused per-partition buffers
(the default), and a buffered sort shuffle writer it switches to above a
user-set partition-count or column-count threshold. The sort writer was
1.33× faster at SF1000 with 40,000 partitions.
This proposal is the automatic mode Comet planned, in the spirit of Gluten's
hybrid. It answers Comet's memory concern three ways:
- it starts in today's layout, so nothing is reserved per partition up front;
- it switches only once a batch per partition's worth of rows has arrived,
and only if the partly filled batches fit a memory budget (not a partition
count);
- it switches back if memory pressure appears later.
Measured peak executor RSS was lower than today at 12 partitions (3.44 → 3.18
GB on TPC-H q3, q9, q17, q18 and q21). The per-partition `take` still makes
many small allocations; gathering straight into the coalescer would avoid
that
(apache/arrow-rs#8957).
**Describe the solution you'd like**
Copy each partition's rows out of an input batch right away, while the batch
is still in cache, into a per-partition `BatchCoalescer`, then drop the
batch. Spill and the final write then emit those batches as they are.
The catch, and the reason #1598 avoided it: each partition holds a partly
filled batch, so memory grows with `partitions × batch_size × row width`.
So keep today's layout as the default, and switch an input over only when it
is safe:
1. it has seen enough rows to fill a batch per partition, and
2. the partly filled batches fit a budget (the smallest of ¼ of the
per-task spill budget, ¼ of the input's share of the memory pool, and
64 MiB).
If a switched input later spills mostly partly filled batches (memory
pressure), it switches back for good.
These defaults are reasoned, not tuned:
- **A quarter** of the limit keeps most of the memory for full batches, so a
spill writes mostly full batches, not near-empty ones.
- **The pool share** stops one input from claiming more than its part of the
pool. Without it, a 2-column shuffle into 200 partitions in a 256 MB pool
spilled 7,390 times, against 88 today, and ran 1.9× slower.
- **64 MiB** is a quarter of the default 256 MiB task budget, so turning the
budget off behaves the same as leaving it at the default.
- **Waiting for a batch's worth of rows per partition** means each partly
filled batch is likely to fill up.
They cause no regressions on TPC-H SF10 (12 and 200 partitions), on probes
up to 1,000 partitions, on small pools or on `benches/sort_shuffle.rs`, but
other values may do better.
Shuffle stats (`num_bytes`) should also report bytes used rather than
allocated capacity, since a partly filled batch would otherwise look full to
AQE.
Prototype results, single node, TPC-H SF10:
- −8.5% total at 12 partitions; shuffle-heavy queries (q3, q4, q12, q17,
q21, q22) −10 to −23%.
- −22% on q1–q8 at 200 partitions.
- No change on `benches/sort_shuffle.rs` (#1598's wide, 200-partition case),
which keeps today's layout.
- Same AQE join decisions on all 22 queries; all pass `--verify`.
**Describe alternatives you've considered**
- Always split: 3× slower on wide rows at 200 partitions (7,390 spills vs
64).
- Reorder each batch by partition on arrival, as Spark does: copies rows
twice; 19–38% slower on `benches/sort_shuffle.rs`.
- Write each partition to the file as soon as every input has encoded it: no
gain when the CPU is already saturated.
**Additional context**
Planned as three PRs:
1. Refactor, no behaviour change: `BufferedBatches` gets `drain()` and
`memory_size()`. (#TBD)
2. Shuffle stats report bytes used. (#TBD)
3. The split-on-arrival layout and its switching rules. (#TBD)
**Expected effect at SF100 / SF1000 (not measured):** smaller. Whether an
input switches depends on the partition count and row width, not on data
size, but larger runs usually use more partitions, so fewer shuffles fit the
64 MiB cap. At 256 partitions the cut-off is about 32 bytes a row: a key and
a decimal (~24 bytes) would still switch, while wider rows, or any column of
long strings, would keep today's layout and run exactly as now. Where it
does switch, the gain should hold or grow, since a bigger buffer makes
today's gather colder.
Open questions: should the 64 MiB cap and ¼ factor be configurable, perhaps
scaling with executor memory, which would extend the gain to larger runs? All
numbers are single-node SF10; a multi-node run at larger scale would help.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
