diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 2293e7b1b8e89..dac0d4b7c527a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -38,9 +38,12 @@ use crate::aggregates::{ evaluate_group_by, }; -use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; +use super::common::{ + AggregateAccumulator, AggregateBatchFn, EvaluatedAggregateBatch, + MaterializeAccumulatorFn, +}; -/// Aggregate table shared by the ordered partial and final paths. +/// Aggregate table shared by the ordered single, partial and final paths. /// /// # Ordering optimization /// @@ -48,10 +51,11 @@ use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; /// are proven complete. Completed groups can be emitted before the input stream /// ends, which keeps memory bounded by the active ordered key range. /// -/// # Partial and final variant difference +/// # Single, partial and final variant difference /// /// The partial and final aggregate tables implement the two stages of grouped -/// aggregation. See +/// aggregation, while the single aggregate table implements both stages in one +/// table. See /// [`OrderedPartialAggregateStream`](crate::aggregates::ordered_partial_stream::OrderedPartialAggregateStream) /// for the high-level plan shape. /// @@ -67,6 +71,11 @@ use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; /// - Table stores: `k, sum(v), count(v)` /// - Output schema: `k, avg(v)` /// +/// Single table ([`AggregateMode::Single`], with optional filter from query): +/// - Input rows: `k, v` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, avg(v)` +/// /// # Marker Type /// /// `OrderedAggrMode` selects the aggregate semantics. For example, @@ -75,7 +84,7 @@ use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; /// `OrderedAggregateTable::::new_with_input_order(...)` /// consumes partial states and emits final values. /// -/// Shared methods live on `impl`; partial/final behavior lives on +/// Shared methods live on `impl`; single/partial/final behavior lives on /// marker-specific impls. pub(in crate::aggregates) struct OrderedAggregateTable { /// Output schema: group columns followed by aggregate state or final values. @@ -129,7 +138,7 @@ pub(super) struct OrderedAggregateTableBuffer { impl OrderedAggregateTable { #[expect( clippy::too_many_arguments, - reason = "keeps ordered partial and final table construction explicit" + reason = "keeps ordered single, partial and final table construction explicit" )] pub(super) fn new_for_mode( agg: &AggregateExec, @@ -259,8 +268,8 @@ impl OrderedAggregateTable { /// /// Unlike normal ordered emission, this operation is allowed to take the /// active (incomplete) groups. Partial aggregation can pass those states to - /// its final stage, while final aggregation sorts and spills them before - /// replay. + /// its final stage, while single and final aggregation sort and spill them + /// before replay. pub(in crate::aggregates) fn take_state_batch( &mut self, ) -> Result> { @@ -304,18 +313,18 @@ impl OrderedAggregateTable { EmitTo::All => (EmitTo::First(self.batch_size), false), } } - /// Aggregates one evaluated input batch. - /// - /// This common utility is used by ordered partial and ordered final aggregation. - /// - /// # Argument: `is_final` + + /// Aggregates one evaluated input batch after selecting the mode-specific + /// accumulator operation. /// - /// - `true`: merge partial aggregate states for final aggregation. - /// - `false`: update aggregate states from raw input for partial aggregation. + /// Each aggregation mode chooses a different `aggregate_fn` according to its + /// semantics. For example, partial aggregation takes raw inputs and updates + /// stored partial states, so it uses + /// [`datafusion_expr::GroupsAccumulator::update_batch`]. pub(super) fn aggregate_evaluated_batch( &mut self, evaluated_batch: &EvaluatedAggregateBatch, - is_final: bool, + aggregate_fn: AggregateBatchFn, ) -> Result<()> { for group_values in &evaluated_batch.grouping_set_args { let starting_num_groups = self.buffer.group_values.len(); @@ -338,19 +347,7 @@ impl OrderedAggregateTable { .iter_mut() .zip(evaluated_batch.accumulator_args.iter()) { - if is_final { - acc.merge_batch( - values, - &self.buffer.group_indices, - total_num_groups, - )?; - } else { - acc.update_batch( - values, - &self.buffer.group_indices, - total_num_groups, - )?; - } + aggregate_fn(acc, values, &self.buffer.group_indices, total_num_groups)?; } drop(timer); } @@ -361,15 +358,13 @@ impl OrderedAggregateTable { /// Emits groups allowed by `GroupOrdering`, leaving only the current /// unfinished ordered-key range buffered. /// - /// This common utility is used by ordered partial and ordered final aggregation. - /// - /// # Argument: `is_final` - /// - /// - `true`: output final aggregate values. - /// - `false`: output partial accumulator states. - pub(super) fn next_output_batch_for_mode( + /// Each aggregation mode chooses a different `materialize_accumulator_fn` + /// according to its semantics. For example, partial aggregation emits + /// partial states to feed the final stage, so it uses + /// [`datafusion_expr::GroupsAccumulator::state`]. + pub(super) fn next_output_batch_inner( &mut self, - is_final: bool, + materialize_accumulator_fn: MaterializeAccumulatorFn, ) -> Result> { if self.buffer.group_values.is_empty() { return Ok(None); @@ -394,11 +389,7 @@ impl OrderedAggregateTable { } for acc in &mut self.buffer.accumulators { - if is_final { - output.push(acc.evaluate(emit_to)?); - } else { - output.extend(acc.state(emit_to)?); - } + output.extend(materialize_accumulator_fn(acc, emit_to)?); } drop(timer); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs index 2c7ec01654a63..fbf3ccc738cb7 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -20,6 +20,7 @@ mod common_ordered; mod final_table; mod ordered_final_table; mod ordered_partial_table; +mod ordered_single_table; mod partial_reduce_table; mod partial_table; mod single_table; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index fd064ebffec12..1b7d419dd5571 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -30,6 +30,7 @@ use crate::aggregates::aggregate_hash_table::FinalMarker; use crate::aggregates::group_values::GroupByMetrics; use crate::aggregates::{AggregateExec, AggregateMode}; +use super::common::HashAggregateAccumulator; use super::common_ordered::OrderedAggregateTable; /// Implementation specific to final aggregation, where the table stores partial @@ -73,13 +74,16 @@ impl OrderedAggregateTable { // `PhysicalGroupBy::as_final()` removes grouping sets while planning // final aggregation, so final ordered aggregation sees one grouping. debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1); - self.aggregate_evaluated_batch(&evaluated_batch, true) + self.aggregate_evaluated_batch( + &evaluated_batch, + HashAggregateAccumulator::merge_batch, + ) } /// See comments in `ordered_partial_stream::next_output_batch` pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_for_mode(true) + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index a04e4dda8fb39..8545289f990cd 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -40,6 +40,7 @@ use crate::aggregates::{ group_values::GroupByMetrics, }; +use super::common::HashAggregateAccumulator; use super::common_ordered::OrderedAggregateTable; /// Implementation specific to partial aggregation, where the table stores @@ -81,7 +82,10 @@ impl OrderedAggregateTable { batch: &RecordBatch, ) -> Result<()> { let evaluated_batch = self.evaluate_batch(batch)?; - self.aggregate_evaluated_batch(&evaluated_batch, false) + self.aggregate_evaluated_batch( + &evaluated_batch, + HashAggregateAccumulator::update_batch, + ) } /// Emits the next batch of partial state rows for groups proven complete by @@ -101,6 +105,6 @@ impl OrderedAggregateTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_for_mode(false) + self.next_output_batch_inner(HashAggregateAccumulator::state) } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs new file mode 100644 index 0000000000000..8ba50e2a59507 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs @@ -0,0 +1,90 @@ +// 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. + +//! Aggregate table for single aggregation when raw input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::aggregate_hash_table::SingleMarker; +use crate::aggregates::group_values::GroupByMetrics; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common::HashAggregateAccumulator; +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to single aggregation, where the table stores final +/// aggregate values and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, avg(x)` +/// - Input rows: `k, x` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + )); + + let input_schema = agg.input().schema(); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + Self::new_for_mode( + agg, + &input_schema, + output_schema, + state_schema, + batch_size, + &agg.input_order_mode, + &agg.mode, + agg.filter_expr.iter().cloned().collect(), + group_by_metrics, + ) + } + + /// Aggregates one raw input batch and updates ordering information for any + /// newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + self.aggregate_evaluated_batch( + &evaluated_batch, + HashAggregateAccumulator::update_batch, + ) + } + + /// Emits the next batch of final aggregate values for groups proven complete + /// by the input ordering. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) + } +} diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6af67d048256e..74a8f56a44207 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -59,8 +59,8 @@ //! AggregateExec (partial, ordered) //! ``` //! -//! See [`OrderedPartialAggregateStream`] and [`OrderedFinalAggregateStream`] for -//! details. +//! See [`OrderedPartialAggregateStream`], [`OrderedFinalAggregateStream`], and +//! [`OrderedSingleAggregateStream`] for details. //! //! Related configuration: //! @@ -81,7 +81,8 @@ //! input //! ``` //! -//! See [`SingleHashAggregateStream`] for details. +//! See [`SingleHashAggregateStream`] and [`OrderedSingleAggregateStream`] for +//! details. //! //! Related configuration: //! @@ -153,6 +154,7 @@ use crate::aggregates::{ hash_stream::{FinalHashAggregateStream, PartialHashAggregateStream}, ordered_final_stream::OrderedFinalAggregateStream, ordered_partial_stream::OrderedPartialAggregateStream, + ordered_single_stream::OrderedSingleAggregateStream, partial_reduce_stream::PartialReduceHashAggregateStream, single_stream::SingleHashAggregateStream, }; @@ -209,6 +211,7 @@ mod hash_stream; pub mod order; mod ordered_final_stream; mod ordered_partial_stream; +mod ordered_single_stream; mod partial_reduce_stream; mod single_stream; mod skip_partial; @@ -672,12 +675,15 @@ enum StreamType { OrderedPartialAggregate(OrderedPartialAggregateStream), /// Final stage of aggregation for ordered input. OrderedFinalAggregate(OrderedFinalAggregateStream), + /// Single stage of aggregation for ordered input. + OrderedSingleAggregate(OrderedSingleAggregateStream), /// Hash aggregation reused for multiple stages /// /// Note this is being incrementally migrated to dedicated streams like /// [`StreamType::PartialHash`], [`StreamType::FinalHash`], /// [`StreamType::OrderedPartialAggregate`], and - /// [`StreamType::OrderedFinalAggregate`] + /// [`StreamType::OrderedFinalAggregate`], and + /// [`StreamType::OrderedSingleAggregate`] /// /// See issue for details: GroupedHash(GroupedHashAggregateStream), @@ -699,6 +705,7 @@ impl From for SendableRecordBatchStream { StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), + StreamType::OrderedSingleAggregate(stream) => Box::pin(stream), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } @@ -1212,6 +1219,12 @@ impl AggregateExec { )?)); } + if self.should_use_ordered_single_aggregate_stream(context) { + return Ok(StreamType::OrderedSingleAggregate( + OrderedSingleAggregateStream::new(self, context, partition)?, + )); + } + if self.should_use_single_hash_stream(context) { return Ok(StreamType::SingleHash(SingleHashAggregateStream::new( self, context, partition, @@ -1277,6 +1290,16 @@ impl AggregateExec { && self.group_by.is_single() } + fn should_use_ordered_single_aggregate_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + ) && self.limit_options.is_none() + && 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 { matches!( self.mode, @@ -4082,6 +4105,83 @@ mod tests { Ok(()) } + /// Ensures `OrderedSingleAggregateStream` is used for ordered raw input. + #[tokio::test] + async fn ordered_single_aggregate_planning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1])), + Arc::new(Int32Array::from(vec![10, 11, 10])), + Arc::new(Int64Array::from(vec![1, 2, 3])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2])), + Arc::new(Int32Array::from(vec![20, 21])), + Arc::new(Int64Array::from(vec![4, 5])), + ], + )?, + ]; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value_col)") + .build()?, + )], + vec![None], + input, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++----------+-----------+----------------+ +| sort_col | group_col | SUM(value_col) | ++----------+-----------+----------------+ +| 1 | 10 | 4 | +| 1 | 11 | 2 | +| 2 | 20 | 4 | +| 2 | 21 | 5 | ++----------+-----------+----------------+ +"); + + let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + let stream = aggregate.execute_typed(0, &finite_memory_task_ctx)?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + + 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/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs new file mode 100644 index 0000000000000..2025bdce307d9 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -0,0 +1,886 @@ +// 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 aggregate stream for ordered raw input. + +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::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::aggregate_hash_table::{OrderedAggregateTable, SingleMarker}; +use super::group_values::GroupByMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::{AggregateExec, create_schema}; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Single aggregate stream for `InputOrderMode::Sorted` and +/// `InputOrderMode::PartiallySorted`. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// If the input is ordered by `k`, and there are existing key partitioning on group +/// by keys, the single mode aggregation with ordering optimization can be used: +/// +/// ## Plan +/// AggregateExec(stage=single, ordered) +/// -- DataSourceExec(t) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final results for all groups (for example, `AVG(x)`) +/// +/// # Order-based Optimization +/// +/// For the aggregation work, the hash aggregation implementation is reused. +/// +/// After each input batch, check whether any groups can be emitted eagerly to +/// improve memory efficiency. For example, if the last group key seen is +/// `k = 100`, it is safe to emit all groups with keys less than 100 because the +/// input is ordered. +/// +/// # Memory Pressure and Spilling +/// +/// ## Fully ordered case +/// +/// If the input is ordered by every group key, for example: +/// +/// - Input order: `a, b` +/// - `GROUP BY`: `a, b` +/// +/// Completed groups can be emitted as soon as the next group is observed. Thus, +/// only the current group remains active after completed groups are emitted, and +/// memory usage does not grow with the total number of groups. +/// +/// If a memory reservation nevertheless fails, the stream returns the error +/// directly, indicating an unexpected behavior. +/// +/// ## Partially ordered case +/// +/// If the input is ordered by only a subset of the group keys, for example: +/// +/// - Input order: `a` +/// - `GROUP BY`: `a, b` +/// +/// If one `a` value contains many distinct `b` values, the table may accumulate +/// enough groups to exceed the memory limit. +/// +/// On reservation failure, the stream sorts the current intermediate states by +/// the complete group key and spills them as one run. After the input ends, it +/// spills any remaining states, performs a sort-preserving merge of all runs, +/// and feeds the merged input into a fully ordered final aggregate stream. +pub(crate) struct OrderedSingleAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + state: Option, +} + +/// Spill configuration and accumulated runs for partially ordered single +/// aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct OrderedSingleSpillContext { + /// Aggregate configuration used to construct the final replay stream. + final_agg: AggregateExec, + /// Task context + context: Arc, + /// Original partition index + partition: usize, + /// Target batch size from configuration + batch_size: usize, + /// Full group-key ordering, such ordering with be kept in: a) individual spill + /// files, b) order after final merging and streaming aggregate + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Fully sorted spill runs waiting to be merged. + spills: Vec, +} + +/// See comments at `poll_next()` for details. +enum OrderedSingleAggregateState { + ReadingInput { + table: OrderedAggregateTable, + /// None if either + /// - Disk Manager doesn't enable temporary file creation + /// - The group keys are fully ordered, it's expected to use bounded memory + spill_context: Option>, + }, + Spilling { + table: OrderedAggregateTable, + spill_context: Box, + }, + ProducingOutput { + table: OrderedAggregateTable, + }, + PreparingMergeInput { + table: OrderedAggregateTable, + spill_context: Box, + }, + MergingSpills { + stream: SendableRecordBatchStream, + }, + Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, +} + +type OrderedSingleAggregatePoll = Poll>>; +type OrderedSingleAggregateStateTransition = ControlFlow< + (OrderedSingleAggregatePoll, OrderedSingleAggregateState), + OrderedSingleAggregateState, +>; + +impl OrderedSingleSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + batch_size: usize, + input_order_mode: &InputOrderMode, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result { + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + let output_ordering = agg.cache.output_ordering(); + let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { + return internal_err!( + "Ordered single spill requires partially ordered input" + ); + }; + let spill_indices = order_indices.iter().copied().chain( + (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), + ); + let spill_sort_exprs = spill_indices.map(|idx| { + let field = group_schema.field(idx); + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Ordered single spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + // Spilled rows contain group keys and intermediate states. Replay must + // merge those states and evaluate the final aggregate values. + let mut final_agg = agg.clone(); + final_agg.mode = match agg.mode { + AggregateMode::Single => AggregateMode::Final, + AggregateMode::SinglePartitioned => AggregateMode::FinalPartitioned, + mode => { + return internal_err!( + "Ordered single aggregate spill cannot replay aggregate mode {mode:?}" + ); + } + }; + final_agg.group_by = Arc::new(agg.group_by.as_final()); + final_agg.input_order_mode = InputOrderMode::Sorted; + + Ok(Self { + final_agg, + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`OrderedSingleAggregateStream`] for spilling details. + fn spill_table( + &mut self, + table: &mut OrderedAggregateTable, + ) -> Result<()> { + let Some(batch) = table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "OrderedSingleAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Ordered single aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run and finalizes it through the fully ordered path. + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + final_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &final_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl OrderedSingleAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + )); + debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let state_schema = Arc::new(create_schema( + input_schema.as_ref(), + &agg.group_by, + &agg.aggr_expr, + AggregateMode::Partial, + )?); + + let table = OrderedAggregateTable::::new( + agg, + partition, + Arc::clone(&schema), + Arc::clone(&state_schema), + batch_size, + )?; + + let can_spill = + matches!(agg.input_order_mode, InputOrderMode::PartiallySorted(_)) + && context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + Some(Box::new(OrderedSingleSpillContext::new( + agg, + context, + partition, + batch_size, + &agg.input_order_mode, + &state_schema, + spill_metrics, + )?)) + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("OrderedSingleAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + state: Some(OrderedSingleAggregateState::ReadingInput { + table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> OrderedSingleAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + OrderedSingleAggregateState::Error, + )) + } + + fn break_with_internal_err(message: &str) -> OrderedSingleAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + table: &OrderedAggregateTable, + spill_context: Option<&OrderedSingleSpillContext>, + ) -> usize { + let table_size = table.memory_size(); + if spill_context.is_some() { + // See `OrderedSingleAggregateStream` comments for how is it estimated + table_size.saturating_add(table.num_groups().saturating_mul(size_of::())) + } else { + table_size + } + } + + /// Consumes one ordered raw input batch, then immediately emits + /// finalized groups if the ordering proves any group is ready. + /// + /// 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<'_>, + original_state: OrderedSingleAggregateState, + ) -> OrderedSingleAggregateStateTransition { + let OrderedSingleAggregateState::ReadingInput { + mut table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered single aggregate stream expected ReadingInput state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedSingleAggregateState::ReadingInput { + table, + spill_context, + }, + )), + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )); + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(spill_context) = spill_context else { + // `None` means spilling is not supported, see comments + // at `OrderedSingleAggregateState` for details. + return Self::break_with_err(e); + }; + if table.is_empty() { + return Self::break_with_internal_err( + "Ordered single aggregate ran out of memory with no aggregated groups", + ); + } + return ControlFlow::Continue( + OrderedSingleAggregateState::Spilling { + table, + spill_context, + }, + ); + } + Err(e) => { + return Self::break_with_err(e); + } + } + + let result = if spill_context + .as_ref() + .is_some_and(|spill_context| spill_context.has_spills()) + { + // Once one incomplete run is spilled, every remaining state + // must participate in replay so no group is finalized twice. + Ok(None) + } else { + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + result + }; + + match result { + // Some finalized groups can be emitted. Yield them, then + // continue aggregating input in the current state. + Ok(Some(batch)) => { + if let Err(e) = + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )) + { + return Self::break_with_err(e); + } + let next_state = OrderedSingleAggregateState::ReadingInput { + table, + spill_context, + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok( + batch.record_output(&self.baseline_metrics) + ))), + next_state, + )) + } + // Can't do early emit, continue aggregating. + Ok(None) => { + ControlFlow::Continue(OrderedSingleAggregateState::ReadingInput { + table, + spill_context, + }) + } + Err(e) => Self::break_with_err(e), + } + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + OrderedSingleAggregateState::PreparingMergeInput { + table, + spill_context, + }, + ) + } + _ => { + table.input_done(); + ControlFlow::Continue( + OrderedSingleAggregateState::ProducingOutput { table }, + ) + } + } + } + } + } + + /// Sorts and spills one complete in-memory state run, then resumes input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_spilling( + &mut self, + original_state: OrderedSingleAggregateState, + ) -> OrderedSingleAggregateStateTransition { + let OrderedSingleAggregateState::Spilling { + mut table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered single aggregate stream expected Spilling state", + ); + }; + + // Sanity check: it's impossible to OOM when the table is empty + if table.is_empty() { + return Self::break_with_internal_err( + "Ordered single aggregation entered Spilling with an empty table", + ); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + result = Err(e); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input + Ok(()) => ControlFlow::Continue(OrderedSingleAggregateState::ReadingInput { + table, + spill_context: Some(spill_context), + }), + Err(e) => Self::break_with_err(e), + } + } + + /// 1. Spills the last in-memory run. + /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// merge to all spills. + /// 3. Constructs a replay stream: an ordered aggregate stream over the fully + /// ordered input constructed from the spills. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_preparing_merge_input( + &mut self, + original_state: OrderedSingleAggregateState, + ) -> OrderedSingleAggregateStateTransition { + let OrderedSingleAggregateState::PreparingMergeInput { + mut table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered single aggregate stream expected PreparingMergeInput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let replay = match spill_context.spill_table(&mut table) { + Ok(()) => { + let group_by_metrics = table.group_by_metrics(); + drop(table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ), + Err(e) => Err(e), + } + } + Err(e) => Err(e), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(OrderedSingleAggregateState::MergingSpills { + stream, + }) + } + Err(e) => Self::break_with_err(e), + } + } + + /// Forwards output from the fully ordered stream that consumes the merged + /// spill runs. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_merging_spills( + &mut self, + cx: &mut Context<'_>, + original_state: OrderedSingleAggregateState, + ) -> OrderedSingleAggregateStateTransition { + let OrderedSingleAggregateState::MergingSpills { mut stream } = original_state + else { + return Self::break_with_internal_err( + "Ordered single aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedSingleAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + OrderedSingleAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => ControlFlow::Continue(OrderedSingleAggregateState::Done), + } + } + + /// Emits one batch after input is exhausted. + /// + /// `table.input_done()` has already made every remaining group safe to emit, + /// so this state keeps draining until the table is empty. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: OrderedSingleAggregateState, + ) -> OrderedSingleAggregateStateTransition { + let OrderedSingleAggregateState::ProducingOutput { table } = original_state + else { + return Self::break_with_internal_err( + "Ordered single aggregate stream expected ProducingOutput state", + ); + }; + + let mut table = table; + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let next_state = if table.is_empty() { + drop(table); + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + OrderedSingleAggregateState::Done + } else { + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + return Self::break_with_err(e); + } + OrderedSingleAggregateState::ProducingOutput { table } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => Self::break_with_err(e), + Ok(None) => { + drop(table); + let next_state = OrderedSingleAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + ControlFlow::Continue(next_state) + } + } + } +} + +impl Stream for OrderedSingleAggregateStream { + type Item = Result; + + /// Entry point for the ordered single aggregate state machine. + /// + /// See comments in [`OrderedSingleAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered raw input and updating the + /// ordered single aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one input batch. If it fits in memory, optionally yield + /// groups proven complete by the input ordering, then read the next batch. + /// -> Spilling + /// The table cannot reserve enough memory. Move all current states into + /// one fully group-key-sorted spill run. + /// -> ProducingOutput + /// Input was exhausted without spilling. Mark every remaining group as + /// complete and produce its final result. + /// -> PreparingMergeInput + /// Input was exhausted after spilling. Spill the last in-memory run and + /// construct the ordered input used to merge all spill files. + /// + /// Spilling + /// -> ReadingInput + /// One sorted run was written; resume reading the original input. + /// + /// PreparingMergeInput + /// Spill the final in-memory run and build the input ordered replay stream. + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. + /// + /// MergingSpills + /// Aggregate the merged spill runs and emit final results. + /// -> MergingSpills + /// Forward one result batch from the fully ordered replay stream that + /// consumes the sort-preserving merge. + /// -> Done + /// The merged spill input was fully aggregated. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One remaining final aggregate batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("OrderedSingleAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ OrderedSingleAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ OrderedSingleAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ OrderedSingleAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ OrderedSingleAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } + state @ OrderedSingleAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ OrderedSingleAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } + state @ OrderedSingleAggregateState::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::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!( + next_state, + OrderedSingleAggregateState::Error + )); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(OrderedSingleAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for OrderedSingleAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +}