diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fb..82b8681efbf3 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 886ffdd3a0b9..f109fcef9514 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_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. @@ -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, -} - -/// 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 - } + /// 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 { @@ -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> { + 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(()) }) } @@ -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, + ) -> 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, + mut emitter: TryEmitter, + ) -> 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; + 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> { - 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(()) } }