diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index f2b1f66c28ba4..c101ba50e4146 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -2086,7 +2086,7 @@ fn test_pushdown_grouping_sets_filter_on_common_column() { - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true output: Ok: - - AggregateExec: mode=Final, gby=[(a@0 as a, b@1 as b), (NULL as a, b@1 as b)], aggr=[cnt], ordering_mode=PartiallySorted([1]) + - AggregateExec: mode=Final, gby=[(a@0 as a, b@1 as b), (NULL as a, b@1 as b)], aggr=[cnt] - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar " ); @@ -2372,7 +2372,7 @@ fn test_pushdown_through_aggregate_grouping_sets_with_reordered_input() { - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true output: Ok: - - AggregateExec: mode=Final, gby=[(a@1 as a, b@2 as b), (NULL as a, b@2 as b)], aggr=[cnt], ordering_mode=PartiallySorted([1]) + - AggregateExec: mode=Final, gby=[(a@1 as a, b@2 as b), (NULL as a, b@2 as b)], aggr=[cnt] - ProjectionExec: expr=[c@2 as c, a@0 as a, b@1 as b] - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar " diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 91e9d6555c3e7..d895b9aeb28aa 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -15,10 +15,11 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::array::{ArrayRef, AsArray, BooleanArray, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; @@ -32,6 +33,7 @@ use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, + group_id_array, max_duplicate_ordinal, }; /// Marker for raw rows -> partial state aggregation. @@ -213,6 +215,75 @@ impl AggregateHashTable { Ok(()) } + /// Creates the required empty grouping-set rows for a raw-input aggregate. + /// + /// For example, this query must still produce one grand-total group even if + /// `t` has no rows: + /// + /// ```sql + /// SELECT COUNT(v) + /// FROM t + /// GROUP BY GROUPING SETS (()); + /// ``` + /// + /// The synthetic row is filtered out before accumulator update so aggregates + /// see the same state they would see for an empty input, rather than a real + /// null-valued row. Only partial and single aggregation call this method; + /// final and partial-reduce aggregation consume states produced upstream. + pub(super) fn init_empty_grouping_sets_for_raw_input(&mut self) -> Result<()> { + let state = self.state.building_mut(); + if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + return Ok(()); + } + + let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let n_expr = state.group_by.expr().len(); + let mut any_interned = false; + + for group in state.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let ordinal = *entry; + *entry += 1; + ordinal + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + let mut cols: Vec = group_schema + .fields() + .iter() + .take(n_expr) + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + state + .group_values + .intern(&cols, &mut state.batch_group_indices)?; + any_interned = true; + } + + if any_interned { + let total_groups = state.group_values.len(); + let false_filter = BooleanArray::from(vec![false]); + for acc in state.accumulators.iter_mut() { + let null_args = acc.null_arguments(&self.input_schema)?; + let values = EvaluatedAccumulatorArgs { + arguments: null_args, + filter: Some(Arc::new(false_filter.clone())), + }; + acc.update_batch(&values, &[0], total_groups)?; + } + } + + Ok(()) + } + /// Materializes the full output once, then returns it downstream incrementally /// by slicing it into `batch_size` chunks. /// 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/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index a64fd32536eeb..4550140abb511 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -15,22 +15,20 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; +use crate::aggregates::AggregateExec; use crate::aggregates::group_values::new_group_values; use crate::aggregates::order::GroupOrdering; -use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, - EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, + HashAggregateAccumulator, PartialMarker, PartialSkipMarker, }; /// Implementation specific to partial aggregation, where the table stores @@ -110,78 +108,10 @@ impl AggregateHashTable { } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { - self.init_empty_grouping_sets()?; + self.init_empty_grouping_sets_for_raw_input()?; self.start_outputting(); Ok(()) } - - /// Creates the required empty grouping-set rows when the input is empty. - /// - /// For example, this query must still produce one grand-total group even if - /// `t` has no rows: - /// - /// ```sql - /// SELECT COUNT(v) - /// FROM t - /// GROUP BY GROUPING SETS (()); - /// ``` - /// - /// The synthetic row is filtered out before accumulator update so aggregates - /// see the same state they would see for an empty input, rather than a real - /// null-valued row. - fn init_empty_grouping_sets(&mut self) -> Result<()> { - let state = self.state.building_mut(); - if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { - return Ok(()); - } - - let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); - let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); - let group_schema = state.group_by.group_schema(&self.input_schema)?; - let n_expr = state.group_by.expr().len(); - let mut any_interned = false; - - for group in state.group_by.groups() { - let ordinal = { - let entry = ordinals.entry(group.as_slice()).or_insert(0); - let ordinal = *entry; - *entry += 1; - ordinal - }; - - if !group.iter().all(|&is_null| is_null) { - continue; - } - - let mut cols: Vec = group_schema - .fields() - .iter() - .take(n_expr) - .map(|field| new_null_array(field.data_type(), 1)) - .collect(); - cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); - - state - .group_values - .intern(&cols, &mut state.batch_group_indices)?; - any_interned = true; - } - - if any_interned { - let total_groups = state.group_values.len(); - let false_filter = BooleanArray::from(vec![false]); - for acc in state.accumulators.iter_mut() { - let null_args = acc.null_arguments(&self.input_schema)?; - let values = EvaluatedAccumulatorArgs { - arguments: null_args, - filter: Some(Arc::new(false_filter.clone())), - }; - acc.update_batch(&values, &[0], total_groups)?; - } - } - - Ok(()) - } } impl AggregateHashTable { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs index 56d601c793206..aae90124fa036 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs @@ -70,6 +70,7 @@ impl AggregateHashTable { } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.init_empty_grouping_sets_for_raw_input()?; self.start_outputting(); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6af67d048256e..26c42cc853e95 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, }; @@ -165,7 +167,8 @@ use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, - InputOrderMode, SendableRecordBatchStream, Statistics, check_if_same_properties, + InputOrderMode, Partitioning, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use datafusion_common::config::ConfigOptions; use parking_lot::Mutex; @@ -209,6 +212,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 +676,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 +706,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), } @@ -1004,7 +1012,9 @@ impl AggregateExec { .filter(|idx| group_by.groups.iter().all(|group| !group[*idx])) .collect(); - let input_order_mode = if indices.len() == groupby_exprs.len() + let input_order_mode = if group_by.has_grouping_set() { + InputOrderMode::Linear + } else if indices.len() == groupby_exprs.len() && !indices.is_empty() && group_by.groups.len() == 1 { @@ -1019,15 +1029,19 @@ impl AggregateExec { let group_expr_mapping = ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?; - let cache = Self::compute_properties( - &input, - Arc::clone(&schema), - &group_expr_mapping, - group_by.is_true_no_grouping(), - &mode, - &input_order_mode, - aggr_expr.as_ref(), - )?; + let cache = if group_by.has_grouping_set() { + Self::compute_grouping_set_properties(&input, Arc::clone(&schema)) + } else { + Self::compute_properties( + &input, + Arc::clone(&schema), + &group_expr_mapping, + group_by.is_true_no_grouping(), + &mode, + &input_order_mode, + aggr_expr.as_ref(), + )? + }; let mut exec = AggregateExec { mode, @@ -1212,6 +1226,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, @@ -1229,8 +1249,6 @@ impl AggregateExec { self.mode == AggregateMode::Partial && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() - && self.group_by.is_single() - && self.limit_options_supported_by_hash_stream() } fn should_use_ordered_partial_aggregate_stream( @@ -1241,15 +1259,13 @@ impl AggregateExec { && self.input_order_mode != InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() - && self.limit_options_supported_by_hash_stream() } fn should_use_final_hash_stream(&self, _context: &TaskContext) -> bool { matches!( self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned - ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode == InputOrderMode::Linear + ) && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1261,7 +1277,6 @@ impl AggregateExec { } self.mode == AggregateMode::PartialReduce - && self.limit_options.is_none() && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() @@ -1271,8 +1286,15 @@ impl AggregateExec { matches!( self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned - ) && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + ) && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + } + + fn should_use_ordered_single_aggregate_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + ) && self.input_order_mode != InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1281,17 +1303,11 @@ impl AggregateExec { matches!( self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned - ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode != InputOrderMode::Linear + ) && self.input_order_mode != InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } - /// See comments in `PartialHashAggregateStream` limit optimization section - fn limit_options_supported_by_hash_stream(&self) -> bool { - self.limit_options.is_none() || self.is_unordered_unfiltered_group_by_distinct() - } - /// Finds the DataType and SortDirection for this Aggregate, if there is one pub fn get_minmax_desc(&self) -> Option<(FieldRef, bool)> { let agg_expr = self.aggr_expr.iter().exactly_one().ok()?; @@ -1411,6 +1427,22 @@ impl AggregateExec { )) } + fn compute_grouping_set_properties( + input: &Arc, + schema: SchemaRef, + ) -> PlanProperties { + // Grouping-set expansion can replace group keys with nulls and adds a + // grouping ID, so input properties do not project through unchanged. + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning( + input.output_partitioning().partition_count(), + ), + EmissionType::Final, + input.boundedness(), + ) + } + pub fn input_order_mode(&self) -> &InputOrderMode { &self.input_order_mode } @@ -3255,12 +3287,14 @@ mod tests { | 2 | 1.0 | 0 | 1 | | 3 | | 1 | 1 | | 3 | | 1 | 2 | - | 3 | 2.0 | 0 | 2 | + | 3 | 2.0 | 0 | 1 | + | 3 | 2.0 | 0 | 1 | | 3 | 3.0 | 0 | 1 | | 4 | | 1 | 1 | | 4 | | 1 | 2 | | 4 | 3.0 | 0 | 1 | - | 4 | 4.0 | 0 | 2 | + | 4 | 4.0 | 0 | 1 | + | 4 | 4.0 | 0 | 1 | +---+-----+---------------+-----------------+ " ); @@ -4082,6 +4116,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), @@ -5528,6 +5639,97 @@ mod tests { Ok(()) } + #[tokio::test] + async fn grouping_sets_use_unordered_hash_streams() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![1])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![2, 1])), + ], + )?, + ]; + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("a", &schema)?)]) + .unwrap(); + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + + let grouping_set = PhysicalGroupBy::new( + vec![ + (col("a", &schema)?, "a".to_string()), + (col("b", &schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Int32(None)), "a".to_string()), + (lit(ScalarValue::Int32(None)), "b".to_string()), + ], + vec![vec![false, true], vec![false, false]], + true, + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + grouping_set.clone(), + vec![], + vec![], + Arc::clone(&input) as Arc, + Arc::clone(&schema), + )?; + + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + assert_eq!(aggregate.maintains_input_order(), vec![false]); + assert!(aggregate.properties().output_ordering().is_none()); + assert!( + aggregate + .properties() + .eq_properties + .constraints() + .is_empty() + ); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let output = collect(stream.into()).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++---+---+---------------+ +| a | b | __grouping_id | ++---+---+---------------+ +| 1 | | 1 | +| 1 | 1 | 0 | +| 1 | 2 | 0 | +| 2 | | 1 | +| 2 | 1 | 0 | ++---+---+---------------+ +"); + + let partial = AggregateExec::try_new( + AggregateMode::Partial, + grouping_set, + vec![], + vec![], + input, + Arc::clone(&schema), + )?; + let stream = partial.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialHash(_))); + + Ok(()) + } + #[test] fn group_exprs_nullable() -> Result<()> { let input_schema = Arc::new(Schema::new(vec![ 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) + } +}