diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 11446137f3ca1..6258576089b51 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -28,6 +28,7 @@ use crate::aggregates::{ ordered_partial_stream::OrderedPartialAggregateStream, partial_reduce_stream::PartialReduceHashAggregateStream, row_hash::GroupedHashAggregateStream, + single_stream::SingleHashAggregateStream, topk_stream::GroupedTopKAggregateStream, }; use crate::execution_plan::{CardinalityEffect, EmissionType}; @@ -84,6 +85,7 @@ mod ordered_final_stream; mod ordered_partial_stream; mod partial_reduce_stream; mod row_hash; +mod single_stream; mod skip_partial; mod topk; mod topk_stream; @@ -539,6 +541,9 @@ enum StreamType { /// Final stage of the hash aggregation /// Input output scheme: partial state -> final result FinalHash(FinalHashAggregateStream), + /// Single stage of the hash aggregation + /// Input output scheme: initial input -> final result + SingleHash(SingleHashAggregateStream), /// Partial stage of aggregation for ordered input. OrderedPartialAggregate(OrderedPartialAggregateStream), /// Final stage of aggregation for ordered input. @@ -547,8 +552,8 @@ enum StreamType { /// /// Note this is being incrementally migrated to dedicated streams like /// [`StreamType::PartialHash`], [`StreamType::FinalHash`], - /// [`StreamType::OrderedPartialAggregate`], and - /// [`StreamType::OrderedFinalAggregate`] + /// [`StreamType::SingleHash`], [`StreamType::OrderedPartialAggregate`], + /// and [`StreamType::OrderedFinalAggregate`] /// /// See issue for details: GroupedHash(GroupedHashAggregateStream), @@ -567,6 +572,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::OrderedPartialAggregate(stream) => Box::pin(stream), StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), StreamType::GroupedHash(stream) => Box::pin(stream), @@ -1071,6 +1077,12 @@ impl AggregateExec { self, context, partition, )?)); } + + if self.should_use_single_hash_stream(context) { + return Ok(StreamType::SingleHash(SingleHashAggregateStream::new( + self, context, partition, + )?)); + } } // Execution paths that have not been migrated use the fallback implementation @@ -1133,6 +1145,21 @@ impl AggregateExec { && self.group_by.is_single() } + fn should_use_single_hash_stream(&self, context: &TaskContext) -> bool { + // TODO: implement memory-limited path and remove this limitation + if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { + return false; + } + + matches!( + self.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + ) && self.limit_options_supported_by_hash_stream() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + fn should_use_ordered_final_aggregate_stream(&self, context: &TaskContext) -> bool { // TODO: implement memory-limited path and remove this limitation if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { @@ -3506,6 +3533,76 @@ mod tests { Ok(()) } + fn single_test_aggregate(mode: AggregateMode) -> Result { + let (schema, batches) = some_data(); + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(b)") + .build()?, + )]; + + AggregateExec::try_new( + mode, + group_by, + aggregates, + vec![None], + input, + Arc::clone(&schema), + ) + } + + /// For single-stage aggregation, ensures `SingleHashAggregateStream` is + /// used for both single modes when enabled by migration config. + #[tokio::test] + async fn single_aggregate_planning() -> Result<()> { + let task_ctx = new_migrated_hash_ctx(2); + + for mode in [AggregateMode::Single, AggregateMode::SinglePartitioned] { + let aggregate = single_test_aggregate(mode)?; + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_eq!( + output.iter().map(RecordBatch::num_rows).collect::>(), + vec![2, 1] + ); + assert_eq!( + batches_to_sort_string(&output), + "\ ++---+----------+ +| a | COUNT(b) | ++---+----------+ +| 2 | 2 | +| 3 | 3 | +| 4 | 3 | ++---+----------+" + ); + } + + Ok(()) + } + + /// Spilling behavior is not implemented for single-stage hash stream yet, + /// so fall back to the existing `GroupedHashAggregateStream`. + #[tokio::test] + async fn single_aggregate_with_memory_limit_planning() -> Result<()> { + let task_ctx = new_finite_memory_migrated_hash_ctx(2, 1)?; + + for mode in [AggregateMode::Single, AggregateMode::SinglePartitioned] { + let aggregate = single_test_aggregate(mode)?; + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + } + + Ok(()) + } + fn partial_reduce_test_aggregate() -> Result { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs new file mode 100644 index 0000000000000..17498a60db590 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -0,0 +1,358 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Single-stage hash aggregation stream implementation. +//! +//! This stream is part of the incremental migration from +//! [`crate::aggregates::row_hash::GroupedHashAggregateStream`]. +//! +//! 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, assert_or_internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, AggregateTableMode}; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can execute in a single stage when raw input rows are +/// aggregated directly into final aggregate values. This stream implements that +/// single stage for both `Single` and `SinglePartitioned` aggregate modes. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=single) +/// -- Scan(t) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final aggregate values for all groups (for example, `AVG(x)` is +/// calculated from the accumulated `SUM(x)` and `COUNT(x)` state) +pub(crate) struct SingleHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing raw rows. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// 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-stage 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 + } +} + +impl SingleHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + super::AggregateMode::Single | super::AggregateMode::SinglePartitioned + )); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + + let hash_table = + AggregateHashTable::new(agg, partition, Arc::clone(&schema), batch_size)?; + assert_or_internal_err!( + hash_table.mode() == AggregateTableMode::Single, + "SingleHashAggregateStream expected single aggregate hash table" + ); + + let reservation = + MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(SingleHashAggregateState::ReadingInput { hash_table }), + }) + } + + /// Handle ReadingInput state - aggregate raw 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( + &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()) + { + 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 = original_state.hash_table_mut().start_output(); + timer.done(); + + match result { + Ok(()) => { + ControlFlow::Continue(original_state.into_producing_output()) + } + Err(e) => { + ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + } + } + } + } + } + + /// 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)) => { + 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 SingleHashAggregateStream { + type Item = Result; + + /// Entry point for the single-stage 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 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; + } + } + } + } +} + +impl RecordBatchStream for SingleHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +}