jayzhan211 opened a new pull request, #25246: URL: https://github.com/apache/datafusion/pull/25246
## Which issue does this PR close? - Part of #24704 — the first step toward blocked group state, as discussed there. Related: #22712, #7065, #19562, #23178. This branch carries three commits. The first (`perf: keep the key in the hash entry of GroupValuesPrimitive`) is independent of the other two and will be filed as its own small PR; it is included here because the blocked group values below rely on it to keep the probe off the value storage. The remaining two commits are the subject of this PR and are best read one at a time. ## Rationale for this change ### Commit 2 — blocked group state addressed by an independent resolve pass Every attempt to store hash-aggregation state in blocks — so it can be emitted, freed and spilled one block at a time — has regressed at high cardinality: #22712, and 1.5–1.7x at 10M groups when retried for this work, for any block size. The reason is mechanical and is neither emission nor the block size: `blocks[g >> SHIFT][g & MASK]` is a **dependent load**. The block pointer has to resolve before the state access can issue, so the update loop keeps fewer cache misses in flight, and a flat `values[g]` loop wins on memory-level parallelism alone. `intern` already produces the whole batch's group indices before any accumulator runs, so the addresses can be resolved in a separate, independent pass and the update loop can run through them. Because that pass knows every address without a dependent load, it can also prefetch each element — which does nothing in the naive loop, where the address itself is the dependency. Interleaved standalone measurement of the update loop (ns/row, M4 Pro, `i64` sum, uniformly random groups; the loop is bound by how many rows fit in the reorder window — even a bounds check is worth 27% here, which is why the resolve pass lives inside `accumulate` rather than in the per-row callback): | groups | flat, unchecked (today) | naive blocked | resolve, no prefetch | **resolve + prefetch** | |---|---|---|---|---| | 1M | 0.87 | 1.18 | 1.09 | 1.11 (single block ⇒ today's loop is used) | | 10M | 2.61 | 4.91 (1.9x) | 2.91 (1.11x) | **2.59 (0.99x)** | | 20M | 3.21 | 5.36 | 3.55 | 3.35 (1.04x) | | 40M | 3.60 | 5.57 | 3.92 | 3.71 (1.03x) | Tables of up to 2^20 groups are one block and run exactly today's loop. With this in place, storing state in blocks costs nothing on the hot loop, and block-wise emission (commit 3) is an output-side change that cannot regress the update loop. ### Commit 3 — emit and free the output one block at a time Hash aggregation materializes its whole output at once and holds it, fully reserved, until the last slice has been polled. When the consumer builds its own state from that output (a second aggregation, a join build side), the peak is *producer output + consumer state*, and under a memory limit the consumer spills even though the producer's memory is dead weight by then. With the state in blocks, `EmitTo::First(BLOCK_LEN)` is a whole-block move with no renumbering of what remains, so the output side can hand over one block, slice it per poll, release it, and emit the next: the reservation shrinks block by block. Nested aggregation (`GROUP BY k` over the output of `GROUP BY k`, 10M random groups, 40M rows, 1 partition, `datafusion-cli`, greedy pool), bytes spilled by the outer / inner aggregate at a fixed limit: | limit | main | commit 2 only | **commits 2 + 3** | |---|---|---|---| | 640 MB | 310 MB / 1241 MB | 310 MB / 0 | **0 / 0** | | 768 MB | 310 / 1241 | 310 / 0 | **0 / 0** | | 896 MB | 310 / 0 | 0 / 0 | 0 / 0 | The inner's 1241 → 0 is commit 2 (no doubling slack in the reservation); the outer's 310 → 0 is commit 3 (the inner's output is released while the outer builds). Put differently, the smallest limit at which this query runs without touching disk goes from over 1 GB to 640 MB. Note on arbitration: at exactly 512 MB under the greedy pool the inner's spill *replay* can be starved once the outer has grown into the pool (main survives there only because its coarser doubling OOMs the outer earlier); the fair pool passes at 448–576 MB and greedy passes from 528 MB — a pool-fairness boundary, not specific to this change. ## What changes are included in this PR? Commit 2: - `BlockedVec<T, SHIFT = 20>` (`functions-aggregate-common::groups_accumulator::blocks`): per-group state in blocks of 2^20; `as_single_block_mut` exposes a flat slice for the single-block case; `AddressResolver::resolve` is the independent pass (with prefetch); `take_all` / `take_first` / `get` / `resize` / `allocated_size`. Blocks grow by doubling, so only the last block carries slack. - `accumulate_resolved` / `NullState::accumulate_resolved`: `accumulate` with a resolve callback per 256-row chunk; all four null/filter arms zip the resolved addresses with the data, so the update loop stays the lean one. `accumulate` itself is untouched. - `PrimitiveGroupsAccumulator` (`SUM`/`MIN`/`MAX`/`BIT_*` on primitives) and `GroupValuesPrimitive` keep their values in a `BlockedVec`. - `benches/blocked_update.rs`: flat vs naive blocked vs resolved at 100k/1M/10M, plus the real accumulator on random (seen-bit path) and first-appearance-dense (partial-stage path) indices. Commit 3: - `GroupsAccumulator::block_len() -> Option<usize>` (default `None`); `PrimitiveGroupsAccumulator` reports its block length. - `GroupValues::block_len()` and `emit_block()` (defaults `None` / not implemented); `GroupValuesPrimitive::emit_block` moves the first block out and drops the hash index instead of renumbering it, since no `intern` follows. - `AggregateHashTable::next_output_batch_inner`: new `OutputtingBlock` state; when every column agrees on a block length, output is emitted a block at a time and `memory_size()` is the remaining state plus the current block. Otherwise today's materialize-all path is used unchanged, so accumulators without block storage (`count`, `avg`, …) behave exactly as before. All three hash streams (partial, final, single) go through it. Commit 1 (to be split out): `GroupValuesPrimitive` stores `(group_index, key)` instead of `(group_index, hash)`, so a probe never reads the values (0.90–0.93x on `GROUP BY l_partkey`); for natives wider than 8 bytes the entry grows to 32 bytes (+13% RSS measured on a Decimal128 key with 10M groups, at unchanged time). No behavior change anywhere; the two new trait methods have defaults, `EmitTo` is unchanged. ## Measured against main Binaries built in separate target dirs, interleaved runs: - Real `PrimitiveGroupsAccumulator<Int64>::update_batch`: seen-bit path 1.00 / 1.00 / **0.90** at 100k / 1M / 10M groups; dense path 1.00 / **0.94** / 0.98–1.02 at 1M / 10M / 20M. - `datafusion-cli`: single-block queries 0.99–1.03x; 10M-group queries 1.00–1.02x with max RSS 6–16% lower; nullable input 0.99x, `FILTER` 1.02x, Float64 keys 1.01x, `count`+`avg` (old output path) 0.96x; SF10 `l_orderkey` (15M clustered groups) 0.98x with RSS 4.7 → 3.8 GB; SF10 `l_partkey` (2M random) 0.90x. - TPC-H SF10 (`dfbench tpch`, 4 runs, min per query): total **0.98x**; q17 0.91x; q18 (`GROUP BY l_orderkey`, ~1.25M groups per partition, the two-block band) +2.3%, the one recurring small regression. - Not run: ClickBench (no `hits.parquet` locally); it is the suite where high-cardinality primitive `GROUP BY` shows most, so a run there would be welcome. ## What is the testing strategy for this PR? - `blocks::tests`: indexing/resize across blocks, `take_all`/`take_first` order and packing, whole-block moves, growth by doubling, resolver addresses. - `accumulate::test::accumulate_resolved_matches_accumulate`: same rows in the same order as `accumulate` for every null/filter combination across chunk boundaries. - `prim_op::tests::multi_block_groups_update_and_emit_in_order`, `primitive::tests::intern_and_emit_across_blocks`, `primitive::tests::emit_block_walks_the_blocks_in_order`. - `hash_stream::tests::test_single_hash_stream_releases_state_block_by_block`: two blocks of `sum(Int64)`; group order preserved across the block boundary, the reservation while the second block is sliced is < 1/4 of what it was during the first, and 0 at the end. - Extended test suite and full sqllogictest suite green. ## Are there any user-facing changes? No. New default trait methods on `GroupsAccumulator` and `GroupValues`; existing implementors compile and behave unchanged. https://claude.ai/code/session_01N4Hye6EyCKWzNDjyoBAozd -- 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]
