alamb commented on code in PR #23408: URL: https://github.com/apache/datafusion/pull/23408#discussion_r3561093507
########## datafusion/physical-plan/src/aggregates/single_stream.rs: ########## @@ -0,0 +1,372 @@ +// 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::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! 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_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, SingleMarker}; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can run the full logical aggregation in one operator. This +/// stream implements the single stage for grouped hash aggregation. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=single) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final aggregate values for all groups (for example, `AVG(x)`) +/// +/// This stream implements the complete aggregation without a partial/final +/// split. It consumes raw input rows and emits final aggregate values. +pub(crate) struct SingleHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + 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<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 + } +} + +impl SingleHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + 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::<SingleMarker>::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + + let reservation = + MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(SingleHashAggregateState::ReadingInput { hash_table }), + }) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable<SingleMarker>, + ) -> Result<()> { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); Review Comment: maybe worth a comment that this eagerly releases the input stream ########## datafusion/physical-plan/src/aggregates/single_stream.rs: ########## @@ -0,0 +1,372 @@ +// 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::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! 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_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, SingleMarker}; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can run the full logical aggregation in one operator. This +/// stream implements the single stage for grouped hash aggregation. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=single) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final aggregate values for all groups (for example, `AVG(x)`) +/// +/// This stream implements the complete aggregation without a partial/final +/// split. It consumes raw input rows and emits final aggregate values. +pub(crate) struct SingleHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + 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<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 + } +} + +impl SingleHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + 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::<SingleMarker>::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + + let reservation = + MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(SingleHashAggregateState::ReadingInput { hash_table }), + }) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable<SingleMarker>, + ) -> 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 `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 = 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)) + } + } + } + } + } + + /// 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 Review Comment: I think we should return the error if the reservation fails (rather than ignoring it), right (as is done above in ```rust if let Err(e) = self .reservation .try_resize(original_state.hash_table().memory_size()) { return ControlFlow::Break(( Poll::Ready(Some(Err(e))), original_state, )); } ``` ########## datafusion/physical-plan/src/aggregates/single_stream.rs: ########## @@ -0,0 +1,372 @@ +// 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::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! 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_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, SingleMarker}; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can run the full logical aggregation in one operator. This +/// stream implements the single stage for grouped hash aggregation. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=single) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final aggregate values for all groups (for example, `AVG(x)`) +/// +/// This stream implements the complete aggregation without a partial/final +/// split. It consumes raw input rows and emits final aggregate values. +pub(crate) struct SingleHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + 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<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 + } +} + +impl SingleHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + 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); Review Comment: This variable is ignored -- why bother to make it? Or is this a TODO when we implement spilling? If so, perhaps we can leave the TODO / link to more context -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
