From b8582d74e603f37579fe02e810cd229792e481ce Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:19:07 +0300 Subject: [PATCH 1/4] chore(`FinalHashAggregateStream`): convert to async generators and cleanup --- .../src/aggregates/hash_stream.rs | 343 ++++++------------ .../physical-plan/src/aggregates/mod.rs | 2 +- 2 files changed, 107 insertions(+), 238 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index e7f0f075b33a5..91b19032c2530 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -31,9 +31,9 @@ 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; @@ -44,7 +44,7 @@ use super::skip_partial::SkipAggregationProbe; use crate::metrics::{ BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, }; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; /// Hash aggregation is implemented in two stages: partial and final. This @@ -204,67 +204,10 @@ pub(crate) struct FinalHashAggregateStream { /// See comments for the same variable in [`PartialHashAggregateStream`]. group_values_soft_limit: Option, - /// Tracks the high-level stream lifecycle. The hash table owns the lower-level - /// state for emitting output batches. - state: Option, -} - -/// States for final hash aggregation processing. -// The typestate pattern is used in case the inner logic becomes more complex in -// the future. -enum FinalHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - }, - ProducingOutput { - hash_table: AggregateHashTable, - }, - Done, -} - -type FinalHashAggregatePoll = Poll>>; -type FinalHashAggregateStateTransition = ControlFlow< - (FinalHashAggregatePoll, FinalHashAggregateState), - FinalHashAggregateState, ->; - -impl FinalHashAggregateState { - 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 - } + /// The hash table owns the lower-level state for emitting output batches. + /// + /// This is option since it will be taken on [`Self::create_stream`] to control the memory + hash_table: Option>, } impl PartialHashAggregateStream { @@ -777,153 +720,22 @@ impl FinalHashAggregateStream { baseline_metrics, reservation, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), - state: Some(FinalHashAggregateState::ReadingInput { hash_table }), + hash_table: Some(hash_table), }) } - /// See comments in [`Self::group_values_soft_limit`] for details. - fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { - self.group_values_soft_limit - .is_some_and(|limit| limit <= hash_table.building_group_count()) - } - - 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_clone = Arc::clone(&self.schema); - /// Handle ReadingInput state - aggregate partial state 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( - &mut self, - cx: &mut Context<'_>, - mut original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - FinalHashAggregateState::ReadingInput { .. } + let cloned_metrics = self.baseline_metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), )); - debug_assert!(original_state.hash_table().is_building()); - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), - 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 self.hit_soft_group_limit(original_state.hash_table()) { - let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); - timer.done(); - - if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - return ControlFlow::Continue(original_state.into_producing_output()); - } - - 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)) - } - 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)) - } - } - } - } + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) } - /// 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: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - FinalHashAggregateState::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)) => { - let _ = self - .reservation - .try_resize(original_state.hash_table().memory_size()); - 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)), - } - } -} - -impl Stream for FinalHashAggregateStream { - type Item = Result; - /// Entry point for the final hash aggregate state machine. /// /// See comments in [`FinalHashAggregateStream`] for high-level ideas. @@ -956,47 +768,104 @@ impl Stream for FinalHashAggregateStream { /// Done /// -> (end) /// ``` - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - loop { - let cur_state = self - .state + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|emitter| async move { + let mut hash_table = self + .hash_table .take() - .expect("FinalHashAggregateStream state should not be None"); + .expect("hash_table should not be None"); + self.handle_reading_input(&mut hash_table).await?; - let next_state = match cur_state { - state @ FinalHashAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ FinalHashAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(state) - } - state @ FinalHashAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Poll::Ready(None); - } - }; + self.produce_output(hash_table, emitter).await?; - 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; - } + Ok(()) + }) + } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= hash_table.building_group_count()) + } + + 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 partial state batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + 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)?; + + if self.hit_soft_group_limit(hash_table) { + self.start_output(hash_table)?; + + // Into producing output + return Ok(()); } + + self.reservation.try_resize(hash_table.memory_size())?; } + + let _timer = elapsed_compute.timer(); + self.start_output(hash_table)?; + + Ok(()) } -} -impl RecordBatchStream for FinalHashAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) + /// emit final aggregate value batches. + /// + /// See comments at [`Self::create_stream`] for details. + async fn produce_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()? { + let _ = self.reservation.try_resize(hash_table.memory_size()); + + debug_assert!(batch.num_rows() > 0); + + if hash_table.is_done() { + drop(hash_table); + self.reservation.try_resize(0)?; + timer.done(); + + emitter.emit(batch).await; + + return Ok(()); + } + + timer.done(); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + drop(hash_table); + self.reservation.try_resize(0)?; + + Ok(()) } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fbc..69c7ee05fa1ac 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -695,7 +695,7 @@ impl From for SendableRecordBatchStream { StreamType::AggregateStream(stream) => Box::pin(stream), StreamType::PartialHash(stream) => Box::pin(stream), StreamType::PartialReduceHash(stream) => Box::pin(stream), - StreamType::FinalHash(stream) => Box::pin(stream), + StreamType::FinalHash(stream) => stream.into_stream(), StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), From 8a49479e3d40a4f39307f8ac49d3f44f47d87e6f Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:45:28 +0300 Subject: [PATCH 2/4] chore(`PartialHashAggregateStream`): convert to async generators and cleanup --- .../src/aggregates/hash_stream.rs | 570 ++++++------------ .../physical-plan/src/aggregates/mod.rs | 2 +- 2 files changed, 178 insertions(+), 394 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 91b19032c2530..7b10a09752367 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -25,9 +25,7 @@ //! //! 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; @@ -41,11 +39,9 @@ use super::aggregate_hash_table::{ AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, }; use super::skip_partial::SkipAggregationProbe; -use crate::metrics::{ - BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, -}; +use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, SpillMetrics}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; -use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; +use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Hash aggregation is implemented in two stages: partial and final. This /// stream implements the partial stage. @@ -132,56 +128,8 @@ pub(crate) struct PartialHashAggregateStream { /// be empty. See struct comments for details. group_values_soft_limit: Option, - /// Tracks the high-level stream lifecycle. The hash table owns the lower-level - /// state for emitting output batches. - state: Option, -} - -/// States for partial hash aggregation processing. -enum PartialHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - }, - ProducingOutput { - hash_table: AggregateHashTable, - /// If `None`, partial skip was never triggered and this state will - /// finish in `Done`. If `Some`, partial skip has triggered and the - /// stream will move to `SkippingAggregation` after these accumulated - /// groups are emitted. - skip_hash_table: Option>, - }, - SkippingAggregation { - hash_table: AggregateHashTable, - }, - Done, -} - -type PartialHashAggregatePoll = Poll>>; -type PartialHashAggregateStateTransition = ControlFlow< - (PartialHashAggregatePoll, PartialHashAggregateState), - PartialHashAggregateState, ->; - -impl PartialHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { - match self { - Self::ReadingInput { hash_table } - | Self::ProducingOutput { hash_table, .. } => hash_table, - Self::SkippingAggregation { .. } | Self::Done => { - unreachable!("state does not hold a partial hash table") - } - } - } - - fn hash_table_mut(&mut self) -> &mut AggregateHashTable { - match self { - Self::ReadingInput { hash_table } - | Self::ProducingOutput { hash_table, .. } => hash_table, - Self::SkippingAggregation { .. } | Self::Done => { - unreachable!("state does not hold a partial hash table") - } - } - } + /// The hash table owns the lower-level state for emitting output batches. + hash_table: Option>, } /// Hash aggregation is implemented in two stages: partial and final. This @@ -270,7 +218,101 @@ impl PartialHashAggregateStream { reduction_factor, skip_aggregation_probe, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), - state: Some(PartialHashAggregateState::ReadingInput { hash_table }), + hash_table: Some(hash_table), + }) + } + + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema_clone = Arc::clone(&self.schema); + + let cloned_metrics = self.baseline_metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + + /// Entry point for the partial hash aggregate state machine. + /// + /// See comments in [`PartialHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling input and aggregating batches into the + /// in-memory hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one batch, update the inner aggregate hash table, and + /// continue with the next input batch. + /// -> ProducingOutput(skip=None) + /// Input was exhausted, or the soft group limit was reached. Move to + /// the next state to start outputting. + /// -> ProducingOutput(skip=Some) + /// Partial skip aggregation was triggered. First move to the + /// `ProducingOutput` state to drain the accumulated state, then move to + /// the `SkippingAggregation` state to convert input directly to partial + /// state without aggregation. + /// + /// ProducingOutput(skip=None) + /// -> ProducingOutput(skip=None) + /// One accumulated output batch was yielded, repeat to continue producing + /// output incrementally. + /// -> Done + /// All accumulated output was emitted. + /// + /// ProducingOutput(skip=Some) + /// -> ProducingOutput(skip=Some) + /// One accumulated output batch was yielded, repeat to continue producing + /// output incrementally. + /// -> SkippingAggregation + /// All accumulated output was emitted. Continue by converting raw + /// input batches directly to partial aggregate state. + /// + /// SkippingAggregation + /// -> SkippingAggregation + /// One `convert_to_state` batch was yielded; repeat to continue + /// processing. + /// -> Done + /// Input was exhausted. + /// + /// Done + /// -> (end) + /// ``` + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + let mut hash_table = self + .hash_table + .take() + .expect("hash_table should not be None"); + + let partial_table = self.consume_input(&mut hash_table).await?; + + self.start_output(&mut hash_table, partial_table.is_none())?; + + self.handle_producing_output(hash_table, &mut emitter) + .await?; + + match partial_table { + // partial skip has triggered and the + // stream will move to `SkippingAggregation` after these accumulated + // groups are emitted. + Some(partial_table) => { + self.handle_skipping_aggregation(partial_table, emitter) + .await?; + } + // Partial skip was never triggered and this state will finish in `Done` + None => { + let _ = self.reservation.try_resize(0); + } + } + + Ok(()) }) } @@ -310,376 +352,116 @@ impl PartialHashAggregateStream { hash_table.start_output() } - /// Handle ReadingInput state - aggregate input batches into the hash table. + /// aggregate input batches into the hash table. /// - /// See comments at `poll_next()` for details. + /// See comments at [`Self::create_stream`] for details. /// - /// Returns the next operator state with control flow decision. - fn handle_reading_input( + /// Returns Some(skip hash table) if should skip aggregation + async fn consume_input( &mut self, - cx: &mut Context<'_>, - mut original_state: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - PartialHashAggregateState::ReadingInput { .. } - )); - debug_assert!(original_state.hash_table().is_building()); - - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), - Poll::Ready(Some(Ok(batch))) => { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let input_rows = batch.num_rows(); - self.reduction_factor.add_total(input_rows); - let result = original_state.hash_table_mut().aggregate_batch(&batch); - timer.done(); + hash_table: &mut AggregateHashTable, + ) -> 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, - )); - } + while let Some(batch) = self.input.next().await.transpose()? { + let _timer = elapsed_compute.timer(); - if self.hit_soft_group_limit(original_state.hash_table()) { - let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut(), true); - timer.done(); - - if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - let PartialHashAggregateState::ReadingInput { hash_table } = - original_state - else { - unreachable!("expected reading input state") - }; - return ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: None, - }, - ); - } + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + hash_table.aggregate_batch(&batch)?; - self.update_skip_aggregation_probe( - input_rows, - original_state.hash_table().building_group_count(), - ); - - // True branch: a decision has been made to skip partial aggregation. - if self.should_skip_aggregation() { - let timer = elapsed_compute.timer(); - let result = match original_state.hash_table().partial_skip_table() { - Ok(skip_hash_table) => self - .start_output(original_state.hash_table_mut(), false) - .map(|()| skip_hash_table), - Err(e) => Err(e), - }; - timer.done(); - - match result { - Ok(skip_hash_table) => { - let PartialHashAggregateState::ReadingInput { hash_table } = - original_state - else { - unreachable!("expected reading input state") - }; - - // Move to `ProducingOutput` first. Its `skip_hash_table` - // field moves the stream to skip-partial aggregation after - // the accumulated batches have been output. - return ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: Some(skip_hash_table), - }, - ); - } - Err(e) => { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - } - } + if self.hit_soft_group_limit(hash_table) { + return Ok(None); + } - // TODO: impl memory-limited aggr, when OOM directly send - // partial state to final aggregate stage - if let Err(e) = self - .reservation - .try_resize(original_state.hash_table().memory_size()) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } + self.update_skip_aggregation_probe( + input_rows, + hash_table.building_group_count(), + ); - ControlFlow::Continue(original_state) - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_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(), true); - timer.done(); + // True branch: a decision has been made to skip partial aggregation. + if self.should_skip_aggregation() { + let skip_hash_table = hash_table.partial_skip_table()?; - match result { - Ok(()) => { - let PartialHashAggregateState::ReadingInput { hash_table } = - original_state - else { - unreachable!("expected reading input state") - }; - ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: None, - }, - ) - } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } - } + // Move to `ProducingOutput` first. Its `skip_hash_table` + // field moves the stream to skip-partial aggregation after + // the accumulated batches have been output. + return Ok(Some(skip_hash_table)); } - } - } - /// Handle ProducingOutput state - emit partial aggregate state 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: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - PartialHashAggregateState::ProducingOutput { .. } - )); - debug_assert!(!original_state.hash_table().is_building()); + // TODO: impl memory-limited aggr, when OOM directly send + // partial state to final aggregate stage + self.reservation.try_resize(hash_table.memory_size())?; + } 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)) => { - let _ = self - .reservation - .try_resize(original_state.hash_table().memory_size()); - self.reduction_factor.add_part(batch.num_rows()); - debug_assert!(batch.num_rows() > 0); - let next_state = if original_state.hash_table().is_done() { - match original_state { - PartialHashAggregateState::ProducingOutput { - skip_hash_table: Some(hash_table), - .. - } => { - PartialHashAggregateState::SkippingAggregation { hash_table } - } - PartialHashAggregateState::ProducingOutput { - skip_hash_table: None, - .. - } => PartialHashAggregateState::Done, - _ => unreachable!("expected producing output state"), - } - } 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); - // If the previous `Aggregating` stage decided to skip partial - // aggregation, go to the `SkippingAggregation` stage; otherwise finish. - let next_state = match original_state { - PartialHashAggregateState::ProducingOutput { - skip_hash_table: Some(hash_table), - .. - } => PartialHashAggregateState::SkippingAggregation { hash_table }, - PartialHashAggregateState::ProducingOutput { - skip_hash_table: None, - .. - } => PartialHashAggregateState::Done, - _ => unreachable!("expected producing output state"), - }; - ControlFlow::Continue(next_state) - } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), - } + Ok(None) } - /// Handle SkippingAggregation state - convert raw input directly to partial states. + /// Handle ProducingOutput state - emit partial aggregate state batches. /// /// See comments at `poll_next()` for details. /// /// Returns the next operator state with control flow decision. - fn handle_skipping_aggregation( + async fn handle_producing_output( &mut self, - cx: &mut Context<'_>, - mut original_state: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - PartialHashAggregateState::SkippingAggregation { .. } - )); + mut hash_table: AggregateHashTable, + emitter: &mut TryEmitter, + ) -> Result<()> { + debug_assert!(!hash_table.is_building()); - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), - Poll::Ready(Some(Ok(batch))) => { - if let Some(probe) = self.skip_aggregation_probe.as_mut() { - probe.record_skipped(&batch); - } + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + while let Some(batch) = hash_table.next_output_batch()? { + let _ = self.reservation.try_resize(hash_table.memory_size()); + self.reduction_factor.add_part(batch.num_rows()); + debug_assert!(batch.num_rows() > 0); - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = match &mut original_state { - PartialHashAggregateState::SkippingAggregation { hash_table } => { - hash_table.convert_batch_to_state(&batch) - } - _ => unreachable!("expected skipping aggregation state"), - }; + if hash_table.is_done() { timer.done(); + emitter.emit(batch).await; - match result { - Ok(batch) => ControlFlow::Break(( - Poll::Ready(Some( - Ok(batch.record_output(&self.baseline_metrics)), - )), - original_state, - )), - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } - } - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } - Poll::Ready(None) => { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - ControlFlow::Continue(PartialHashAggregateState::Done) - } + return Ok(()); + }; + + timer.done(); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); } - } -} -impl Stream for PartialHashAggregateStream { - type Item = Result; + let _ = self.reservation.try_resize(0); - /// Entry point for the partial hash aggregate state machine. - /// - /// See comments in [`PartialHashAggregateStream`] for high-level ideas. - /// - /// State transition graph: - /// - /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling input and aggregating batches into the - /// in-memory hash table. - /// - /// ReadingInput - /// -> ReadingInput - /// Aggregate one batch, update the inner aggregate hash table, and - /// continue with the next input batch. - /// -> ProducingOutput(skip=None) - /// Input was exhausted, or the soft group limit was reached. Move to - /// the next state to start outputting. - /// -> ProducingOutput(skip=Some) - /// Partial skip aggregation was triggered. First move to the - /// `ProducingOutput` state to drain the accumulated state, then move to - /// the `SkippingAggregation` state to convert input directly to partial - /// state without aggregation. - /// - /// ProducingOutput(skip=None) - /// -> ProducingOutput(skip=None) - /// One accumulated output batch was yielded, repeat to continue producing - /// output incrementally. - /// -> Done - /// All accumulated output was emitted. - /// - /// ProducingOutput(skip=Some) - /// -> ProducingOutput(skip=Some) - /// One accumulated output batch was yielded, repeat to continue producing - /// output incrementally. - /// -> SkippingAggregation - /// All accumulated output was emitted. Continue by converting raw - /// input batches directly to partial aggregate state. - /// - /// SkippingAggregation - /// -> SkippingAggregation - /// One `convert_to_state` batch was yielded; repeat to continue - /// processing. - /// -> Done - /// Input was exhausted. + Ok(()) + } + + /// Handle SkippingAggregation state - convert raw input directly to partial states. /// - /// Done - /// -> (end) - /// ``` - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - loop { - let cur_state = self - .state - .take() - .expect("PartialHashAggregateStream state should not be None"); + /// See comments at [`Self::create_stream`] for details. + async fn handle_skipping_aggregation( + mut self, + mut hash_table: AggregateHashTable, + mut emitter: TryEmitter, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let next_state = match cur_state { - state @ PartialHashAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ PartialHashAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(state) - } - state @ PartialHashAggregateState::SkippingAggregation { .. } => { - self.handle_skipping_aggregation(cx, state) - } - state @ PartialHashAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Poll::Ready(None); - } - }; + while let Some(batch) = self.input.next().await.transpose()? { + let timer = elapsed_compute.timer(); - 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; - } + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.record_skipped(&batch); } + let result = hash_table.convert_batch_to_state(&batch)?; + + timer.done(); + emitter.emit(result).await; } - } -} -impl RecordBatchStream for PartialHashAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) + Ok(()) } } @@ -970,7 +752,8 @@ mod tests { // Execute and collect results let mut stream = - PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)? + .into_stream(); let mut results = Vec::new(); while let Some(result) = stream.next().await { @@ -1114,7 +897,8 @@ mod tests { // Execute and collect results let mut stream = - PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)? + .into_stream(); let mut results = Vec::new(); while let Some(result) = stream.next().await { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 69c7ee05fa1ac..9dbca37a76a1e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -693,7 +693,7 @@ impl From for SendableRecordBatchStream { fn from(stream: StreamType) -> Self { match stream { StreamType::AggregateStream(stream) => Box::pin(stream), - StreamType::PartialHash(stream) => Box::pin(stream), + StreamType::PartialHash(stream) => stream.into_stream(), StreamType::PartialReduceHash(stream) => Box::pin(stream), StreamType::FinalHash(stream) => stream.into_stream(), StreamType::SingleHash(stream) => Box::pin(stream), From 158f330a97eb56268b831153ffb3439ab600a1d2 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:50:06 +0300 Subject: [PATCH 3/4] update comments --- .../physical-plan/src/aggregates/hash_stream.rs | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 7b10a09752367..0b96b64c9e290 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -295,7 +295,7 @@ impl PartialHashAggregateStream { self.start_output(&mut hash_table, partial_table.is_none())?; - self.handle_producing_output(hash_table, &mut emitter) + self.produce_output(hash_table, &mut emitter) .await?; match partial_table { @@ -404,10 +404,8 @@ impl PartialHashAggregateStream { /// Handle ProducingOutput state - emit partial aggregate state batches. /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - async fn handle_producing_output( + /// See comments at [`Self::create_stream`] for details. + async fn produce_output( &mut self, mut hash_table: AggregateHashTable, emitter: &mut TryEmitter, @@ -556,7 +554,7 @@ impl FinalHashAggregateStream { .hash_table .take() .expect("hash_table should not be None"); - self.handle_reading_input(&mut hash_table).await?; + self.consume_input(&mut hash_table).await?; self.produce_output(hash_table, emitter).await?; @@ -581,10 +579,8 @@ impl FinalHashAggregateStream { /// Handle ReadingInput state - aggregate partial state batches into the hash table. /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - async fn handle_reading_input( + /// See comments at [`Self::create_stream`] for details. + async fn consume_input( &mut self, hash_table: &mut AggregateHashTable, ) -> Result<()> { From df7e482c865ce87f6b383bd282d55cb3d95e3ecf Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:04:46 +0300 Subject: [PATCH 4/4] format --- datafusion/physical-plan/src/aggregates/hash_stream.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 0b96b64c9e290..d79e05dc5e3f7 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -295,8 +295,7 @@ impl PartialHashAggregateStream { self.start_output(&mut hash_table, partial_table.is_none())?; - self.produce_output(hash_table, &mut emitter) - .await?; + self.produce_output(hash_table, &mut emitter).await?; match partial_table { // partial skip has triggered and the