You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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:
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().
SpillableAggregateTable — a small object-safe trait capturing exactly those per-batch seams, implemented in the existing per-marker table files.
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.
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 chore(OrderedFinalAggregateStream): refactor to async generator #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 ([EPIC] Reduce binary size #24727). Switching later is mechanical.
Ordering vs. removal of the legacy stream. I'd like these to land beforegrouped_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.
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 [EPIC] Use blocked / chunked memory management in hash aggregation #24704 is reworking, so I'd leave it as a later, separately-benchmarked step.
Is your feature request related to a problem or challenge?
#22710 split
GroupedHashAggregateStreaminto 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 —— in
single_stream.rs,ordered_single_stream.rs,ordered_final_stream.rsand the Final half ofhash_stream.rs. Concretely:SingleSpillContext,OrderedSingleSpillContext,OrderedFinalSpillContextandFinalSpillContexthave the same seven fields and the samespill_table/into_replay_streambodies (all four end in the sameStreamingMergeBuilderchain feedingOrderedFinalAggregateStream::new_with_input_and_metrics).ReadingInput,Spilling,ProducingOutput,PreparingMergeInput,MergingSpills,Done,Error), the same handler names and the sameControlFlowalias pattern — roughly 2,600 non-test lines for one state machine written three times.Errorstate orDoneis reused.#23974 is converting streams to async generators, and #24008 / #24016 do that for
OrderedFinalandSingleindividually. 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()vsinput_done(),is_done()vsis_empty(), which group count feeds the spill-index overhead, where replay metrics come from, …) plus constructor-time configuration. That suggests:AggregateSpill— one non-generic spill context (spill_state_batch(batch),has_spills(),into_replay_stream(..)). Non-generic because everyspill_tableonly needstable.take_state_batch(). The spill sort key isorder_indices ++ remaining group columns, of whichLinearis the empty-prefix case; the replay config mapsSingle → Finalwithgroup_by.as_final().SpillableAggregateTable— a small object-safe trait capturing exactly those per-batch seams, implemented in the existing per-marker table files.FinalHashAggregateStream::create_stream, replacing the three hand-rolled state machines and the Final generator.StreamTypevariants 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 + migrateFinalHash(already a generator, so the smallest semantic diff), (3)SingleHash, (4)OrderedSingle, thenOrderedFinal(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:
OrderedFinalAggregateStream): refactor to async generator #24008 / chore(SingleHashAggregateStream): refactor to async generator implementation #24016. This would supersede them. Happy for their authors to take the corresponding stages if they prefer.record_outputat in-memory emit sites and forward replay batches unrecorded, notObservedStream— the replay stream sharesBaselineMetrics, so wrapping double-countsoutput_rows(raised in review of chore(OrderedFinalAggregateStream): refactor to async generator #24008).ordered_partial_stream.rscurrently usesObservedStream; 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 existingdyn GroupValues/dyn GroupsAccumulatorcalls, anddynkeeps one copy of the async state machine instead of four ([EPIC] Reduce binary size #24727). Switching later is mechanical.grouped_hash_stream.rsandenable_migration_aggregateare removed: the aggregation fuzzer's baseline context runs with the flag off, so the legacy stream is a free differential oracle for the migration.core/tests/memory_limit) and thewith_can_spill(true)workaround for External sort failing with non-spillable operators as input (RepartitionExec) #17334 on the replay stream.Describe alternatives you've considered
AggregateHashTableandOrderedAggregateTable(they overlap heavily, andGroupOrdering::Nonealready exists). Probably worthwhile, but it touches the output-materialisation code that [EPIC] Use blocked / chunked memory management in hash aggregation #24704 is reworking, so I'd leave it as a later, separately-benchmarked step.Additional context
Follow-up to #22710 (closed); related to #23974.