[EPIC] Use blocked / chunked memory management in hash aggregation

[alamb]

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 BY queries in DataFusion suffer from several challenges that look unrelated at first, but all trace back to the same root cause. The problems:

1. 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).

2. Queries with more than 2 GiB of overall string data in group keys can crash outright. All the bytes of Utf8/Binary group keys are interned into one contiguous buffer addressed by i32 offsets; once more than 2 GiB of key bytes accumulate, the offsets overflow and the query fails with offset overflow, buffer size > 2147483647:

   • Grouping operations on large datasets can overflow i32 offsets #23694 (reported by @maxburke)
   • Comet aggregation task crashes with offset overflow datafusion-comet#4718 (report from Comet by @comphead)
   • fix: experiments with large types for aggregated values datafusion-comet#4791 (workaround by @comphead)

3. 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 as RepartitionExec and TopK:

   • Hash aggregation produces batches reporting huge memory size #22526 (reported by @ariel-miculas)
   • Generate GroupByHash output in multiple RecordBatches rather than one large one #9562 (reported by @alamb)

4. 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:

   • [EPIC] Eliminate Long Polls in HashAggregate via Chunked Storage and Incremental Emission #19906 (reported by @ahmed-mez)

5. 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.

   • Manage group values and states by blocks in aggregation #11931 (reported by @Rachelint)

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

GroupedHashAggregateStream stores all per-group state in single contiguous buffers that grow by doubling:

• Group keys are stored by a GroupValues implementation (e.g. PrimitiveGroupValueBuilder or ByteGroupValueBuilder). The hash table itself only stores group indexes — usize offsets into these buffers.
• Aggregate state is stored by one GroupsAccumulator per aggregate expression (not per group). Each accumulator manages the state for all groups, typically as a single Vec<T> (plus a null buffer), again indexed by a single group index (a usize).

                                         ┌──────────────┐   ┌──────────────┐   ┌──────────────┐
                                         │┌────────────┐│   │┌────────────┐│   │┌────────────┐│
    ┌─────┐                              ││accumulator ││   ││accumulator ││   ││accumulator ││
    │  5  │                              ││     0      ││   ││     0      ││   ││     0      ││
    ├─────┤                              ││ ┌────────┐ ││   ││ ┌────────┐ ││   ││ ┌────────┐ ││
    │  9  │                              ││ │ state  │ ││   ││ │ state  │ ││   ││ │ state  │ ││
    ├─────┤                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
    │     │                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
    ├─────┤                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
    │  1  │                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
    ├─────┤                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
    │     │                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
    └─────┘                              ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
                                         ││ │        │ ││   ││ │        │ ││   ││ │        │ ││
                                         ││ └────────┘ ││   ││ └────────┘ ││   ││ └────────┘ ││
                                         │└────────────┘│   │└────────────┘│   │└────────────┘│
    Hash Table                           └──────────────┘   └──────────────┘   └──────────────┘


stores "group indexes"                     There is one GroupsAccumulator per aggregate
which are indexes into                     (NOT PER GROUP). Internally, each
the state vectors                          GroupsAccumulator manages the state for
                                           multiple groups

This contiguous layout explains each symptom:

1. Memory held until drained: output is produced via EmitTo::All: at end of input, all groups are materialized into one giant RecordBatch, which is then handed downstream as batch.slice(..) chunks of batch_size rows. The slices share the giant batch's buffers, so no memory is freed until the last slice is dropped.
2. Crash on wide string keys: for Utf8/Binary group keys, ByteGroupValueBuilder<i32> stores all key bytes in a single contiguous buffer addressed by i32 offsets, so accumulating more than 2 GiB of key bytes overflows.
3. Spilling / accounting: every slice of the giant batch reports the full underlying allocation to the memory accounting, not just its own rows.
4. Runtime stall: materializing all groups in a single EmitTo::All call is one large CPU-bound operation with no await point.
5. Copying performance: growing the contiguous buffers requires reallocating and copying all existing data.

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:

• Blocks store some number of rows (most likely target_batch_size)
• Blocks are not resized; when a block fills up, a new block is allocated (which avoids copying and the growth is predictable and incremental)
• A group index becomes some form of (block_id, offset) rather than a single group_idx
• Emission now happens one block at a time (e.g. something like EmitTo::NextBlock), so memory is freed incrementally. The blocks now back individual output RecordBatches (fixing the accounting).
• Since each block is capped at target_batch_size rows, they are far more likely to stay below 2 GiB of total string values (the i32 offset limit), avoiding the string offset overflow.

The idea is illustrated here:

                                         ┌──────────────┐   ┌──────────────┐
                                         │┌────────────┐│   │┌────────────┐│
    ┌─────────┐                          ││accumulator ││   ││accumulator ││
    │  (0,5)  │                          ││    AGG     ││   ││    SUM     ││
    ├─────────┤                          ││ ┌────────┐ ││   ││ ┌────────┐ ││
    │  (1,3)  │                          ││ │ block  │ ││   ││ │ block  │ ││
    ├─────────┤                          ││ │   0    │ ││   ││ │   0    │ ││
    │         │                          ││ │        │ ││   ││ │        │ ││
    ├─────────┤                          ││ │        │ ││   ││ │        │ ││
    │  (0,1)  │                          ││ │        │ ││   ││ │        │ ││
    ├─────────┤                          ││ └────────┘ ││   ││ └────────┘ ││
    │         │                          ││            ││   ││            ││
    └─────────┘                          ││ ┌────────┐ ││   ││ ┌────────┐ ││
                                         ││ │ block  │ ││   ││ │ block  │ ││
    Hash Table                           ││ │   1    │ ││   ││ │   1    │ ││
                                         ││ │        │ ││   ││ │        │ ││
                                         ││ │        │ ││   ││ │        │ ││
                                         ││ │        │ ││   ││ │        │ ││
                                         ││ └────────┘ ││   ││ └────────┘ ││
                                         │└────────────┘│   │└────────────┘│
                                         └──────────────┘   └──────────────┘


stores "group indexes"                     Each accumulator stores its state in
as (block_id, offset)                      fixed size blocks: a full block is
pairs into the block                       never resized; instead a new block
storage                                    is allocated as needed, and whole
                                           blocks can be emitted / freed
                                           one at a time

Why this is hard to fix

The reason this is so hard to implement is that the entire API is designed around a single usize group index that is assumed to be contiguous and directly addressable:

• GroupValues::intern assigns each distinct group key a dense index 0..n, and the hash table stores those raw indexes.
• GroupsAccumulator::update_batch receives group_indices: &[usize] and total_num_groups: usize; implementations index their state Vecs directly with the group index and grow them with a single resize(total_num_groups).
• EmitTo::First(n): after emitting the first n groups, every remaining group index is renumbered down by n — a notion that only makes sense when the state is one contiguous array.

For example GroupValues::intern:

pub trait GroupValues: Send {
    // Required methods
    fn intern(
        &mut self,
        cols: &[Arc<dyn Array>],
        groups: &mut Vec<usize>,   // <---- groups are identified by contiguous `usize`
    ) -> Result<(), DataFusionError>;

    ...
}

Because this assumption is spread across every GroupValues and GroupsAccumulator implementation — including user-defined aggregates and the FFI bindings — moving to a blocked (block_id, offset) index:

1. Potentially touches the whole ecosystem at once (aka is a massive change)
2. Likely involves an extra memory lookup in the hottest critical path: once to find the base pointer for the block_id and once to find the actual value within that block.

We have discussed ways to address both issues:

1. Incremental rollout: neither option found so far is great — supporting blocked and contiguous layouts in each implementation leads to the dual code paths and generics that made Intermediate result blocked approach to aggregation memory management #15591 so complex, while switching the index semantics outright is a breaking change that is very hard to stage incrementally (see the discussion on Intermediate result blocked approach to aggregation memory management #15591).
2. Different strategies for small and large aggregates: use direct indexing while the hash table is small (where the extra indirection would hurt most), and switch to two part (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:

PR	Year	What it showed	Outcome
#11758	2024	Generate GroupByHash output in multiple RecordBatches (@JasonLi-cn)	Closed unmerged
#11943	2024	First sketch of blocked management (@Rachelint)	Closed, superseded by #15591; motivated the aggregation fuzz test framework (#12114)
#15591	2025	Full blocked implementation (@Rachelint): supports_blocked_groups / alter_block_size trait additions, blocked PrimitiveGroupsAccumulator + GroupValuesPrimitive	Open. Extensive review concluded the dual-mode (blocked + contiguous) design is too complex, and some aggregates regress ~10%
#20964	2026	BatchedVec<T> bench (@Dandandan): O(1) per-block emission is achievable in small steps	Closed (proof of concept)
#22712	2026	PoC on refactored streams (@2010YOUY01): 10–16% faster at medium/high cardinality; memory curve becomes bell-shaped instead of monotonically growing; ClickBench Q5 +61% pending the skip-partial-aggregation fast path	Closed (proof of concept; demonstrated the #22710 refactor is necessary first)
#23274	2026	EmitTo::FirstBlock as an API-only first step (@hhhizzz)	Closed: the API should follow the blocked physical layout rather than precede it

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 GroupedHashAggregateStream into 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::NextBlock is being discussed on #15591. It is not yet a PR, but the work seems to be on the add-blocks-impl branch 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:

• Implement cache-efficient partial aggregation by Leis et al #20773

Related issues

• [EPIC] Split Aggregation Logic into Dedicated Streams #22710
• Manage group values and states by blocks in aggregation #11931
• Grouping operations on large datasets can overflow i32 offsets #23694
• Fixed blocks in aggregation state for better memory management #19649
• [EPIC] Eliminate Long Polls in HashAggregate via Chunked Storage and Incremental Emission #19906
• Track long-term replacement for destructive EmitTo::First in aggregation #23251
• Hash aggregation produces batches reporting huge memory size #22526
• Generate GroupByHash output in multiple RecordBatches rather than one large one #9562
• Grouped approx_distinct can overflow Arrow BinaryArray state offsets #24072
• GroupValuesRows::emit untracked decode buffer leaks past MemoryReservation #22739
• Long poll (async runtime stall) in HashAggregate when emitting large number of groups #18907

Related 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:

• TPCDS Q72 Fails with OffsetOverflowError(2147731589) (AQE) datafusion-ballista#1826 (reported by @milenkovicm): TPC-DS Q72 in Ballista fails with the same OffsetOverflowError, 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.
• Avoid concatenating record batches in joins to alleviate memory pressure #23031 (reported by @maxburke) tracks avoiding concatenating record batches in joins to alleviate memory pressure.
• Operators should respect batch_size instead of Emit::All #19481 (reported by @EmilyMatt) generalizes "operators should respect batch_size instead of emitting everything at once" across operators.

[comphead]

Thanks for summarizing it