fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter - #24858
Open
adriangb wants to merge 2 commits into
Open
fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter#24858adriangb wants to merge 2 commits into
adriangb wants to merge 2 commits 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.92% +0.32%
==========================================
Files 1123 1123
Lines 408898 414792 +5894
Branches 408898 414792 +5894
==========================================
+ Hits 333670 339833 +6163
+ Misses 55625 55502 -123
+ Partials 19603 19457 -146 ☔ 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 |
A `Single` mode aggregate spills rather than emitting groups early, so the scratch capacity the adapter now charges is observable as a spill: at 128 groups of 8192 rows the retained `indices` hold 4 MiB against a 1 MiB limit, which the base commit runs straight past with `spill_count` 0.
adriangb
marked this pull request as ready for review
September 2, 2026 21:16
adriangb
requested review from
2010YOUY01
and
a balanced review from Copilot
September 2, 2026 21:16
Contributor
Author
|
@2010YOUY01 mind reviewing this PR? |
Contributor
There was a problem hiding this comment.
🟢 Approval recommended
The accounting fix and regression coverage have no unresolved issues.
Pull request overview
Fixes memory accounting for retained scratch indices in GroupsAccumulatorAdapter, allowing memory limits and spilling to work correctly.
Changes:
- Tracks and releases retained index-vector capacity.
- Adds unit and end-to-end spill regression tests.
File summaries
| File | Description |
|---|---|
datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs |
Corrects scratch-memory accounting and adds unit tests. |
datafusion/core/tests/memory_limit/mod.rs |
Verifies retained scratch memory triggers spilling. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
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.
GroupsAccumulatorAdapternever charges the capacity of its scratchindicesvector to theMemoryPool. An aggregate holds megabytes that the pool does not see, so a memory limit does not stop it.Reproduction
This needs only
datafusion-cli. There is no patch, no custom allocator and no data file.covar_samphas no specializedGroupsAccumulator, so it runs throughGroupsAccumulatorAdapter. The query makes 128 groups. Each group gets one full 8192-row batch. The scratch vectors hold128 * 8192 * 4bytes, which is 4 MiB against a 1 MiB limit.Merge base
da89c7c85b. The aggregate runs past the limit and does not spill:This branch. The aggregate sees the same bytes and spills:
da89c7c85bspill_countspilled_bytesspilled_rowsThe query returns the same 128 rows on both builds. The run takes under a second. Both numbers repeat exactly across runs.
target_partitions = 1makes the effect visible. The planner then folds the aggregate into oneAggregateMode::Singlenode, which spills. AnAggregateMode::Partialnode usesOutOfMemoryMode::EmitEarlyand sheds the bytes instead.Which issue does this PR close?
No existing issue. I found this when I investigated a production out of memory. I can file an issue if you want it in the changelog.
Rationale for this change
The adapter keeps a running total in
allocation_bytes. It measuresAccumulatorState::size()before and after the accumulator work, then charges the difference.The scratch vector grows in the per-row push loop. That loop runs before the adapter measures
sizes_pre. Theindices.clear()call after the work keeps the capacity. Both measurements therefore see the same capacity, the difference is always zero, and the adapter never charges the capacity.evaluateandstatehave the opposite error. Both callfree_allocation(state.size())and release a capacity that the adapter never charged.allocation_bytesthus falls to zero across the partial emits.The size of the hole is
groups * rows_per_batch * 4bytes.How large the error is
An instrumented allocator measured these numbers, so the CLI cannot reproduce them. A counting
GlobalAllocgives the heap that the query holds. A peak-recordingMemoryPoolgives the reported bytes.The peak heap agrees between the two builds to within 8 bytes. The memory use does not change. Only the reported number moves.
What changes are included in this PR?
A new private field
indices_allocation_bytesrecords the capacity that the adapter already charged. Each batch totals the current capacity in the loop that already visits every group, then charges only the growth. An emit removes the capacity of the emitted state from that total.This adds no
size()call and no per-row work. It adds oneusizeaddition per group per batch to an existing loop.The invariant is
allocation_bytes == sum(state.size()) + states.allocated_size(). Four new tests assert it against an oracle that they recompute from the states. All four fail onda89c7c85band pass here.Are there any user-facing changes?
No public API change and no change to query results. Only the accounting arithmetic changes.
A memory-limited aggregate now reports its true size to the
MemoryPool. It can therefore spill, or fail where it cannot spill, in cases where it previously ran past its limit.A note on metrics
grouped_hash_stream.rsrecords apeak_mem_usedgauge from the pool reservation. That gauge is the exact number this PR corrects.EXPLAIN ANALYZEdoes not print it, andEXPLAIN ANALYZE VERBOSEdoes not print it either. The reproduction above therefore usesspill_countunder a fixed limit. If the aggregate exposedpeak_mem_used, a reviewer could see this bug with no memory limit at all.