Skip to content

chore(SingleHashAggregateStream): refactor to async generator implementation #24016

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
buraksenn wants to merge 2 commits into
base: main
Choose a base branch
from
+87 −245
Open

chore(SingleHashAggregateStream): refactor to async generator implementation #24016

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
44cf67a
move single hash agg stream to generators and less state
buraksenn Jul 30, 2026
e9eaa1f
address PR review
buraksenn Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion datafusion/physical-plan/src/aggregates/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,7 +696,7 @@ impl From<StreamType> for SendableRecordBatchStream {
StreamType::PartialHash(stream) => Box::pin(stream),
StreamType::PartialReduceHash(stream) => Box::pin(stream),
StreamType::FinalHash(stream) => Box::pin(stream),
StreamType::SingleHash(stream) => Box::pin(stream),
StreamType::SingleHash(stream) => stream.into_stream(),
StreamType::OrderedPartialAggregate(stream) => stream.into_stream(),
StreamType::OrderedFinalAggregate(stream) => Box::pin(stream),
StreamType::GroupedHash(stream) => Box::pin(stream),
Expand Down Expand Up @@ -4085,7 +4085,17 @@ mod tests {
assert!(matches!(stream, StreamType::SingleHash(_)));
let stream: SendableRecordBatchStream = stream.into();
let output = collect(stream).await?;
assert_eq!(output.len(), 2);
assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);

let metrics = single.metrics().expect("aggregate metrics should exist");
assert_eq!(metrics.output_rows(), Some(3));
assert_eq!(
metrics
.sum(|metric| matches!(metric.value(), MetricValue::OutputBatches(_)))
.map(|value| value.as_usize()),
Some(2)
);
assert_snapshot!(batches_to_sort_string(&output), @r"
+---+--------+
| a | SUM(b) |
Expand Down
320 changes: 76 additions & 244 deletions datafusion/physical-plan/src/aggregates/single_stream.rs
View file
Open in desktop

@rluvaton rluvaton Jul 30, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would keep the code in their original functions but just change to async and remove state

Original file line number Diff line number Diff line change
Expand Up @@ -22,22 +22,20 @@
//!
//! See issue for details: <https://github.com/apache/datafusion/issues/22710>

use std::ops::ControlFlow;
use std::sync::Arc;
use std::task::{Context, Poll};

use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use datafusion_execution::TaskContext;
use datafusion_common::{DataFusionError, Result};
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
use datafusion_execution::{TaskContext, TryEmitter, async_try_stream};
use futures::stream::{Stream, StreamExt};

use super::AggregateExec;
use super::aggregate_hash_table::{AggregateHashTable, SingleMarker};
use crate::metrics::{BaselineMetrics, RecordOutput};
use crate::stream::EmptyRecordBatchStream;
use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream};
use crate::metrics::BaselineMetrics;
use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter};
use crate::{InputOrderMode, SendableRecordBatchStream};

/// Hash aggregation can run the full logical aggregation in one operator. This
/// stream implements the single stage for grouped hash aggregation.
Expand Down Expand Up @@ -68,67 +66,11 @@ pub(crate) struct SingleHashAggregateStream {
/// Memory reservation for group keys and accumulators.
reservation: MemoryReservation,

/// Tracks the high-level stream lifecycle. The hash table owns the lower-level
/// state for emitting output batches.
state: Option<SingleHashAggregateState>,
}

/// States for single hash aggregation processing.
// The typestate pattern mirrors the final stream and keeps the input/output
// semantics explicit for this mode.
enum SingleHashAggregateState {
ReadingInput {
hash_table: AggregateHashTable<SingleMarker>,
},
ProducingOutput {
hash_table: AggregateHashTable<SingleMarker>,
},
Done,
}

type SingleHashAggregatePoll = Poll<Option<Result<RecordBatch>>>;
type SingleHashAggregateStateTransition = ControlFlow<
(SingleHashAggregatePoll, SingleHashAggregateState),
SingleHashAggregateState,
>;

impl SingleHashAggregateState {
fn hash_table(&self) -> &AggregateHashTable<SingleMarker> {
match self {
Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => {
hash_table
}
Self::Done => unreachable!("Done state does not hold a hash table"),
}
}

fn hash_table_mut(&mut self) -> &mut AggregateHashTable<SingleMarker> {
match self {
Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => {
hash_table
}
Self::Done => unreachable!("Done state does not hold a hash table"),
}
}

fn into_hash_table(self) -> AggregateHashTable<SingleMarker> {
match self {
Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => {
hash_table
}
Self::Done => unreachable!("Done state does not hold a hash table"),
}
}

fn into_producing_output(self) -> Self {
Self::ProducingOutput {
hash_table: self.into_hash_table(),
}
}

fn into_done(self) -> Self {
Self::Done
}
/// The hash table owns the lower-level state for emitting output batches.
///
/// This is optional so `create_stream` can take ownership and control when
/// the table's memory is released.
hash_table: Option<AggregateHashTable<SingleMarker>>,
}

impl SingleHashAggregateStream {
Expand Down Expand Up @@ -164,7 +106,34 @@ impl SingleHashAggregateStream {
input,
baseline_metrics,
reservation,
state: Some(SingleHashAggregateState::ReadingInput { hash_table }),
hash_table: Some(hash_table),
})
}

pub(crate) fn into_stream(self) -> SendableRecordBatchStream {
let schema = Arc::clone(&self.schema);
let baseline_metrics = self.baseline_metrics.clone();
let stream =
Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream()));

Box::pin(ObservedStream::new(stream, baseline_metrics, None))
}

/// State transitions are implemented using the generator pattern; see the
/// comments in [`async_try_stream`].
///
/// Conceptually: ReadingInput -> ProducingOutput -> Done.
fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
async_try_stream(|emitter| async move {
let mut hash_table = self
.hash_table
.take()
.expect("SingleHashAggregateStream hash table should not be None");

self.handle_reading_input(&mut hash_table).await?;
self.handle_producing_output(hash_table, emitter).await?;

Ok(())
})
}

Expand All @@ -183,197 +152,60 @@ impl SingleHashAggregateStream {

/// Handle ReadingInput state - aggregate input batches into the hash table.
///
/// See comments at `poll_next()` for details.
///
/// Returns the next operator state with control flow decision.
fn handle_reading_input(
/// See comments at [`Self::create_stream`] for details.
async fn handle_reading_input(
&mut self,
cx: &mut Context<'_>,
mut original_state: SingleHashAggregateState,
) -> SingleHashAggregateStateTransition {
debug_assert!(matches!(
&original_state,
SingleHashAggregateState::ReadingInput { .. }
));
debug_assert!(original_state.hash_table().is_building());

match self.input.poll_next_unpin(cx) {
Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)),
// Get a new input batch, aggregate it in the hash table
Poll::Ready(Some(Ok(batch))) => {
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
let timer = elapsed_compute.timer();
let result = original_state.hash_table_mut().aggregate_batch(&batch);
timer.done();
hash_table: &mut AggregateHashTable<SingleMarker>,
) -> Result<()> {
debug_assert!(hash_table.is_building());
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();

if let Err(e) = result {
return ControlFlow::Break((
Poll::Ready(Some(Err(e))),
original_state,
));
}

if let Err(e) = self
.reservation
.try_resize(original_state.hash_table().memory_size())
{
return ControlFlow::Break((
Poll::Ready(Some(Err(e))),
original_state,
));
}

ControlFlow::Continue(original_state)
}
Poll::Ready(Some(Err(e))) => {
ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
}
// Input ends, move to output state
Poll::Ready(None) => {
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
let timer = elapsed_compute.timer();
let result = self.start_output(original_state.hash_table_mut());
timer.done();
while let Some(batch) = self.input.next().await.transpose()? {
let timer = elapsed_compute.timer();
hash_table.aggregate_batch(&batch)?;
timer.done();

match result {
Ok(()) => {
ControlFlow::Continue(original_state.into_producing_output())
}
Err(e) => {
ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
}
}
}
self.reservation.try_resize(hash_table.memory_size())?;
}

let timer = elapsed_compute.timer();
let result = self.start_output(hash_table);
timer.done();
result
}

/// Handle ProducingOutput state - emit final aggregate value batches.
///
/// See comments at `poll_next()` for details.
///
/// Returns the next operator state with control flow decision.
fn handle_producing_output(
/// See comments at [`Self::create_stream`] for details.
async fn handle_producing_output(
&mut self,
mut original_state: SingleHashAggregateState,
) -> SingleHashAggregateStateTransition {
debug_assert!(matches!(
&original_state,
SingleHashAggregateState::ProducingOutput { .. }
));
debug_assert!(!original_state.hash_table().is_building());
mut hash_table: AggregateHashTable<SingleMarker>,
mut emitter: TryEmitter<RecordBatch, DataFusionError>,
) -> Result<()> {
debug_assert!(!hash_table.is_building());

let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
let timer = elapsed_compute.timer();
let result = original_state.hash_table_mut().next_output_batch();
timer.done();
let mut timer = elapsed_compute.timer();

match result {
Ok(Some(batch)) => {
if let Err(e) = self
.reservation
.try_resize(original_state.hash_table().memory_size())
{
return ControlFlow::Break((
Poll::Ready(Some(Err(e))),
original_state,
));
}
debug_assert!(batch.num_rows() > 0);
let next_state = if original_state.hash_table().is_done() {
original_state.into_done()
} else {
original_state
};

ControlFlow::Break((
Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))),
next_state,
))
}
Ok(None) => {
let _ = self.reservation.try_resize(0);
ControlFlow::Continue(original_state.into_done())
}
Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)),
}
}
}
while let Some(batch) = hash_table.next_output_batch()? {
self.reservation.try_resize(hash_table.memory_size())?;
debug_assert!(batch.num_rows() > 0);

impl Stream for SingleHashAggregateStream {
type Item = Result<RecordBatch>;
if hash_table.is_done() {
drop(hash_table);
timer.done();

/// Entry point for the single hash aggregate state machine.
///
/// See comments in [`SingleHashAggregateStream`] for high-level ideas.
///
/// State transition graph:
///
/// ```text
/// (start)
/// -> ReadingInput
/// The stream starts by polling raw input rows and aggregating those
/// rows into the single-stage hash table.
///
/// ReadingInput
/// -> ReadingInput
/// Aggregate one raw input batch, update the inner aggregate hash
/// table, and continue with the next input batch.
///
/// -> ProducingOutput
/// Input was exhausted. Move to the next state to start outputting
/// final aggregate values.
///
/// ProducingOutput
/// -> ProducingOutput
/// One final output batch was yielded; repeat to continue producing
/// output incrementally.
///
/// -> Done
/// All final output was emitted.
///
/// Done
/// -> (end)
/// ```
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Self::Item>> {
loop {
let cur_state = self
.state
.take()
.expect("SingleHashAggregateStream state should not be None");

let next_state = match cur_state {
state @ SingleHashAggregateState::ReadingInput { .. } => {
self.handle_reading_input(cx, state)
}
state @ SingleHashAggregateState::ProducingOutput { .. } => {
self.handle_producing_output(state)
}
state @ SingleHashAggregateState::Done => {
let _ = self.reservation.try_resize(0);
self.state = Some(state);
return Poll::Ready(None);
}
};

match next_state {
ControlFlow::Continue(next_state) => {
self.state = Some(next_state);
continue;
}
ControlFlow::Break((poll, next_state)) => {
self.state = Some(next_state);
return poll;
}
emitter.emit(batch).await;

return Ok(());
}

timer.done();

emitter.emit(batch).await;
timer = elapsed_compute.timer();
}
}
}

impl RecordBatchStream for SingleHashAggregateStream {
fn schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
Ok(())
}
}
Loading