jayzhan211 opened a new issue, #25537: URL: https://github.com/apache/datafusion/issues/25537
### Is your feature request related to a problem or challenge? #22710 split `GroupedHashAggregateStream` into dedicated streams over a shared core (`AggregateHashTable<M>` / `OrderedAggregateTable<M>`). The core is nicely factored, but the *driver* layer on top of it is not: the four streams that can spill each re-implement the same algorithm — > read input → on OOM spill the table as one sorted run → at end of input merge the runs → replay them through an ordered final aggregation — in `single_stream.rs`, `ordered_single_stream.rs`, `ordered_final_stream.rs` and the Final half of `hash_stream.rs`. Concretely: - `SingleSpillContext`, `OrderedSingleSpillContext`, `OrderedFinalSpillContext` and `FinalSpillContext` have the same seven fields and the same `spill_table` / `into_replay_stream` bodies (all four end in the same `StreamingMergeBuilder` chain feeding `OrderedFinalAggregateStream::new_with_input_and_metrics`). - Three of the four are hand-rolled poll state machines with the same states (`ReadingInput`, `Spilling`, `ProducingOutput`, `PreparingMergeInput`, `MergingSpills`, `Done`, `Error`), the same handler names and the same `ControlFlow` alias pattern — roughly 2,600 non-test lines for one state machine written three times. - The copies have already drifted in ways that look accidental: whether an OOM without a spill context gets extra error context, whether OOM on an empty table is an internal error or the OOM itself, whether there is an `Error` state or `Done` is reused. #23974 is converting streams to async generators, and #24008 / #24016 do that for `OrderedFinal` and `Single` individually. Done one by one, we end up with the same generator written three times instead of the same state machine written three times. ### Describe the solution you'd like Looking at where the four drivers *really* differ, it is nine small per-batch decisions (soft group limit, early emit for ordered input, `start_output()` vs `input_done()`, `is_done()` vs `is_empty()`, which group count feeds the spill-index overhead, where replay metrics come from, …) plus constructor-time configuration. That suggests: 1. **`AggregateSpill`** — one non-generic spill context (`spill_state_batch(batch)`, `has_spills()`, `into_replay_stream(..)`). Non-generic because every `spill_table` only needs `table.take_state_batch()`. The spill sort key is `order_indices ++ remaining group columns`, of which `Linear` is the empty-prefix case; the replay config maps `Single → Final` with `group_by.as_final()`. 2. **`SpillableAggregateTable`** — a small object-safe trait capturing exactly those per-batch seams, implemented in the existing per-marker table files. 3. **One generator-style driver**, modelled on today's `FinalHashAggregateStream::create_stream`, replacing the three hand-rolled state machines and the Final generator. `StreamType` variants stay as they are (planner tests match on them); only their payload becomes the shared driver. Proposed as a sequence of behaviour-preserving PRs: (1) `AggregateSpill`, (2) trait + driver + migrate `FinalHash` (already a generator, so the smallest semantic diff), (3) `SingleHash`, (4) `OrderedSingle`, then `OrderedFinal` (the replay target), (5) normalise the accidental divergences separately since that changes user-visible error text. Roughly −2,100 lines net. Things I would like feedback on before starting: - **Relationship to #24008 / #24016.** This would supersede them. Happy for their authors to take the corresponding stages if they prefer. - **Metrics convention.** The driver would use explicit `record_output` at in-memory emit sites and forward replay batches unrecorded, not `ObservedStream` — the replay stream shares `BaselineMetrics`, so wrapping double-counts `output_rows` (raised in review of #24008). `ordered_partial_stream.rs` currently uses `ObservedStream`; it has no replay path so it is correct today, but it would be good to settle on one convention. - **`Box<dyn SpillableAggregateTable>` vs a generic driver.** All trait calls are per batch, next to existing `dyn GroupValues` / `dyn GroupsAccumulator` calls, and `dyn` keeps one copy of the async state machine instead of four (#24727). Switching later is mechanical. - **Ordering vs. removal of the legacy stream.** I'd like these to land *before* `grouped_hash_stream.rs` and `enable_migration_aggregate` are removed: the aggregation fuzzer's baseline context runs with the flag off, so the legacy stream is a free differential oracle for the migration. - Preserved verbatim: memory consumer names (asserted in `core/tests/memory_limit`) and the `with_can_spill(true)` workaround for #17334 on the replay stream. ### Describe alternatives you've considered - An enum over the four table types instead of a trait: nine methods × four variants of match boilerplate for no gain at per-batch call frequency. - Going further and merging `AggregateHashTable` and `OrderedAggregateTable` (they overlap heavily, and `GroupOrdering::None` already exists). Probably worthwhile, but it touches the output-materialisation code that #24704 is reworking, so I'd leave it as a later, separately-benchmarked step. ### Additional context Follow-up to #22710 (closed); related to #23974. -- 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]
