fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter - #24858
Draft
adriangb wants to merge 1 commit into
Draft
fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter#24858adriangb wants to merge 1 commit into
adriangb wants to merge 1 commit into
Conversation
…dapter `GroupsAccumulatorAdapter` tracks per-group memory in `allocation_bytes` by measuring each `AccumulatorState::size()` before and after accumulator work and applying the delta. `size()` includes the scratch `indices` vector's capacity, but that capacity is never charged, because: 1. `indices` grows in the per-row push loop, which runs before `sizes_pre` is measured; 2. `indices.clear()` after the accumulator call retains the capacity. So `sizes_pre` and `sizes_post` observe the identical `allocated_size()` on every batch and the delta is always zero. The capacity is charged exactly zero times, permanently, while `size()` is what the aggregate stream reports to the `MemoryPool`, so the pool under-counts and memory-pressure handling is delayed. The same asymmetry has a second effect at emit time: `evaluate` and `state` call `free_allocation(state.size())`, which releases capacity that was never charged, so `allocation_bytes` drifts down (and saturates at zero) across partial emits. Charge the growth explicitly. `indices_allocation_bytes` records the capacity already charged; each batch totals the current capacity in the pass that already visits every group and charges only the difference, so a group whose `indices` grew once and was then cleared stays charged without being charged again. Emitting a state drops its capacity from that total. The invariant is now that `allocation_bytes` equals the sum of `AccumulatorState::size()` plus the `states` vector allocation, which is what the added tests assert. No new per-row work: the push loop is untouched, and no `size()` call is added (`size()` was historically a bottleneck with many distinct groups, which is why deltas are used). The added cost is one `usize` addition per group per batch in an existing loop. Measured on a 16384-row batch across 1000 groups, with 8192 rows in group 0 and the rest spread evenly over the remaining 999, using a 16-byte accumulator: 168,096 bytes truly retained, 96,960 reported before, so 71,136 bytes (42%) went unaccounted. Query results are unchanged.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24858 +/- ##
=======================================
Coverage 81.60% 81.60%
=======================================
Files 1123 1123
Lines 408898 408987 +89
Branches 408898 408987 +89
=======================================
+ Hits 333670 333769 +99
+ Misses 55625 55617 -8
+ Partials 19603 19601 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Author
|
run benchmark clickbench_partitioned external_aggr |
Contributor
Author
|
run benchmark clickbench_extended |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/groups-accumulator-indices-accounting (6b5b27b) to da89c7c (merge-base) diff Run configurationrun benchmark external_aggr
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/groups-accumulator-indices-accounting (6b5b27b) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/groups-accumulator-indices-accounting (6b5b27b) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU 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 Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/groups-accumulator-indices-accounting (6b5b27b) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/groups-accumulator-indices-accounting (6b5b27b) to da89c7c (merge-base) diff Run configurationrun benchmark external_aggr
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU 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 |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/groups-accumulator-indices-accounting (6b5b27b) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU 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 Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
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
3 participants
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 existing issue. This was found while investigating a production out of memory. Happy to file one if you would like it tracked for the changelog.
Rationale for this change
GroupsAccumulatorAdapterreports per-group memory to theMemoryPoolfromallocation_bytes, maintained as a delta ofAccumulatorState::size()measured before and after accumulator work.size()includes the scratchindicesvector's capacity, but that capacity is never charged:indicesgrows in the per-row push loop, which runs beforesizes_preis taken, andindices.clear()afterwards retains the capacity. Both measurements therefore see the sameallocated_size()and the delta is always zero, so the capacity is charged exactly zero times, permanently. Withdatafusion.execution.batch_size = 16384a single hot group can hold a 64 KBVec<u32>the pool never sees, and it gets worse as group skew increases.The same asymmetry bites at emit time:
evaluateandstatecallfree_allocation(state.size()), releasing capacity that was never charged, soallocation_bytesdrifts down and saturates at zero across partial emits. In the added test the adapter reports 0 bytes while genuinely holding 224.Measured on this base with a 16384-row batch across 1000 groups, 8192 rows in group 0 and the rest spread evenly over the remaining 999, using a 16-byte accumulator: 168,096 bytes truly retained, 96,960 reported before, 71,136 bytes (42%) unaccounted.
Production motivation
This is one of a group of fixes prompted by a production process running DataFusion dying at 10.98 GB, with roughly 75% of the heap in per-group
COUNT(DISTINCT)accumulators. This one is independent of the others: it is an accounting correction, not a memory reduction. But an aggregate whose reported size collapses toward zero across partial emits gives the pool exactly the wrong signal on the memory-pressure path.What changes are included in this PR?
The growth is charged explicitly. A new private
indices_allocation_bytesfield records what has already been charged. Each batch totals the current capacity in the pass that already visits every group and charges only the difference, so a group whoseindicesgrew once and was then cleared stays charged without being re-charged. Emitting a state drops its capacity from that total.The established invariant is
allocation_bytes == sum(state.size()) + states.allocated_size(), which the tests assert against an oracle recomputed directly from the states.sizes_prewas deliberately not moved before the push loop:groups_with_rowsis not known until after it, so a pre-measurement there would needsize()for every group (the bottleneck the existing code comment warns about) or a per-row branch in the push loop. This change adds nosize()call and no per-row work, only oneusizeaddition per group per batch in an existing loop.What is the testing strategy for this PR?
Three tests are added in
datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs, all failing before this change and passing after, covering the skewed update path, the merge path, and release on partial and full emit. Each asserts the invariant above against an oracle recomputed from the states, rather than against a hardcoded byte count.cargo check,cargo clippy --tests,cargo testfordatafusion-functions-aggregate-common(50 passed) andcargo fmt --checkall pass locally.Why the coverage sits at unit level and not in
memory_limitThis change only ever makes reported usage larger, so no query can go from failing to succeeding, and there is no "fails before, passes after" integration test to write. We looked for the opposite observable instead, a query that now notices a budget it previously ran past, and measured that it does not exist.
The shape was built to maximise the effect:
covar_samp, chosen because it has no specializedGroupsAccumulatorand so goes throughGroupsAccumulatorAdapterwith a constant size per group accumulator, over one full 8192-row batch per group so every group's scratchindicesvector grows to 8192u32and keeps that capacity.target_partitionspinned to 1, repartition rules off. Peak memory pool reservation, measured with a peak-recordingMemoryPool:The difference is exactly
groups x rows_per_batch x 4bytes in every case, which is the retained scratch, and the base under-reports by up to 81x at 256 groups.That 81x still produces no integration-level observable. Swept memory limits from 250 KB to 8 MB at 128 groups, where this branch reports a 4.28 MB peak: the query succeeds on both base and this branch at every limit, and with the disk manager enabled
spill_countis 0 on both at every limit. The partialAggregateExecruns inOutOfMemoryMode::EmitEarly, so the newly visible bytes are shed by emitting groups early rather than by erroring or spilling. The only difference anywhere in the metrics is the partial aggregate'soutput_batches, 1 on base against 5 here, with identicaloutput_rowsand identical results. That is an internal scheduling detail with no stability guarantee, so asserting on it would be a test that looks stronger than it is.The three unit tests are therefore the right level. They assert
allocation_bytes == sum(state.size()) + states.allocated_size()against an oracle recomputed from the states rather than against a hardcoded byte count, which is a strictly stronger statement than any budget threshold. Run against the base commit they fail withleft: 256, right: 33040on the update and merge paths, andreleases_retained_indices_capacity_on_emitfails withleft: 0, right: 224, which is exactly the reported-zero-while-genuinely-holding-224 case.Are there any user-facing changes?
No public API change (the new field is private) and no change to query results. Only the accounting arithmetic changed, so a memory-limited aggregate now reports its true usage to the
MemoryPooland may spill or fail where it previously ran past its limit undetected.Noted while reading, not fixed here
invoke_per_accumulatorreturns?ontake_arraysfailure withstate.indicesalready populated and never cleared, so a retried call on the same adapter would double-push. An error aborts the query today so it is unreachable, but it is a latent trap if that path ever becomes recoverable.