Repository navigation
[EPIC] Use blocked / chunked memory management in hash aggregation #24704
Description
Activity
- marked Improve aggregate performance with adaptive sizing in accumulators / avoiding reallocations in accumulators #7065 as a duplicate of this issue
on Aug 27, 2026 Thanks for summarizing it
Reacted by Andrew LambI would like to help on this topic too 🚀
Reacted by Andrew Lamb, Kumar Ujjawal and Peter LeeHave an new idea that maybe can support blocked approach in a simpler way
#15591 (comment)Have an new idea that maybe can support blocked approach in a simpler way #15591 (comment)
I like the idea of
EmitTo::Last(n), especially for migration toward the end state (blocked storage) — although introducingEmitTo::Last(n)directly would add complexity for downstream users and the existing emission mechanism. So the proposal below keeps its semantics but doesn't touchEmitToat first.Emission: the tail contract. The contract is emit the state of the
nhighest group indexes; indexes of all remaining groups are unchanged. Nothing shifts and nothing is renumbered — avoiding the ~15x cost #19562 measured for incrementalFirst(n)— and both layouts can serve the same request: today's contiguous storage by truncation (split_offon the tail), blocked storage by popping its last block in O(1) with zero copy. That symmetry is what makes it a migration path. Instead of a new enum variant, it ships as default-implemented trait methods behind a runtime flag (default off):// GroupsAccumulator — all default-implemented; no existing impl breaks fn supports_tail_emit(&self) -> bool { false } fn evaluate_tail(&mut self, n: usize) -> Result<ArrayRef> { /* default: internal_err! */ } fn state_tail(&mut self, n: usize) -> Result<Vec<ArrayRef>> { /* default: internal_err! */ } // GroupValues fn supports_tail_emit(&self) -> bool { false } fn emit_tail(&mut self, n: usize) -> Result<Vec<ArrayRef>> { /* default: internal_err! */ }
Zero downstream compile impact (same pattern as
supports_convert_to_state), and the stream negotiates: tail drain only if all state owners in the plan support it, otherwise fall back toAll— so unconverted UDAF/FFI accumulators keep working unchanged, and a plan mixing blocked built-ins with a merely-truncatable external accumulator still drains incrementally. The countnsits at the call site because the group values and every accumulator must release the same groups each step anyway (or output columns don't align). Scope: end-of-input drain on the unordered path only —First(n)and its renumbering contract stay untouched for ordered streams and the spill/merge read-back. Whether to later fold this into a realEmitTo::Last(n)variant (ideally together with markingEmitTo#[non_exhaustive]) can be a major-release decision after the contract has soaked.Storage. As sketched in this issue: fixed-capacity blocks of
B = target_batch_sizerows (power of two), never resized, allocated on demand, filling strictly front-to-back. Two refinements: block 0 is sized like today's initial capacity, so low-cardinality queries stay effectively flat (aimed at the ~10% regression observed in #15591); and byte/byte-view blocks each own their values buffer and theiri32offsets, so no offset ever spans more than one block — which closes #23694 structurally rather than by widening types.Indexing. The group index stays a single dense
usize: with power-of-two blocks filling front-to-back,(block, offset) = (idx >> k, idx & (B - 1))is a bijection.intern,update_batch,total_num_groups, and the hash table are all unchanged — no trait signature changes anywhere, including UDAFs and FFI. The known cost is one extra indirection per state access (the #15591 concern); hot-loop parity at low cardinality is the declared benchmark gate for the storage PRs, so this gets retired with numbers rather than argument.
Phases toward the end state
1. Benchmarks first — minimum only. For regressions, nothing new is built: ClickBench + the h2o groupby suite already in
bench.share the gate, because they are what aggregation PRs are judged against today and h2o spans the dimensions this touches (int vs. string keys, low vs. high cardinality). For the wins, exactly one new artifact: a drain-phase benchmark that pre-builds ~10k and ~10M groups and records five numbers — total drain time, time-to-first-batch, max inter-batch gap (the long-poll proxy, #19906), peak reserved memory, and reserved memory at the 50%-drained mark (≈100% of peak underAll, ≈50% under tail — a one-sample proof of incremental release). Gate numbers agreed on this issue before any behavioral PR.2. Tail contract. One vertical slice, targeting the refactored streams from #22710 (the declared prerequisite — happy to help review there): the default trait methods + the drain loop (drop the hash table at drain start, truncate-serve, shrink the reservation per step, one await per step) + tail arms for primitive group values, the primitive accumulators, and
count— behind a config flag, with a fuzz mode assertingAll≡ tail modulo row order, and extending the aggregation fuzz framework to run under constrained memory pools (the #19649 scenario — a query that exhausts a fixed pool today and completes with incremental release — becomes a regression test). Everything else negotiates back toAll. Flip the flag after soak; release notes call out thatGROUP BY … LIMITwithoutORDER BYmay return different (equally valid) groups.3. Blocked internals.
Blocks<T>+ microbenches; blockedPrimitiveGroupsAccumulatorandGroupValuesPrimitive(gate: hot-loop parity); blockedByteGroupValueBuilderwith per-blocki32offsets (closes #23694, with a gated >2 GiB regression test — and worth validating against Comet's suite, since #4718 came from there); byte-view; multi-column ideally by helping land @rluvaton'sadd-blocks-implin pieces; then the remaining accumulator families. One requirement I missed earlier: the same implementations also serve the ordered streams and the spill read-back, which emit viaFirst(n)— so each blocked conversion must implementFirst(n)(andAll) over blocks too, with the renumbering semantics intact. so the per-implementation sequence is: blocked storage →First(n)/Allover blocks → delete contiguous. The cross-block compactionFirst(n)needs can be written once inBlocks<T>(those paths keep their tables small, so it needn't be fast), leaving each conversion'sFirstarm thin. Default is one atomic conversion PR per implementation; if one grows too big to review, the escape hatch is a temporary sibling type — old contiguous and new blocked coexisting as separate single-mode implementations selected by negotiation, with a tracked deletion — which is not the dual-mode-inside-one-implementation that sank #15591. Dual storage inside an implementation still never exists (the #15591 lesson). Contiguous survives only as the negotiated compatibility tier for external accumulators — a feature kept indefinitely, not a deprecation candidate — and a small permanent coverage gap remains for types that only theRows-based fallback handles: those stay onAllunless someone later blocks that path.4. Order-sensitive follow-ups. Ordered and partially-ordered streams keep
First(n)and its renumbering contract untouched, as does the spill read-back (which re-aggregates a sorted merge). Mid-stream emission under memory pressure is a separate design — the hash table is live there, so tail eviction would cost a full table sweep — and it is where the interaction with skip-partial aggregation (the pending ClickBench Q5 +61% from #22712) lives. Blocked state also opens a follow-up the epic hints at: block-granular spilling (spill whole blocks under pressure instead of emit-sort-spill of everything), which is a step toward the memory pool proactively requesting release from consumers. If a hard front-first requirement surfaces, the same blocked layout supports a front cursor instead, at the documented cost of forkingFirst's contract and the dense0..total_num_groupsinvariant — recorded as the fallback design, not the default.5. Convert to
EmitTo::Last(n)— the formalization release. At the next major, once the contract has soaked through phases 2–3: fold the trait methods into a realEmitTo::Last(n)variant, markEmitTo#[non_exhaustive]in the same release so this class of breakage never recurs, deprecate the tail methods for one release cycle with a mechanical migration path, and ship the upgrade guide for external accumulator authors (how to move a tail arm from the methods to thematch). The contract is never deprecated — only its provisional surface is folded into the enum. If soak instead shows the methods are fine as a permanent home, this phase collapses to documenting that decision; either way it is decided on schedule, not by drift.
Open question: does anything require front-first drain? My answer is no for the unordered path — nothing forces emitting from the front there: SQL guarantees no output order without
ORDER BY, today's insertion order is incidental, and the paths that genuinely need ordering (ordered/partially-ordered streams, the spill read-back) keepFirst(n)and its renumbering contract untouched. The end state is indifferent either way — over blocks, with the hash table dropped at drain start, either end pops in O(1) — so a front-first requirement would only threaten the phase-2 migration slice, where the tail is the only cheap direction on contiguous buffers. If anyone knows a constraint I'm missing — spill file ordering, partial-stage eviction, assumptions inadd-blocks-implReacted by kamille- added 2 commits that reference this issue
on Aug 30, 2026 #24795 sets the baseline for the follow-up PR, and the numbers show exactly what issue we're facing.
Reacted by kamilleI don't like
EmitTo::Lastsince sorted input is common, for example when usingGroupAccumulatorsinWindowwhich can be sorted by the partition keys.@alamb Thank you for opening this issue!
Also, couple of things to do right away, is marking the
GroupValuesand the implementations of it as private, the reason they are public from what I checked is to allow doing micro benchmarks.
My branch
add-blocks-implcurrently goes the easy-ish way, what if we could do as many breaking changes as we want, and we only need to support blocked approach, the reason I took that path is that:- I want to see what it would take
- Is the blocked approach easily implemented by existing accumulators
- To be able to run benchmark on the best case scenario
EmitTo::All/EmitTo::Firstand emit blocks does not work together very well if the data structure is built in blocks:
KeepingEmitTo::Allmeans that if we emit the output as a single array it requires concatenation which I don't like
KeepingEmitTo::First(n)without shifting the data between blocks means:
1. Whennis larger than block size, it will require concatenation which I don't like as it is expensive and open for issues like we have
2. that blocks are now not fixed size so next emit block will have either partial or would have to concat 2 batches
3. indexing become harder and force every implementors of acc to deal with that shifting
KeepingEmitTo::First(n)with shifting data between blocks:
1. for nested types, likeList, the inner values blocks should be managed by the list itself (which is my current impl), so taking the firstnvalues makes shifting the inner values really complex since the block size is dynamic for the inner values and you must move group of values as 1 to another prev block. so this is really really complex- To surface any problems we might have (like the
EmitTo::FirstI explain below) (which is also why I'm doing this manually and not using LLM)
A problem I have with not having
EmitTo::Firstis that think of this scenario.
the block size is 1000, we can emit early (because the input is ordered for example)
and we got all the values in the first 500 groups, now we got 2 more batches for group 501, and when we try to get another batch we don't have enough memory, if we don't haveEmitTo::First, we can't emit the first 500 finished groups, we can emit the whole block which in our case is partial and group 501 is not finished, which I don't have solution for this yet.Actually, we can keep the
EmitTo::Firstand it would just require shifting the remaining data in blocks rather than 1 big continues allocation, this is more complex to implement (and for nested types is really really complex)Side note: for variable length types (like string, binary, list, etc.) the shifting will require more working memory while shifting, but for fixed size types (if done right) this could be done with no extra memory
Reacted by Andrew LambWe can have
EmitTo::Allif we have a signature that return Vec in that case
Emit all can still be valuable since accumulators might include more tracking data (like unique values for distinct) and emitting all will release those and we can move the blocks memory tracking to the callerEmitTo::First(n)if we have a signature that return Vec in the case thannis larger thanblock_size
but I think wenmust not be larger than block size and it should return Block so users would need to call emit next block with n as the remainder (altough allowing n to be larger than block size would allow the users to avoid cleanup after each block (for example if there is some data that is tracked globally and after every emit, that tracked is get cleaned, so to avoid cleaning for each block can emit first n and one cleanup would be performed)Drafted another approach. It seems performance neutral or positive overal (will try to run the benchmarks in the runner).
Some notable items:
- It has both a flat and and a blocked mode, to avoid regressing the low cardinality cases
- Flat and blocked mode with branching hoisted outside of the loop
- Size is a power of two, so indexing in blocked mode is cheap
Reacted by Jay Zhan@rluvaton I think we should process the sorted case and the normal case(unsorted) in different ways like what is doing currently:
1. For the sorted case
Maybe we can continue to use
Emit::First(n).
And if we decide to remove theEmit::First(n)finally (due to always unacceptable performance, and easy to be misused), I think we can make theGroupValuessorted input aware and useEmit::Allto replaceEmit::First(n).
The logic can be:- Assume the
group valueskeepinga, and it is aware of the input is sorting, and has the api to returncurrent key - When following batch
a, b, b, ccomes, we getcurrent keyfromGroupValuesfirst - Compare the current key
awitha,b,b,c， and found the batch will lead to key switching - Split the the batch to
aandb,b,c，and only putaintoGroupValues - Call
Emit::Alland return the freezedagroup - Input the
b,b,cto resetGroupValues, and go to next loop
2. For the unsorted case
After many experiment, I found blocked approach is only a better memory management approach to reduce the peak memory usage.
And I think our main target just be:- Support really return aggr state block by block, rather than holding a large guy + slicing
- Don't lead to unacceptable performance regression
And I think the totally blocked approach #15591 may be not worthy:
- Really large code changes
- A massive impact on the existing architecture
- And after many many hard works, still lead to unavoidable performance regression
So after struggling with #15591 , I think something simple like
Emit::Lastmay be a way more worthy trying.- Assume the
After many experiment, I found blocked approach is only a better memory management approach to reduce the peak memory usage.
It's also relevant for excessive spilling in the downstream operator, as described here: #22526 (comment)
Even more relevant for the partitioning in datafusion-comet because threre's a shuffle after the partial aggregation that will always hit this spilling issue.
Reacted by kamille17 remaining items
@Dandandan tried this in #24815. I haven't reviewed it fully, but the change looks quite small — I think it's worth trying to find a simpler implementation here.
any implementation of blocked index that require having fixed block size must try to implement the multi group by for list, since the child
GroupValuesblock size is dynamic. so this should be checked if the impl support thatIf multi group by for nested types is the concern, how about we split them into separate paths? We could use the blocked approach for non-nested types and work out a specialized design for nested types.
It is not a concern just saying that having blocks index that is flat (i.e. if the data was single continues memory it would just be an index rather than a BlocksIndex that contain block index and index in block in it) won't work for the nested types.
Reacted by Jay Zhan@alamb @Dandandan @comphead @jayzhan211 I want to push this forward, the one question remaining is:
Do we want to change the existing trait for
GroupsAccumulatoror create a new one for theBlockedGroupsAccumulatorwith an adapter forGroupsAccumulatorso existing implementation would still work and it can be incrementally migrated.From my benchmarks doing continues state/evaluate with emit to first which shifts is really expensive (from what I remember can be 4 times slower).
Do we want to change the existing trait for GroupsAccumulator or create a new one for the BlockedGroupsAccumulator with an adapter for GroupsAccumulator so existing implementation would still work and it can be incrementally migrated
No strong preference on my side — either approach works for me as long as we get the result. I'm also fine if the first version is harder to review; we can look at simplifying it afterwards.
No strong preference on my side — either approach works for me as long as we get the result. I'm also fine if the first version is harder to review; we can look at simplifying it afterwards.
I agree too. I think for this work I would suggest:
- Get a PR up with the base pattern for one or two accumulators, so we can validate the design/pattern carefully
- then we can mechanically switch over the other accumulators to the new pattern.
Reacted by Oleks Vcc @sunchao who also experimenting with hybrid hash join.
Agree with @alamb lets move trait discussion to specific PR, from more wide angle I would consider a way where we can upgrade or rollback without affecting users, exposing them to the painful migration and/or keeping parts of the code used not effectively.
I'm more leaning towards second trait and once its stable to merge it to
GroupsAccumulatorbut I might be missing some detailsHere is the implementation:
It contain
BlockedGroupsAccumulatortrait, supportcountin blocked so you see the example usage, change the entire aggregate to work with blocked (this is possible due to the added adapters), implement BlockedGroupValues for primitive so you will see how it is being usedSo instead of having an adapter and having large performance hit until all is supported I created another PR with alternative approach which support both and have less performance hit for unsupported cases
Please review so I can move this forward
I am thinking about how the block approach might fit with the recent aggregation optimizations we’ve been exploring, and wanted to share a couple of observations:
-
for final aggregation, @jayzhan211 has proposed a promising bucket-based approach that could improve performance while offering memory-management benefits similar to the block approach.
feat: aggregate a large final hash table as hash buckets (experimental, off by default) #25724 -
for partial aggregation, I’ve seen encouraging results from proactive flushing (inspired by part of @jayzhan211 's work in 25724 ). Flushing earlier seems to avoid accumulating more state than is beneficial, improving performance while keeping the working set small. This could also address much of the memory pressure at this stage.
perf: flush partial batches with fixed group capacity (experiment) #26116
If we pursue both optimizations, I wonder would the block approach still offer enough additional benefits.
-
If we pursue both optimizations, I wonder would the block approach still offer enough additional benefits.
I do think it could lower the peak memory requirements for queries with high cardinality aggregates -- for example plans with subqueries that have large numbers of aggregates that are then fed into their own 🤔
But the ideas from @jayzhan211 look quite promishing -- maybe the extra benefit wouldn't be worth the effort
What do you think @rluvaton ?
Reacted by kamilleI do think it could lower the peak memory requirements for queries with high cardinality aggregates...
Seems in final when bucketing is on, the memory peak will just be
1.015 * state sizerather than original2 * state size; and when bucketing is off, the memory usage is always low? Maybe it is already good enough for reducing memory usage?I think the Utf8 case is still open after #26116 (Partial) and #25724 (Final)
Yes, for utf8 or more complex type, bucketing seems still possible lead to slight performance regression, I think we should continue to optimize.
But if take its effect of memory usage reducing into consider, the slight regression seems acceptable (compared with blocked approach) ? And it is actually much simpler than blocked approach.Reacted by Andrew Lamb19. I think the Utf8 case is still open after #26116 (Partial) and #25724 (Final). Maybe we could look into whether blocked storage, or an alternative approach, would help with the string case.
I agree this sounds like a good idea @jayzhan211
So what are the next steps here? Shall we get #26116 (Partial) and #25724 (Final) ready for review?
It would probably be good to hear from @rluvaton to know if the combination of those two PRs would solve his usecase
I'll wait a few days for more input. If there's no further discussion, let's move ahead with reviewing #26116 and #25724. @Rachelint, feel free to start getting #26116 ready for review in the meantime if you'd like. I'll also take a look at #25877.
Note: this issue consolidates and replaces #7065, which dates from 2023 and predates much of the discussion.
What is going on (symptoms)
High-cardinality
GROUP BYqueries in DataFusion suffer from several challenges that look unrelated at first, but all trace back to the same root cause. The problems:Aggregation memory is held until the hash table is fully drained. All group state is emitted as slices of one giant batch, so none of it is released until the last output batch. For a typical two-stage aggregation, this shows up as ~2x peak memory: while the first stage drains its 1 GB of state, the final stage is simultaneously building its own ~1 GB (measured by @2010YOUY01 in PoC: Blocked state management for hash aggregation #22712).
Queries with more than 2 GiB of overall string data in group keys can crash outright. All the bytes of
Utf8/Binarygroup keys are interned into one contiguous buffer addressed byi32offsets; once more than 2 GiB of key bytes accumulate, the offsets overflow and the query fails withoffset overflow, buffer size > 2147483647:offset overflowdatafusion-comet#4718 (report from Comet by @comphead)Downstream memory accounting is off / operators spill when they don't need to. Each output batch reports the memory of the entire aggregation output via
get_array_memory_size(), which causes unnecessary spilling in operators such asRepartitionExecandTopK:RecordBatches rather than one large one #9562 (reported by @alamb)The async runtime stalls when output begins. Producing output for >500k groups block a tokio worker thread for hundreds of milliseconds to seconds, causing latency spikes for everything else on that thread:
Potential copying performance. As groups accumulate, internal buffers repeatedly double in size and copy all existing data (up to 2 copies per element on average), which is likely expensive and cache/TLB unfriendly and could in theory be avoided.
What is causing the problem
Note you can read more about the current group state in the blog Aggregating Millions of Groups Fast in Apache Arrow DataFusion 28.0.0
GroupedHashAggregateStreamstores all per-group state in single contiguous buffers that grow by doubling:GroupValuesimplementation (e.g.PrimitiveGroupValueBuilderorByteGroupValueBuilder). The hash table itself only stores group indexes —usizeoffsets into these buffers.GroupsAccumulatorper aggregate expression (not per group). Each accumulator manages the state for all groups, typically as a singleVec<T>(plus a null buffer), again indexed by a single group index (ausize).This contiguous layout explains each symptom:
EmitTo::All: at end of input, all groups are materialized into one giantRecordBatch, which is then handed downstream asbatch.slice(..)chunks ofbatch_sizerows. The slices share the giant batch's buffers, so no memory is freed until the last slice is dropped.Utf8/Binarygroup keys,ByteGroupValueBuilder<i32>stores all key bytes in a single contiguous buffer addressed byi32offsets, so accumulating more than 2 GiB of key bytes overflows.EmitTo::Allcall is one large CPU-bound operation with no await point.The high level solution sketch
The fix that everyone seems to agree on is to store group keys and accumulator state in multiple blocks instead of one contiguous
Vec. This is the approach used by DuckDB and most other databases with buffer managers, which don't have the luxury of large contiguous arrays. It was originally suggested for DataFusion by @yjshen in 2023.At a high level, the idea is that:
target_batch_size)(block_id, offset)rather than a singlegroup_idxEmitTo::NextBlock), so memory is freed incrementally. The blocks now back individual outputRecordBatches (fixing the accounting).target_batch_sizerows, they are far more likely to stay below 2 GiB of total string values (thei32offset limit), avoiding the string offset overflow.The idea is illustrated here:
Why this is hard to fix
The reason this is so hard to implement is that the entire API is designed around a single
usizegroup index that is assumed to be contiguous and directly addressable:GroupValues::internassigns each distinct group key a dense index0..n, and the hash table stores those raw indexes.GroupsAccumulator::update_batchreceivesgroup_indices: &[usize]andtotal_num_groups: usize; implementations index their stateVecs directly with the group index and grow them with a singleresize(total_num_groups).EmitTo::First(n): after emitting the firstngroups, every remaining group index is renumbered down byn— a notion that only makes sense when the state is one contiguous array.For example
GroupValues::intern:Because this assumption is spread across every
GroupValuesandGroupsAccumulatorimplementation — including user-defined aggregates and the FFI bindings — moving to a blocked(block_id, offset)index:block_idand once to find the actual value within that block.We have discussed ways to address both issues:
(block_id, offset)indexes once the table grows past a threshold — at that point accesses are cache misses anyway, so the extra lookup matters less. See Intermediate result blocked approach to aggregation memory management #15591 (comment) and the earlier version of the same idea in PoC: Blocked state management for hash aggregation #22712 (comment).Past attempts and prototypes
There is a long and distinguished history of trying to address this problem:
RecordBatches (@JasonLi-cn)supports_blocked_groups/alter_block_sizetrait additions, blockedPrimitiveGroupsAccumulator+GroupValuesPrimitiveBatchedVec<T>bench (@Dandandan): O(1) per-block emission is achievable in small stepsEmitTo::FirstBlockas an API-only first step (@hhhizzz)One seemingly obvious alternative would be to use
EmitTo::First(n)to incrementally emit data from the front of the state. However, this does not work either as it is destructive: it requires shifting all remaining elements to the start of the buffers and renumbering every remaining group index. @ahmed-mez tried incremental emission with these mechanics in #19562 and measured it ~15x slower at high cardinality.As part of #22710, @2010YOUY01 has been refactoring the monolithic
GroupedHashAggregateStreaminto dedicated per-path streams, in large part to make changes like blocked state management feasible to implement and review. I (@alamb) thinks completing this refactor is a prerequisite for beginning the blocked state work in earnest: implementing it against the old multiplexed stream would be far more complex and would conflict with the refactoring itself.The same blocked approach was also proposed independently by @alchemist51 in #19649, which includes an experiment (based on @Rachelint's #15591) where a high-cardinality query that fails with resource exhaustion in a 16 GB memory pool today completes once blocked state management is enabled.
An in-progress implementation for multi-column group-by (@rluvaton) using only
EmitTo::NextBlockis being discussed on #15591. It is not yet a PR, but the work seems to be on theadd-blocks-implbranch of their fork.A complementary approach that reduces partial-stage state without changing these traits is cache-efficient (morsel-driven) partial aggregation, proposed by @Dandandan:
Related issues
RecordBatches rather than one large one #9562Related symptoms that will NOT be addressed by this issue
Note that other operators produce giant contiguous intermediate batches too, and show the same failure modes. This issue only covers aggregation; something similar will be needed for joins:
OffsetOverflowError(2147731589)(AQE) datafusion-ballista#1826 (reported by @milenkovicm): TPC-DS Q72 in Ballista fails with the sameOffsetOverflowError, but profiling in that thread suggests the oversized intermediate is produced primarily by the join chain (~1.12B rows, ~96 GB at the join output) before it reaches the partial aggregate — blocked aggregation state will not help there.batch_sizeinstead of emitting everything at once" across operators.