alamb opened a new issue, #24704: URL: https://github.com/apache/datafusion/issues/24704
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 has been sent downstream. For a typical [two-stage (`Partial` → `Final`) 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 of state from those batches, so peak reaches partial + final ≈ 2 GB where ~1 GB should suffice (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, with a reproducer: `CUBE` + `COUNT(DISTINCT)` over ~384 byte string keys) - https://github.com/apache/datafusion-comet/pull/4791 (workaround by @comphead that Comet is shipping to unblock users: promote group keys to `LargeUtf8`/`LargeBinary` so the buffer uses `i64` offsets) 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 (or complex keys such as strings) is a single CPU-bound operation that can 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) ## 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. ## What is causing the problem [`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`). ```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, in the same order as above: 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<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. 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. ## 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<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 #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<T>` 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 | 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 - [ ] 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 [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<i32>`]: 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 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
