Skip to content

Unify the spilling aggregate streams behind one spill-replay driver #25537

Description

@jayzhan211

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:

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 [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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions