Repository navigation
fix: count buffers shared by sliced batches once in shuffle buffering - #2796
Merged
andygrove merged 1 commit intoNov 7, 2026
Merged
Conversation
Aggregates emit their output as one large array sliced into batch-size pieces. The sort shuffle writer summed get_array_memory_size per buffered batch, which counts the whole shared buffer for every slice, so after the first slice the 256 MiB per-task budget looked exceeded and the writer spilled on almost every batch: TPC-H q18's partial aggregate stage spilled 1,145 times per input partition at SF100. Use DataFusion's RecordBatchMemoryCounter, which counts each buffer once across batches, in the sort shuffle buffer. The counter is reset when the buffer is drained.
andygrove
force-pushed
the
fix-shuffle-shared-buffer-accounting
branch
from
November 7, 2026 17:18
7bf2f03 to
a7ac90d
Compare
andygrove
marked this pull request as ready for review
November 7, 2026 17:22
Member
Author
|
Thanks for the review @milenkovicm |
NoahKusaba
added a commit
to NoahKusaba/datafusion-ballista
that referenced
this pull request
Nov 7, 2026
…plit-on-arrival Resolves the conflict with apache#2796 in sort_shuffle/buffer.rs by moving its RecordBatchMemoryCounter into DeferredBuffer, with its slice test. stays_deferred_when_the_pool_refuses_the_switch now grants 12 MiB rather than 16: its inputs reach the writer as slices of shared arrays, so the deferred rows now count 4 MiB instead of 10, and 16 MiB fit both layouts.
This was referenced Nov 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
No single issue. Split out of #2792. Related: #2520 (performance review tracking), #2318 (buffering shuffle output in memory).
Rationale for this change
Aggregates emit their output as one large array sliced into batch-size pieces. The sort shuffle writer summed
get_array_memory_sizeper buffered batch, which counts the whole shared buffer for every slice. After the first slice the 256 MiB per-task budget looked exceeded, and the writer spilled on almost every batch: TPC-H q18's partial aggregate stage spilled 1,145 times per input partition at SF100. At SF1000 the extra spill volume pushes q5 and q18 past DataFusion's 100 GB spill cap (see the separate spill-limit PR).What changes are included in this PR?
The sort shuffle buffer (
BufferedBatches) uses DataFusion'sRecordBatchMemoryCounter, which counts each underlying buffer once across batches, instead of summingget_array_memory_sizeper batch. The counter is reset when the buffer is drained. A test checks that eight slices of one array are counted once, and again after a drain.Not changed:
BufferExec's dam sumsget_array_memory_sizeper batch the same way, so it also over-charges sliced input (conservatively, by breaking the dam early). It releases each batch's charge as the batch is drained, which a count-once scheme would release too early while later slices still hold the buffer, so it needs a different design and is left out.Also not changed: the byte statistics in
writer.rsandspill.rsstill sum per batch. They describe freshly materialised output batches and do not drive spill decisions.The counter hashes each distinct buffer pointer per batch, which is small next to the partitioning work done on the same batch.
Are there any user-facing changes?
No config or API changes. Shuffle writes of sliced batches spill far less.