NoahKusaba opened a new pull request, #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579
## Which issue does this PR close?
Part of #2578 (hybrid sort-shuffle buffering: split rows by partition on
arrival when memory allows). This is the first of three PRs; it changes no
behaviour on its own.
## Rationale for this change
`SortShuffleWriterExec` reaches into `BufferedBatches` from three places:
- the final encode (`encode_buffered_partitions`) and the spill
(`spill_all_partitions`) each call `take()` and rebuild the same
`PartitionedBatchIterator` gather loop;
- the input loop estimates the buffer's memory by hand: the input batch's
size, plus the index lists' capacity growth measured with
`indices_allocated_size()` before and after each push, accumulated into a
local counter.
The buffer's layout therefore leaks into the writer. Moving these two
operations into `BufferedBatches` gives the writer a narrow interface
(`push_batch`, `drain`, `memory_size`, `is_empty`) and keeps the gather loop
and the memory accounting in one place.
That interface is what the follow-up needs. It adds a second layout that
copies each partition's rows out of an input batch as it arrives, while the
batch is still in cache, and falls back to today's layout when the
per-partition buffers would not fit a memory budget. This is the "auto" mode
Comet planned for its immediate-mode shuffle (apache/datafusion-comet#3845)
but set aside over memory at high partition counts, and close to Gluten's
default hash shuffle writer, which falls back to a buffered writer for many
partitions or wide tables (apache/gluten#6475). A single-node TPC-H SF10
prototype is 8.5% faster in total, and 10–23% faster on shuffle-heavy
queries. With this PR in place, the follow-up adds the new layout inside
`BufferedBatches`; its only writer changes are computing the budget and two
metrics, leaving encode, spill and memory accounting untouched.
## What changes are included in this PR?
- `BufferedBatches::drain(f)`: empties the buffer, passing each output batch
of at most `batch_size` rows to `f` with its partition, in partition order
and then arrival order. Encode and spill now just call it.
- `BufferedBatches::memory_size()`: the buffered batches'
`get_array_memory_size` plus the index lists' allocated capacity, kept as a
running total as the lists grow. The writer mirrors it in the pool with
`try_resize` instead of accumulating growth itself.
- `BufferedBatches::new` takes the batch size, so
`encode_buffered_partitions`
and `spill_all_partitions` no longer need it passed in.
`memory_size()` is the same number the writer computed before (index lists
grow identically), so the pool reservation, the per-task budget check and
therefore every spill decision are unchanged, and the gather produces the
same batches in the same order. A side effect: the index-capacity total is no
longer recomputed over every partition twice per input batch.
## Are there any user-facing changes?
No behaviour change: same shuffle files, same spills, same metrics.
**API change** (`sort_shuffle::BufferedBatches` is public):
- `new(num_partitions, schema)` → `new(num_partitions, schema, batch_size)`
- `take()` / `BufferedTake`, `indices_for()`, `indices_allocated_size()` and
`num_buffered_rows()` are removed, replaced by `drain()` and
`memory_size()`.
## How was this tested?
- Buffer unit tests rewritten against the new API, comparing whole outputs:
drained batches per partition (order, `batch_size` chunking, emptiness
afterwards), and `memory_size()` against the batches plus index capacity.
- Existing writer tests unchanged and passing, including the spill tests
(per-task budget, pool rejection, zero budget) and the round trips.
- `cargo test -p ballista-core -p ballista-executor -p ballista`, clippy
(`-D warnings`) and fmt are clean.
--
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]
