fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter - #71
fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter#71adriangb 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.
4255114 to
6b5b27b
Compare
|
Rebased from the DataFusion 55 fork branch ( One conflict, in the file this PR touches. Upstream added a Upstream has not addressed any of this. Revalidated on the new base rather than carried over: the three added tests still fail before the change and pass after it (the emit test still reports 0 bytes while holding 224). The measurement in the description moved slightly on the new base and has been updated: 168,096 bytes retained, 96,960 reported, 71,136 unaccounted (42%), on a precisely specified batch shape. |
|
Macroscope has since reviewed this pull request. An earlier review was skipped by a cost limit; a review has now completed, so that notice no longer applies. |
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.
|
Superseded by the upstream pull request: apache#24858 This copy existed only to run the change through review here before sending it upstream. That is done, so closing this one. The branch is unchanged and still backs the upstream pull request. |
1 participant
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.