apache / apache/datafusion

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

Open
#24,704 16 comments 9 reactions 0 assignees View on GitHub
enhancement
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

Note: this issue consolidates and replaces https://github.com/apache/datafusion/issues/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][multi-phase], 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 https://github.com/apache/datafusion/pull/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`:
- https://github.com/apache/datafusion/issues/23694 (reported by @maxburke)
- https://github.com/apache/datafusion-comet/issues/4718 (report from Comet by @comphead)
- https://github.com/apache/datafusion-comet/pull/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`]:
- https://github.com/apache/datafusion/issues/22526 (reported by @ariel-miculas)
- https://github.com/apache/datafusion/issues/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:
- https://github.com/apache/datafusion/issues/19906 (reported by @ahmed-mez)

5. **Potential copying performance.** As groups accumulate, internal buffers [repeatedly double in size and copy all existing data][vec-growth] (up to 2 copies per element on average), which is likely expensive and cache/TLB unfriendly and could in theory be avoided.
- https://github.com/apache/datafusion/issues/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](https://datafusion.apache.org/blog/2023/08/05/datafusion_fast_grouping/)

[`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` (plus a null buffer), again indexed by a single group index (a `usize`).

```text
┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│┌────────────┐│ │┌────────────┐│ │┌────────────┐│
┌─────┐ ││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][emit-slice]. 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`] 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][duckdb-agg] 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][yjshen-suggestion].

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 `RecordBatch`es (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:
```text
┌──────────────┐ ┌──────────────┐
│┌────────────┐│ │┌────────────┐│
┌─────────┐ ││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][group-index-contiguous]:

- [`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 `Vec`s 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`][`EmitTo::First(n)`] — a notion that only makes sense when the state is one contiguous array.

For example [`GroupValues::intern`]:

```rust
pub trait GroupValues: Send {
// Required methods
fn intern(
&mut self,
cols: &[Arc],
groups: &mut Vec, // <---- 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 #15591 so complex, while switching the index semantics outright is a breaking change that is very hard to stage incrementally (see the discussion on https://github.com/apache/datafusion/pull/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 https://github.com/apache/datafusion/pull/15591#issuecomment-5333397682 and the earlier version of the same idea in https://github.com/apache/datafusion/pull/22712#issuecomment-4672476038.

## Past attempts and prototypes

There is a long and distinguished history of trying to address this problem:

| PR | Year | What it showed | Outcome |
|----|------|----------------|---------|
| https://github.com/apache/datafusion/pull/11758 | 2024 | Generate GroupByHash output in multiple `RecordBatch`es (@JasonLi-cn) | Closed unmerged |
| https://github.com/apache/datafusion/pull/11943 | 2024 | First sketch of blocked management (@Rachelint) | Closed, superseded by #15591; motivated the aggregation fuzz test framework (#12114) |
| https://github.com/apache/datafusion/pull/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% |
| https://github.com/apache/datafusion/pull/20964 | 2026 | `BatchedVec` bench (@Dandandan): O(1) per-block emission is achievable in small steps | Closed (proof of concept) |
| https://github.com/apache/datafusion/pull/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) |
| https://github.com/apache/datafusion/pull/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][take-n] and [renumbering every remaining group index][`EmitTo::First(n)`]. @ahmed-mez tried incremental emission with these mechanics in https://github.com/apache/datafusion/pull/19562 and measured it ~15x slower at high cardinality.

As part of https://github.com/apache/datafusion/issues/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 https://github.com/apache/datafusion/issues/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 https://github.com/apache/datafusion/pull/15591. It is not yet a PR, but the work seems to be on the [`add-blocks-impl` branch](https://github.com/rluvaton/datafusion/tree/add-blocks-impl) 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:
- https://github.com/apache/datafusion/issues/20773

## Related issues

- [x] https://github.com/apache/datafusion/issues/22710
- [ ] https://github.com/apache/datafusion/issues/11931
- [ ] https://github.com/apache/datafusion/issues/23694
- [ ] https://github.com/apache/datafusion/issues/19649
- [ ] https://github.com/apache/datafusion/issues/19906
- [ ] https://github.com/apache/datafusion/issues/23251
- [ ] https://github.com/apache/datafusion/issues/22526
- [ ] https://github.com/apache/datafusion/issues/9562
- [ ] https://github.com/apache/datafusion/issues/24072
- [ ] https://github.com/apache/datafusion/issues/22739
- [ ] https://github.com/apache/datafusion/issues/18907

[multi-phase]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.Accumulator.html#multi-phase-grouping
[vec-growth]: https://doc.rust-lang.org/std/vec/struct.Vec.html#capacity-and-reallocation
[emit-slice]: https://github.com/apache/datafusion/blob/d1fe98894238dd908a874a88c21c12a102ef5b92/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs#L826-L827
[take-n]: https://docs.rs/datafusion/latest/datafusion/physical_plan/aggregates/group_values/multi_group_by/trait.GroupColumn.html#tymethod.take_n
[group-index-contiguous]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.GroupsAccumulator.html#details
[duckdb-agg]: https://duckdb.org/2022/03/07/aggregate-hashtable
[yjshen-suggestion]: https://github.com/apache/datafusion/pull/6800#discussion_r1251142165
[`GroupedHashAggregateStream`]: https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
[`GroupValues`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/aggregates/group_values/trait.GroupValues.html
[`GroupValues::intern`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/aggregates/group_values/trait.GroupValues.html#tymethod.intern
[`GroupValuesPrimitive`]: https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
[`PrimitiveGroupValueBuilder`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/aggregates/group_values/multi_group_by/primitive/struct.PrimitiveGroupValueBuilder.html
[`ByteGroupValueBuilder`]: https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs
[`ByteGroupValueBuilder`]: https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs
[`GroupsAccumulator`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.GroupsAccumulator.html
[`GroupsAccumulator::update_batch`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.GroupsAccumulator.html#tymethod.update_batch
[`PrimitiveGroupsAccumulator`]: https://docs.rs/datafusion-functions-aggregate-common/latest/datafusion_functions_aggregate_common/aggregate/groups_accumulator/prim_op/struct.PrimitiveGroupsAccumulator.html
[`EmitTo::All`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.EmitTo.html#variant.All
[`EmitTo::First(n)`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.EmitTo.html#variant.First
[`get_array_memory_size()`]: https://docs.rs/arrow/latest/arrow/array/struct.RecordBatch.html#method.get_array_memory_size
[`RepartitionExec`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/repartition/struct.RepartitionExec.html
[`TopK`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/struct.TopK.html

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

- https://github.com/apache/datafusion-ballista/issues/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.
- https://github.com/apache/datafusion/issues/23031 (reported by @maxburke) tracks avoiding concatenating record batches in joins to alleviate memory pressure.
- https://github.com/apache/datafusion/issues/19481 (reported by @EmilyMatt) generalizes "operators should respect `batch_size` instead of emitting everything at once" across operators.

Contributor guide

Open the contributing guide

Research direction

Start with the named GroupedHashAggregateStream, GroupValues::intern, GroupsAccumulator::update_batch, and EmitTo entry points, tracing how contiguous group indices and state storage flow through their implementations. Done means group keys and accumulator state use fixed-size blocks, emission releases one block at a time, memory accounting reflects individual outputs, and the design covers user-defined aggregates and FFI bindings.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data-engineering, databases
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.