NoahKusaba opened a new pull request, #2774:
URL: https://github.com/apache/datafusion-ballista/pull/2774
## Which issue does this PR close?
Closes #2578. Second of two PRs, after #2579.
## Rationale for this change
The sort-shuffle writer gathers every output batch with
`interleave_record_batch` at spill or write time. On large shuffles that
gather reads rows that have left the cache, so nearly every row is a miss. On
TPC-H SF10 it is the largest part of the shuffle write. #2578 has the
measurements and the design, including how Gluten handles the same trade-off.
## What changes are included in this PR?
`BufferedBatches` gets a second layout, split on arrival. Each partition's
rows are copied into that partition's `BatchCoalescer` while the input batch
is still in cache, and the batch is dropped.
- **When it splits.** Every input starts in today's deferred layout. Once it
has buffered `partitions × batch_size` rows, it decides once. It splits if
one partly filled batch per partition fits a budget, at the measured row
width plus a 1 MiB string block per long-string column. Otherwise it stays
deferred for good.
- **The budget** is the smallest of:
- the new `split_on_arrival_max_bytes` (default 64 MiB);
- ¼ of `memory_limit_per_task_bytes`;
- ¼ of the input's share of the pool's free memory.
- **Going back.** A split input whose spill is mostly partly filled batches,
a sign of memory pressure, goes back to deferred for good.
- **Stats.** Partition `num_bytes` and spill bytes count the bytes rows use,
not allocated capacity. Otherwise a partly filled batch reports the size of
a full one to AQE. For deferred output this is about the same number as
before.
- **Metrics.** `buffer_time` and `split_inputs` are new writer metrics.
The commits:
1. `perf(core): split sort-shuffle rows by partition on arrival when it
fits`:
the layout, the switching rules, the stats and the metrics.
2. `feat(core): make the sort-shuffle split-on-arrival cap configurable`:
the setting, a new optional proto field, and the docs.
## Are there any user-facing changes?
- **New setting:** `ballista.shuffle.sort_based.split_on_arrival_max_bytes`.
0 turns splitting off and restores today's behaviour.
- **Docs:** documented in `configs.md` and the tuning guide.
- **Proto:** adds an optional field to `SortShuffleWriterExecNode`. Older
plans
decode to the default. This needs the `proto-change` label.
- **Explain output:** `SortShuffleWriterExec` shows two more metrics.
- **API:** no public API changes; `BufferedBatches` is crate-private.
## How was this tested?
- **Buffer unit tests:**
- when the buffer switches, and when it doesn't;
- going back to deferred;
- inline and sparse strings;
- a check of the coalescer's 1 MiB string block.
- **Writer tests:**
- split and deferred runs write identical partitions and report identical
bytes;
- a cap of 0 never splits;
- `num_bytes` excludes spare capacity;
- a small memory pool behaves exactly like a writer that never splits,
spills included;
- the budget follows the pool's free memory.
**Performance**, from a release-build probe on a laptop with ±10% noise:
| Case | Layout | vs `main` |
|---|---|---|
| 16 partitions, strings | split | 5–30% faster |
| 16 partitions, strings, spilling (128 MiB) | split | 27–41% faster |
| 200 partitions, ints (± zstd) | split | same |
| 200 partitions, strings (± zstd) | deferred | within noise |
On TPC-H SF10 with one executor, the prototype was 8.5% faster overall at 12
partitions, and 22% faster on q1–q8 at 200 partitions. All 22 queries pass
`--verify` with the same AQE join decisions.
Still to run: the criterion `sort_shuffle` bench (#1598's 200-partition case)
against `main`.
--
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]
