Repository navigation
fix: size aggregate spill batches by bytes as well as rows - #25853
Conversation
Aggregate spill files were written in batches of `batch_size` rows, so a few groups with large state, such as `array_agg` over a low-cardinality key, were written as one batch holding every group. Replay then had to hold that batch and all of its groups at once, and failed with `ResourcesExhausted` even though each group fit in memory. Write spill batches of about 1 MiB, based on the average row size of the spilled state, and replay the spill files with a merge batch size no larger than the smallest spilled batch, so the merge does not combine the small batches back into one large batch. Closes apache#25851
github-actions
Bot
added
sqllogictest
The first version cut spill batches by the average row size of the spilled state, and lowered the replay merge's batch size to the fewest rows per batch of any run. That throttled whole replays to tiny batches whenever one run had large rows, and still wrote huge batches when large groups were adjacent in key order or when string views shared bytes. - Cut spill batches by the size each row takes once spilled (SpilledRowSizes), so every batch stays near the target unless it holds a single larger group. - Emit replay merge batches by bytes as well as rows, so runs of small groups keep batches of batch_size rows. - Combine small replay output batches up to batch_size rows.
- Pass batches of at least half the target through the replay output coalescing, so only small batches are copied, and always concatenate the pending ones. - Reuse the batch size that BatchBuilder::push_batch already measures instead of measuring it again in the merge. - Take Option<usize> in StreamingMergeBuilder::with_max_batch_bytes, like the other merge layers. - Build SpilledRowSizes for each chunk instead of keeping it alongside the sort indices, and cut chunks of fixed-width rows without visiting every row. - Model the batch as a struct of its columns in SpilledRowSizes.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25853 +/- ##
========================================
Coverage 82.89% 82.89%
========================================
Files 1153 1154 +1
Lines 457487 457960 +473
Branches 457487 457960 +473
========================================
+ Hits 379213 379629 +416
- Misses 55264 55302 +38
- Partials 23010 23029 +19 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Measuring every row of a spilled batch made spilling aggregates with variable-width state, such as string keys, about 9% slower. Measure an upper bound on the row size once per batch instead, and only measure rows one by one when batch_size rows of that size could exceed the target. The rows of fixed-width columns are no longer measured one by one either.
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @andygrove. A few suggestions inline, plus some small ones:
IncrementalSortIterator's doc (sorts/stream.rs:292-294) still says each chunk hasbatch_sizerows. Withwith_max_batch_bytesa chunk can be shorter.ColumnSize::max_rowandColumnSize::rowsboth decode a view's length and apply theMAX_INLINE_VIEW_LENrule. A smallview_bytes(view: u128) -> usizehelper would keep the two in sync.- nit:
SortPreservingMergeStream::row_bytesholds per-input averages, soavg_row_byteswould match its doc.
| } | ||
|
|
||
| #[test] | ||
| fn largest_row_bounds_every_row() { |
There was a problem hiding this comment.
This test builds the same view and list arrays as views_count_their_bytes_unless_inlined and lists_count_their_elements, and nearly the same batch as fixed_width_and_offsets. Its max_row asserts could move into those three tests, next to the row asserts for the same arrays.
other_types and rows_of_fixed_width_columns_have_one_size also build the same FixedSizeBinary(3) array, so they could be one test.
There was a problem hiding this comment.
Done in b7cd911. row_sizes now returns the max_row bound along with the row sizes, so each test checks both for the same arrays, and largest_row_bounds_every_row is gone. other_types is folded into rows_of_fixed_width_columns_have_one_size.
| let num_rows = |batches: &[RecordBatch]| -> Vec<usize> { | ||
| batches.iter().map(RecordBatch::num_rows).collect() | ||
| }; | ||
| for from_spill_files in [true, false] { |
There was a problem hiding this comment.
The byte cap is only set by aggregate replay, which always merges spill files. Each multi-level pass (create_new_merge_sort) then builds this same SortPreservingMergeStream over the spill-file streams, so the false run covers no code path that the true run misses. Keeping only the spill-file case would drop the from_spill_files flag and the in-memory branch of merge_wide_and_narrow_rows.
There was a problem hiding this comment.
Done in b7cd911. The test now merges only spill files, without the from_spill_files flag or the in-memory branch.
|
run benchmark tpch tpcds |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix-issue-258521 (832a7eb) to bdad988 (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix-issue-258521 (832a7eb) to bdad988 (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix-issue-258521 (832a7eb) to bdad988 (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix-issue-258521 (832a7eb) to bdad988 (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @andygrove nothing blocks this PR
2010YOUY01
left a comment
2010YOUY01
left a comment
There was a problem hiding this comment.
Thank you for this fix.
I’m only familiar with aggregates/. For that part, my suggestion is to handle the small-output-batch problem differently. I think it’s separate from this PR, so we could remove that part here and address it in a follow-up PR. See my comments for details.
It would also be great to have someone familiar with sorts/ review this. @rluvaton and @EmilyMatt should have more context on that code.
| /// Combines consecutive small batches of `input` into batches of up to | ||
| /// `batch_size` rows and `max_bytes`. Batches of at least half that size are | ||
| /// passed on as they are. | ||
| fn coalesce( |
There was a problem hiding this comment.
IIRC, producing smaller batches is an internal issue for ordered aggregation, and it would be better solved within the operator rather than by adding a wrapper here.
There was a problem hiding this comment.
Removed in b7cd911, so replay returns the OrderedFinalAggregateStream output as on main. It emits the groups that each input batch completes, so even on main replay emits about batch_size / runs rows per batch. Agreed that combining them belongs in the operator, I'll follow up separately.
| .with_batch_size(batch_size) | ||
| .with_max_batch_bytes(Some(SPILL_BATCH_TARGET_BYTES)) | ||
| .with_reservation(merge_reservation) | ||
| .with_replay_headroom() |
There was a problem hiding this comment.
Just a rough intuition: we’ve accumulated quite a few special cases to make memory-limited execution work. There may be an opportunity to unify and simplify some of them.
There was a problem hiding this comment.
Agreed. The spill merge now has three options that only aggregate spilling sets (with_replay_headroom, with_intermediate_merge_sizing and with_max_batch_bytes). I'd rather revisit them together in a follow-up than grow this PR.
There was a problem hiding this comment.
Filed
I think once we clean up the old aggregate code some of the older configs will go away
|
While this can be a useful mechanism, I don't think it's the right solution for the problem at hand, I think the aggregate estimations can and should be improved, this ties us into behaviour that can be both bad for perf, and complex to maintain (and is also not very accurate since it uses averages) |
- Remove the coalescing of replay output. Ordered aggregation emits the groups that each input batch completes, which can be few, and that is better handled within the operator in a follow-up. - Check with EXPLAIN ANALYZE that the last two Case H queries spill. - Apply the view size rule of SpilledRowSizes in one view_bytes helper, and test max_row next to the row sizes of the same arrays. - Test byte-limited merges only from spill files, the only merges that set a byte limit. - Rename row_bytes to avg_row_bytes, and document that IncrementalSortIterator chunks can be shorter than batch_size.
|
Addressed the review feedback in b7cd911:
|
|
@EmilyMatt I agree the aggregate's estimates are worth improving, since they decide when it spills and whether it keeps room to sort and write the spill. But they don't fix this issue, which happens after the spill, when the spilled state is read back. A spill file is read back one batch at a time, and spilled batches are bounded only by On the other points:
|
My point is that no 184MB should be spilled. for a 128MB pool we should be spilling by a calculation that would allow the 2x readback expected(so a much earlier spill), and this should be calculated before inserting a new input batch, and if not possible - spill first, and only then insert the new batch into the cleared state. |
|
@EmilyMatt I prototyped your suggestion on current main to check it: before inserting a batch, the final and single-stage hash aggregates spill their table if the table plus the batch would exceed a budget. I also tried inserting the input one row at a time, so that a large input batch can't force a large spill either. On the reproducer from #25851 (
This PR passes both. When the input arrives in small batches, as in the single-stage queries of Case H, spilling at 2–4 MB does let them pass on main, because the merge's split path eventually shrinks the batches of the runs it splits. But replay then runs one row per batch: the last Case H query spills 50 times instead of 6, emits 340K batches instead of 75, and takes about 1.1 s against 0.1 s with this PR. So the size of the batches replay has to hold is set by how spill files and merges cut rows into batches, not by when the aggregate spills. That is what this PR changes. Reserving before inserting would still be worth doing on its own, since today the aggregate can overshoot its budget by a whole input batch before it spills, which is what #25606 tracks. |
Resolve the conflicting test imports in sorts/stream.rs. apache#23565 made VIEW_SIZE_BYTES test-only, but SpilledRowSizes uses it to size spilled views, so it is defined for all builds again.
A merged batch references the whole data buffers of the spilled batches its rows come from, and replay keeps the list states of each group, such as those of `array_agg`, until the group completes, so it held about twice the state of each group. Since apache#23565, spilling no longer compacts view arrays whose buffers are fully used, so merged batches can reference more of them. On main, Case H's query with 14 adjacent large groups failed with ResourcesExhausted at 24 MB. Compact the view arrays of nested columns before replay. Top-level columns are left as they are, since group keys and the accumulators of top-level states typically copy the values they keep. Also note in SpilledRowSizes that views sharing bytes are counted as if compaction copied them.
|
Merged main into the branch. Beyond the conflict, this needed two changes:
@comphead both came after your approval, in case you want to look at 3289c00. |
Dictionary rows were sized by their keys only, and `take` keeps every value of a dictionary, so a chunk of `array_agg` states over dictionary-encoded values carried the whole state dictionary, which replay could not hold. Dictionaries whose values are mostly unique, with fewer than two keys per value, now count the values their keys reference, and each spilled chunk holds only those values. Dictionaries whose values many keys share still count their keys only and stay shared, since compacting each chunk would write the shared values once per chunk. Replay also compacts the dictionaries of the nested columns of merged batches.
|
Merged main again and addressed the latest reviews:
@comphead, the dictionary handling is new since your approval, in case you want to look at c5cc3af. |
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
Thanks @andygrove, I think the issues are all addressed.
# Conflicts: # datafusion/physical-plan/src/spill/mod.rs
|
@alamb @2010YOUY01 do you have any additional feedback, or can we merge this one (and maybe consider for 56.0.0)? |
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Thanks @andygrove. I have nothing to add beyond the discussion so far, and to @alamb's question, I think this is ready to merge. As one more data point, a single-stage array_agg over 5M rows in 32 groups (about 17 MB of state each) fails on main at 64 and 96 MB, and returns all 5,000,000 values with this PR.
I don't fully understand the underlying logic, so I can't offer concrete feedback. Just sharing an additional perspective here. I have no objections and trust the reviewers' approval judgment. Suggestion: If anyone has a good understanding of the overall sort spilling design, it would be helpful to add a top-level comment summarizing how the existing mechanisms fit together. Concern: Given the PR history in fixing sort spilling, I have the impression that more fixes may actually introduce more bugs rather than make spilling more stable. The process feels a bit like gradient descent programming—fixing corner cases one by one, the review also focus on letting AI to identify more corner cases, rather than refining a clear high-level design. This makes it hard to tell whether we're converging toward a stable implementation. |
|
I put together this page to help me understand the overall sort spilling design. It has diagrams, code links and measurements, in case it helps others here. When memory runs out, This PR adds a byte limit for aggregates, at the spill writer and the merge output. Remaining work, as untested ideas that could go under #26711:
Please correct anything I got wrong. |
# Conflicts: # datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
|
Filed the remaining dictionary case (values that repeat within a group but not across groups, with a large dictionary) as #27240, with a reproducer. I'd keep it out of this PR and handle it as a follow-up, alongside the consolidation in #26711. @2010YOUY01 @jayzhan211 thanks for the design notes. This PR stays a bounded fix for #25851, and the design questions (one byte limit for every caller, removing the three workarounds) belong under #26711. |
Main now disallows DictionaryArray::with_values in clippy.toml. The dictionary was just garbage collected, so it holds no more values than rows.
|
Pushed 1132fb7: the merge of main from today failed clippy, because main now disallows |
…ws can merge The multi-level merge reserves twice the largest batch of each spill run. When an external sort spills batches of wide rows, two such reservations can exceed the memory pool, and the merge fails or re-spills repeatedly. This commit bounds the in-memory size of each batch that a spill file holds. `SpillManager::with_max_batch_bytes` sets the bound, and `InProgressSpillFile` writes each appended batch as `SpillPieces`: - The pieces are cut by `SpilledRowSizes`, the per-row spilled size from apache#25853, after the IPC padding of each piece. - Each piece compacts its columns with `gc_columns`, so a dictionary that many keys share stays shared. - A piece whose recorded size is still over the bound is halved. - The recorded size counts IPC padding, so a batch read back never exceeds it. The external sort sets the bound so that the smallest merge fits in `sort_spill_reservation_bytes`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…pill merge The multi-level merge re-spills a run at half its batch size when two runs do not fit in memory, and it tracks a limit per run to drive that. The sort and the aggregate now bound their spill batches so that the smallest merge fits, so the fallback is no longer needed. This commit deletes the re-spill, the per-run limits and their tests. The merge keeps the output cap of apache#25853. It takes the change from apache#26987, adapted to keep that cap. The upgrade guide says to keep spilled batches to at most a quarter of the merge memory. Two aggregate memory call-sequence snapshots lose one `shrink`: the merge no longer releases its reservation to try a re-spill. Co-Authored-By: Jay Zhan <jayzhan211@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Which issue does this PR close?
Rationale for this change
A grouped aggregate whose groups carry large state, such as
array_aggover a low-cardinality key, fails withResourcesExhaustedafter it spills, even though each group fits in memory and the same data spread over many small groups succeeds.Spilling writes the aggregate state in batches of at most
batch_sizerows, whatever their size in bytes. With fewer groups thanbatch_size, every group lands in one batch, and everything after the spill has to hold that batch at once:batch_sizeof them, rebuilding a batch of large groups.What changes are included in this PR?
Aggregate spilling and replay now bound their batches by bytes as well as rows, with a target of 1 MiB (
SPILL_BATCH_TARGET_BYTES):IncrementalSortIterator::with_max_batch_bytescuts the sorted state into chunks by the size each row takes once spilled, so a batch only exceeds the target when it holds a single larger group. Row sizes come from the newSpilledRowSizes, which counts view values as they are written after compaction and nested values by their elements.StreamingMergeBuilder::with_max_batch_bytesmakes the merge emit a batch once its rows reach the target, estimated from the average row size of the batches they come from, as well as atbatch_sizerows. It applies to every pass of the multi-level merge, so rows of small groups still merge in batches ofbatch_sizerows.takekeeps every value of a dictionary. Dictionaries whose values many keys share count only their keys and stay shared, as onmain, so that each chunk does not write its own copy of them.gc_nested_arrays). A merged batch references the buffers of every batch its rows come from, andarray_aggkeeps the lists it merges until its group completes.Both new options are crate-private and only used by aggregate spilling and replay; other merges are unchanged. The intermediate merge sizing from #25428 now receives the largest row count of the batches each run actually wrote. This applies to every stream built on
AggregateSpill(hash and ordered, final and single-stage). The legacyGroupedHashAggregateStream, used only whendatafusion.execution.enable_migration_aggregateisfalse, is unchanged.What is the testing strategy for this PR?
aggregate_memory_spill.slt, with a 24 MB limit, whereEXPLAIN ANALYZEassertions check that each query spilled. All of these fail onmainwithResourcesExhausted:array_aggover 32 groups of about 2 MB each, with two partitions (one spill file per partition) and with one partition (several spill files, merged).array_aggover dictionary-encoded values, with two partitions and with one.string_aggwith one partition.array_aggover dictionary values that each repeat within one group, at 8 MB.SpilledRowSizes(including views that share bytes, nested views, fixed-width rows, and dictionaries of unique, shared and no values), for byte-limited chunks inIncrementalSortIteratorand their dictionaries, for byte-limited merges of spill files, and for dictionary compaction.With
datafusion-cli -m 128M, the reproducer from the issue now returns 16 groups and 3,000,000 values, and so does the variant from the issue's additional context. The reproducer with 3,000 one-value groups added,array_aggover dictionary-encoded Parquet strings, and the reproducer with its values cast toDictionary(Int32, Utf8), which fail onmain, also succeed.Performance compared with
main:external_aggr(TPC-H SF1, integer keys): no case is slower; 7 are unchanged and 1 is 1.06x faster.GROUP BYover 5M string keys at 256 MB is unchanged (0.61 s against 0.62 s).array_agggroups take 0.80 s at 128 MB, against 2.99 s onmain. Theremainspills 26M rows (1 GB) while merging its runs, and this PR spills each of the 8M rows once. At 512 MB, where both spill twice, this PR takes 1.00 s against 0.96 s, from measuring the rows of the spilled batch that holds the large groups.sort_preserving_mergeshows no change beyond run-to-run noise. Merges without a byte limit only check that none is set.array_aggover a dictionary of 100 distinct values (2M rows, 100K groups, 24 MB) spills the same 10.5 MB (20-byte values) and 14.7 MB (200-byte values) asmain, in the same time.Are there any user-facing changes?
Aggregations whose groups have large state can now complete under memory limits where they previously failed after spilling. No API or configuration changes.