Repository navigation
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
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
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
There was a problem hiding this comment.
Thanks @andygrove nothing blocks this PR
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. |
|
run benchmark external_aggr |
|
run benchmark spill_views |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix-issue-258521 (3289c00) to 1da3cf0 (merge-base) diff Run configurationrun benchmark external_aggrResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix-issue-258521 (3289c00) to 1da3cf0 (merge-base) diff Run configurationrun benchmark spill_viewsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix-issue-258521 (3289c00) to 1da3cf0 (merge-base) diff Run configurationrun benchmark external_aggrCPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageexternal_aggr — base (merge-base)
external_aggr — branch
File an issue against this benchmark runner |
alamb
left a comment
There was a problem hiding this comment.
I took a look at this code and ran the tests. It seems ok to me but I don't really understand the spilling code enough to approve.
Do other reviewers think it is ready to merge?
I verified that the slt tests fail without this code change:
warning: `datafusion-sqllogictest` (test "sqllogictests") generated 1 warning
Finished `test` profile [unoptimized + debuginfo] target(s) in 29.04s
Running bin/sqllogictests.rs (target/debug/deps/sqllogictests-616efeba2509d40d)
Running with 16 test threads (available parallelism: 16)
Completed 1 test files in 7 seconds Peak memory drift: drift=103.0 MB allocated=103.0 MB reserved=10.9 KB pool=aggregate_memory_spill.slt consumer=PartialHashAggregateStream[0]
Test files ran concurrently, so the file and consumer above can be wrong. Run with --test-threads 1 to attribute drift to one file.
External error: 8 errors in file /Users/andrewlamb/Software/datafusion/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
1. query failed: DataFusion error: Resources exhausted: Additional allocation failed for FinalHashAggregateStream[1] with top memory consumers (across reservations) as:
FinalHashAggregateStream[1]#144(can spill: true) consumed 5.4 KB, peak 5.4 KB,
FinalHashAggregateStream[0]#148(can spill: true) consumed 5.4 KB, peak 5.4 KB,
AggregateStream[1]#146(can spill: false) consumed 48.0 B, peak 48.0 B,
AggregateStream[0]#142(can spill: false) consumed 48.0 B, peak 48.0 B,
AggregateStream[0]#150(can spill: false) consumed 48.0 B, peak 48.0 B.
Error: Failed to allocate additional 34.8 MB for FinalHashAggregateStream[1] with 5.4 KB already allocated for this reservation - 24.0 MB remain available for the total memory pool: greedy(used: 11.0 KB, pool_size: 24.0 MB)
[SQL] SELECT count(*), sum(n)
FROM (
SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n
FROM generate_series(1, 500000) AS t(v)
GROUP BY v % 32
)
at test_files/aggregate_memory_spill.slt:261
...| # --- Case H: a few groups with large state --- | ||
| # The array_agg states of all 32 groups do not fit in the memory limit, but a | ||
| # single group does. Spilling must write the groups in batches that can be | ||
| # read back within the limit, instead of one batch holding every group. |
There was a problem hiding this comment.
When i ran this test locally
cargo test --test sqllogictests -- aggregate_memory_spillIt produced several warnings about peak memory drift
warning: `datafusion-sqllogictest` (test "sqllogictests") generated 1 warning
Finished `test` profile [unoptimized + debuginfo] target(s) in 1m 10s
Running bin/sqllogictests.rs (target/debug/deps/sqllogictests-616efeba2509d40d)
Running with 16 test threads (available parallelism: 16)
Completed 1 test files in 8 seconds Peak memory drift: drift=102.7 MB allocated=102.7 MB reserved=10.9 KB pool=aggregate_memory_spill.slt consumer=PartialHashAggregateStream[0]
Test files ran concurrently, so the file and consumer above can be wrong. Run with --test-threads 1 to attribute drift to one file.
Is that expected?
There was a problem hiding this comment.
It comes from the Case I queries, but not from the spilling changes. Run on its own, the first Case I query (two partitions) peaks at 103.6 MB allocated with 10.9 KB reserved, sampled when a partial aggregate changes its reservation. With a 2 GB limit, where nothing spills and none of this PR's code runs, it's the same (107 MB). With target_partitions = 1 and 2 GB, the drift is 103.5 MB with 48 B reserved, sampled when the aggregate releases its reservation after emitting. So it's memory that the aggregation holds outside any reservation, such as the batches an aggregate has emitted and the next operator hasn't reserved, which is what #25650 tracks.
| /// counted. Each view counts the bytes of its value, so views that share bytes | ||
| /// are counted as if compaction copied them, although it can write shared bytes | ||
| /// only once. | ||
| pub(crate) struct SpilledRowSizes { |
There was a problem hiding this comment.
This seems oddly specific for aggregation, but I think we can consolidate as part of a follow on
There was a problem hiding this comment.
Agreed. Nothing in it is specific to aggregation, since it sizes rows by how they're written to a spill file, but only aggregate spilling uses it so far. A sort could use it the same way to bound its spill batches. Happy to consolidate it in the #26711 follow-up.
| .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.
Filed
I think once we clean up the old aggregate code some of the older configs will go away
| } | ||
| } | ||
|
|
||
| /// Also emits a batch once its rows take about `max_batch_bytes`, as |
There was a problem hiding this comment.
why is this an optional setting on just the merge stream? It seems like maybe we would want a max_batch_bytes as a global config setting, similar to batchs_size (though it would take quite a bit of effort to thread it through)
There was a problem hiding this comment.
Only aggregate spilling needs it so far. batch_size bounds rows, and a row of aggregate state can hold a whole group, such as an array_agg list, so the batches that spilling writes and replay merges can be far larger than their row count suggests. A global max_batch_bytes like batch_size makes sense to me, but every operator that builds batches would need to estimate row sizes, as this PR does for the two that aggregate spilling uses: IncrementalSortIterator sizes each row as it is spilled, and the merge uses the average row size of each input batch. I kept the option crate-private so it can become a config later without an API change. It could be part of #26711.
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix-issue-258521 (3289c00) to 1da3cf0 (merge-base) diff Run configurationrun benchmark spill_viewsCPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usagespill_views — base (merge-base)
spill_views — branch
File an issue against this benchmark runner |
jayzhan211
left a comment
There was a problem hiding this comment.
@andygrove thanks for working on this! The byte bound fixes the #25851 reproducer and every non-dictionary state I tried (Utf8, Utf8View, LargeUtf8, struct, nested list and string_agg), but dictionary-typed state still fails the same way. Could you take a look at the dictionary case inline?
| DataType::Dictionary(key, _) => { | ||
| Self::Fixed(key.primitive_width().unwrap_or_default()) | ||
| } |
There was a problem hiding this comment.
Dictionary-typed state still hits #25851: a dictionary row is sized by its key only, and take keeps every dictionary value, so a List<Dictionary> state ends up in one chunk, and even a one-row chunk carries the whole dictionary. Since Comet passes dictionaries, I think this belongs in this PR.
Fails on this PR with Failed to allocate additional 160.4 MB (datafusion-cli -m 128M):
SET datafusion.execution.target_partitions = 2;
SELECT count(*), sum(n) FROM (
SELECT k, cardinality(array_agg(v)) AS n FROM (
SELECT value % 16 AS k,
arrow_cast(concat(CAST(value AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)') AS v
FROM generate_series(1, 3000000))
GROUP BY k);I tried the change below locally. The query above returns 16 groups and 3,000,000 values, and so does a variant with 1000-byte values over 300,000 rows, which still fails with only the GC part. dictionaries_count_their_keys needs updating. The GC in gc_array applies to every spill; if you'd rather keep sort unchanged, it could run only on the chunks IncrementalSortIterator yields when max_batch_bytes is set (I haven't tested that variant).
--- a/datafusion/physical-plan/src/spill/spilled_row_sizes.rs
+++ b/datafusion/physical-plan/src/spill/spilled_row_sizes.rs
/// A row of each field.
Struct(Vec<ColumnSize>),
+ /// A key per row, plus the dictionary value it references.
+ Dictionary(usize, Vec<usize>, Box<ColumnSize>),
}
@@ fn new
DataType::Dictionary(key, _) => {
- Self::Fixed(key.primitive_width().unwrap_or_default())
+ let dictionary = array.as_any_dictionary();
+ Self::Dictionary(
+ key.primitive_width().unwrap_or_default(),
+ dictionary.normalized_keys(),
+ Box::new(Self::new(dictionary.values().as_ref())),
+ )
}
@@ fn fixed
- Self::Bytes(_) | Self::Views(_) | Self::List(_, _) => None,
+ Self::Bytes(_)
+ | Self::Views(_)
+ | Self::List(_, _)
+ | Self::Dictionary(_, _, _) => None,
@@ fn max_row
+ Self::Dictionary(width, _, values) => width.saturating_add(values.max_row()),
@@ fn rows
+ Self::Dictionary(width, keys, values) => keys[start..end]
+ .iter()
+ .map(|&key| width + values.rows(key, key + 1))
+ .sum(),
--- a/datafusion/physical-plan/src/spill/mod.rs
+++ b/datafusion/physical-plan/src/spill/mod.rs
@@ fn gc_array
+ DataType::Dictionary(_, _) => {
+ let dictionary = array.as_any_dictionary();
+ let gc =
+ arrow_select::dictionary::garbage_collect_any_dictionary(dictionary)?;
+ let shrunk =
+ gc.as_any_dictionary().values().len() < dictionary.values().len();
+ let (values, values_mutated) = gc_array(gc.as_any_dictionary().values())?;
+ if !shrunk && !values_mutated {
+ return Ok((Arc::clone(array), false));
+ }
+ Ok((gc.as_any_dictionary().with_values(values), true))
+ }
_ => gc_array_children(array),(This also needs arrow-select = { workspace = true } in datafusion/physical-plan/Cargo.toml.) Could you add the query above to Case H as well?
There was a problem hiding this comment.
Thanks, that reproduces here, and so does a single-stage version (target_partitions = 1), which fails with Failed to allocate additional 233.6 MB. Fixed in c5cc3af, but differently from the diff above:
normalized_keys()asserts that the dictionary has values, so it panics on the empty dictionary that a batch of null lists holds.- Compacting the dictionaries of every chunk makes repeated values expensive to spill. The
array_aggstate dictionary holds a copy of each input batch's dictionary, and each chunk references a different subset of those copies, so each chunk writes its own. Onarray_aggover 100 distinct 200-byte values (2M rows, 100K groups, 24 MB limit), the diff above spills 371 MB, against 14.7 MB on main, where every chunk shares the dictionary and the spill file holds it once.
So SpilledRowSizes now decides per dictionary. When fewer than two keys reference each value on average, a row counts the values its keys reference, and IncrementalSortIterator compacts each chunk to those values. Other dictionaries count only their keys and stay shared, as on main. Replay also compacts the dictionaries of nested columns, as it does their views, since a merged batch can hold the dictionaries of all the batches its rows come from. The compaction only runs in aggregate spilling, so the spill files of other operators don't change.
Your query and the 1000-byte variant now pass at 128 MB, as do single-stage and max_spill_merge_fan_in = 2 versions, and the low-cardinality case above spills 10.5 MB and 14.7 MB (20- and 200-byte values), as on main. Case I now has a version of your query at the file's 24 MB limit, with target_partitions 2 and 1, and a query whose values repeat within each group, which needs the replay compaction.
This still leaves one case as on main: a large dictionary whose values repeat within a group but not across groups. For example, array_agg over 5M rows whose values each repeat 10 times in one group still fails at 24 and 64 MB. Compacting would be cheap there, but telling it apart from values that groups share needs counting distinct values per row, which I'd leave for a follow-up.
| 08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>early_emit_count=<slt:ignore>] | ||
| <slt:ignore> | ||
|
|
||
| # --- Case H: a few groups with large state --- |
There was a problem hiding this comment.
Optional: every Case H query uses array_agg, so only the list-of-views path is covered end to end. A string_agg case exercises top-level bytes state; it fails on main and passes with this PR at -m 128M:
SET datafusion.execution.target_partitions = 2;
SELECT count(*), sum(n) FROM (
SELECT value % 16 AS k, length(string_agg(concat(CAST(value AS VARCHAR), repeat('x', 100)), ',')) AS n
FROM generate_series(1, 3000000)
GROUP BY value % 16);There was a problem hiding this comment.
Added in c5cc3af, as a single-stage query. At the file's 24 MB limit with 500K rows, the two-partition version already passes on main, while the single-stage one fails there with Failed to allocate additional 48.3 MB.
# Conflicts: # datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
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. |
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.