From 44cf67add2668d2c8cfaa24f753ea15f8f6a4923 Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Thu, 30 Jul 2026 22:26:41 +0300 Subject: [PATCH 1/2] move single hash agg stream to generators and less state --- .../physical-plan/src/aggregates/mod.rs | 12 +- .../src/aggregates/single_stream.rs | 316 ++++-------------- 2 files changed, 67 insertions(+), 261 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fbc..82b8681efbf32 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -696,7 +696,7 @@ impl From 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), @@ -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::(), 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) | diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 886ffdd3a0b99..7278a2591a72d 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -22,22 +22,20 @@ //! //! See issue for details: -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_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, 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::{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. @@ -68,67 +66,8 @@ 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, -} - -/// 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, - }, - ProducingOutput { - hash_table: AggregateHashTable, - }, - Done, -} - -type SingleHashAggregatePoll = Poll>>; -type SingleHashAggregateStateTransition = ControlFlow< - (SingleHashAggregatePoll, SingleHashAggregateState), - SingleHashAggregateState, ->; - -impl SingleHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { - 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 { - 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 { - 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 - } + /// Hash table containing group keys and accumulators. + hash_table: AggregateHashTable, } impl SingleHashAggregateStream { @@ -164,216 +103,73 @@ impl SingleHashAggregateStream { input, baseline_metrics, reservation, - state: Some(SingleHashAggregateState::ReadingInput { hash_table }), + hash_table, }) } - /// Moves the aggregate hash table's inner state to `Outputting`. - /// - /// The caller guarantees that input is fully consumed, so this function can - /// eagerly release the input stream. - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - hash_table.start_output() + 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)) } - /// Handle ReadingInput state - aggregate input batches into the hash table. - /// - /// See comments at `poll_next()` for details. + /// State transitions are implemented using the generator pattern; see the + /// comments in [`async_try_stream`]. /// - /// Returns the next operator state with control flow decision. - 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(); - - 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()) + /// Conceptually: ReadingInput -> ProducingOutput -> Done. + fn create_stream(self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + let Self { + mut input, + baseline_metrics, + reservation, + mut hash_table, + .. + } = self; + let elapsed_compute = baseline_metrics.elapsed_compute().clone(); + + debug_assert!(hash_table.is_building()); + while let Some(batch) = input.next().await.transpose()? { { - 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(); - - match result { - Ok(()) => { - ControlFlow::Continue(original_state.into_producing_output()) - } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + let _timer = elapsed_compute.timer(); + hash_table.aggregate_batch(&batch)?; } + reservation.try_resize(hash_table.memory_size())?; } - } - } - /// 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( - &mut self, - mut original_state: SingleHashAggregateState, - ) -> SingleHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - SingleHashAggregateState::ProducingOutput { .. } - )); - debug_assert!(!original_state.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(); - - 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 - }; + { + let _timer = elapsed_compute.timer(); - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) + // Input is exhausted. Release the upstream pipeline before draining output. + drop(input); + hash_table.start_output()?; } - 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)), - } - } -} -impl Stream for SingleHashAggregateStream { - type Item = Result; + debug_assert!(!hash_table.is_building()); + loop { + let next_batch = { + let _timer = elapsed_compute.timer(); + hash_table.next_output_batch()? + }; + let Some(batch) = next_batch else { + return Ok(()); + }; - /// 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> { - loop { - let cur_state = self - .state - .take() - .expect("SingleHashAggregateStream state should not be None"); + reservation.try_resize(hash_table.memory_size())?; + debug_assert!(batch.num_rows() > 0); - let next_state = match cur_state { - state @ SingleHashAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) + if hash_table.is_done() { + drop(hash_table); + reservation.free(); + emitter.emit(batch).await; + return Ok(()); } - 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; } - } - } -} - -impl RecordBatchStream for SingleHashAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) + }) } } From e9eaa1ffad967f17dac6dc7c5808216532eab8de Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Thu, 30 Jul 2026 23:50:11 +0300 Subject: [PATCH 2/2] address PR review --- .../src/aggregates/single_stream.rs | 138 +++++++++++------- 1 file changed, 87 insertions(+), 51 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 7278a2591a72d..f109fcef95149 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -26,15 +26,15 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_execution::{TaskContext, async_try_stream}; +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; -use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, SendableRecordBatchStream}; /// Hash aggregation can run the full logical aggregation in one operator. This @@ -66,8 +66,11 @@ pub(crate) struct SingleHashAggregateStream { /// Memory reservation for group keys and accumulators. reservation: MemoryReservation, - /// Hash table containing group keys and accumulators. - hash_table: AggregateHashTable, + /// 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>, } impl SingleHashAggregateStream { @@ -103,7 +106,7 @@ impl SingleHashAggregateStream { input, baseline_metrics, reservation, - hash_table, + hash_table: Some(hash_table), }) } @@ -120,56 +123,89 @@ impl SingleHashAggregateStream { /// comments in [`async_try_stream`]. /// /// Conceptually: ReadingInput -> ProducingOutput -> Done. - fn create_stream(self) -> impl Stream> { - async_try_stream(|mut emitter| async move { - let Self { - mut input, - baseline_metrics, - reservation, - mut hash_table, - .. - } = self; - let elapsed_compute = baseline_metrics.elapsed_compute().clone(); - - debug_assert!(hash_table.is_building()); - while let Some(batch) = input.next().await.transpose()? { - { - let _timer = elapsed_compute.timer(); - hash_table.aggregate_batch(&batch)?; - } - reservation.try_resize(hash_table.memory_size())?; - } + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|emitter| async move { + let mut hash_table = self + .hash_table + .take() + .expect("SingleHashAggregateStream hash table should not be None"); - { - let _timer = elapsed_compute.timer(); + self.handle_reading_input(&mut hash_table).await?; + self.handle_producing_output(hash_table, emitter).await?; - // Input is exhausted. Release the upstream pipeline before draining output. - drop(input); - hash_table.start_output()?; - } + Ok(()) + }) + } + + /// Moves the aggregate hash table's inner state to `Outputting`. + /// + /// The caller guarantees that input is fully consumed, so this function can + /// eagerly release the input stream. + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + hash_table.start_output() + } + + /// Handle ReadingInput state - aggregate input batches into the hash table. + /// + /// See comments at [`Self::create_stream`] for details. + async fn handle_reading_input( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + debug_assert!(hash_table.is_building()); + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + + while let Some(batch) = self.input.next().await.transpose()? { + let timer = elapsed_compute.timer(); + hash_table.aggregate_batch(&batch)?; + timer.done(); + + 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 [`Self::create_stream`] for details. + async fn handle_producing_output( + &mut self, + mut hash_table: AggregateHashTable, + mut emitter: TryEmitter, + ) -> Result<()> { + debug_assert!(!hash_table.is_building()); + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + while let Some(batch) = hash_table.next_output_batch()? { + self.reservation.try_resize(hash_table.memory_size())?; + debug_assert!(batch.num_rows() > 0); - debug_assert!(!hash_table.is_building()); - loop { - let next_batch = { - let _timer = elapsed_compute.timer(); - hash_table.next_output_batch()? - }; - let Some(batch) = next_batch else { - return Ok(()); - }; - - reservation.try_resize(hash_table.memory_size())?; - debug_assert!(batch.num_rows() > 0); - - if hash_table.is_done() { - drop(hash_table); - reservation.free(); - emitter.emit(batch).await; - return Ok(()); - } + if hash_table.is_done() { + drop(hash_table); + timer.done(); emitter.emit(batch).await; + + return Ok(()); } - }) + + timer.done(); + + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + Ok(()) } }