diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index f0be10bc6c797..4280e938ef8a2 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -796,10 +796,10 @@ config_namespace! { /// Temporary switch for aggregate stream implementations that are being /// migrated from `GroupedHashAggregateStream`. /// - /// When set to true, DataFusion tries the migrated implementations when - /// their preconditions are satisfied. When set to false, grouped - /// aggregation falls back to `GroupedHashAggregateStream`. This option - /// will be removed after the migration is finished. + /// When set to true, grouped hash aggregation executes only migrated + /// implementations. When set to false, grouped hash aggregation uses + /// `GroupedHashAggregateStream`. This option will be removed after the + /// migration is finished. /// /// See for details. pub enable_migration_aggregate: bool, default = true 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..f53810e6dc9b0 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; @@ -181,7 +184,6 @@ use datafusion_common::{ internal_err, not_impl_err, }; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::MemoryLimit; use datafusion_expr::{Accumulator, Aggregate}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -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), } @@ -1004,7 +1011,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 +1028,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,14 +1225,27 @@ 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, )?)); } + + return internal_err!( + "No migrated aggregate stream supports mode={:?}, input_order_mode={:?}, grouping_sets={}", + self.mode, + self.input_order_mode, + self.group_by.has_grouping_set() + ); } - // Execution paths that have not been migrated use the fallback implementation + // Migration-disabled execution uses the legacy implementation. Ok(StreamType::GroupedHash(GroupedHashAggregateStream::new( self, context, partition, )?)) @@ -1229,8 +1255,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,27 +1265,19 @@ 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() } - fn should_use_partial_reduce_hash_stream(&self, context: &TaskContext) -> bool { - // TODO: implement memory-limited path and remove this limitation - if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { - return false; - } - + fn should_use_partial_reduce_hash_stream(&self, _context: &TaskContext) -> bool { 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 +1287,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 +1304,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 +1428,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 +3288,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 +4117,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), @@ -4151,10 +4263,10 @@ mod tests { Ok(()) } - /// Spilling behavior is not implemented for partial-reduce stream yet, so fall - /// back to the existing `GroupedHashAggregateStream` + /// Partial-reduce hash aggregation returns `ResourcesExhausted` when its + /// reservation cannot grow. #[tokio::test] - async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { + async fn partial_reduce_aggregate_with_memory_limit_returns_oom() -> Result<()> { let partial_reduce = partial_reduce_test_aggregate()?; let runtime = RuntimeEnvBuilder::new() .with_memory_limit(1, 1.0) @@ -4170,7 +4282,13 @@ mod tests { ); let stream = partial_reduce.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::GroupedHash(_))); + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let err = collect(stream).await.unwrap_err(); + assert!( + matches!(err.find_root(), DataFusionError::ResourcesExhausted(_)), + "expected ResourcesExhausted, got: {err}" + ); Ok(()) } @@ -5528,6 +5646,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) + } +} diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 2f4535e66f4ef..d6dcf19368b34 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -28,7 +28,7 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; @@ -68,6 +68,12 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// This stage is useful for tree-reduce plans. It consumes the same schema as /// a final aggregate stage, but emits the same schema as a partial aggregate /// stage. +/// +/// # Memory Management +/// +/// If memory usage exceeds the budget, the stream returns an execution error. +/// +/// Larger-than-memory execution is left for future work. pub(crate) struct PartialReduceHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -97,6 +103,10 @@ enum PartialReduceHashAggregateState { hash_table: AggregateHashTable, }, 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 PartialReduceHashAggregatePoll = Poll>>; @@ -114,7 +124,9 @@ impl PartialReduceHashAggregateState { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table } - Self::Done => unreachable!("Done state does not hold a hash table"), + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") + } } } @@ -123,7 +135,9 @@ impl PartialReduceHashAggregateState { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table } - Self::Done => unreachable!("Done state does not hold a hash table"), + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") + } } } @@ -132,7 +146,9 @@ impl PartialReduceHashAggregateState { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table } - Self::Done => unreachable!("Done state does not hold a hash table"), + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") + } } } @@ -184,13 +200,18 @@ impl PartialReduceHashAggregateStream { }) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { + fn close_input(&mut self) { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - hash_table.start_output() + } + + fn break_with_err( + error: DataFusionError, + ) -> PartialReduceHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialReduceHashAggregateState::Error, + )) } /// Handle ReadingInput state - aggregate partial state batches into the hash table. @@ -219,41 +240,33 @@ impl PartialReduceHashAggregateStream { timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } + // If OOM, return execution error. if let Err(e) = self .reservation .try_resize(original_state.hash_table().memory_size()) { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } ControlFlow::Continue(original_state) } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), // Input ends, move to output state Poll::Ready(None) => { + self.close_input(); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); + let result = original_state.hash_table_mut().start_output(); timer.done(); match result { Ok(()) => { ControlFlow::Continue(original_state.into_producing_output()) } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Err(e) => Self::break_with_err(e), } } } @@ -281,9 +294,12 @@ impl PartialReduceHashAggregateStream { match result { Ok(Some(batch)) => { - let _ = self + if let Err(e) = self .reservation - .try_resize(original_state.hash_table().memory_size()); + .try_resize(original_state.hash_table().memory_size()) + { + return Self::break_with_err(e); + } debug_assert!(batch.num_rows() > 0); let next_state = if original_state.hash_table().is_done() { original_state.into_done() @@ -300,7 +316,7 @@ impl PartialReduceHashAggregateStream { let _ = self.reservation.try_resize(0); ControlFlow::Continue(original_state.into_done()) } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + Err(e) => Self::break_with_err(e), } } } @@ -337,6 +353,13 @@ impl Stream for PartialReduceHashAggregateStream { /// -> Done /// All merged partial-state output was emitted. /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// /// Done /// -> (end) /// ``` @@ -357,6 +380,12 @@ impl Stream for PartialReduceHashAggregateStream { state @ PartialReduceHashAggregateState::ProducingOutput { .. } => { self.handle_producing_output(state) } + state @ PartialReduceHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } state @ PartialReduceHashAggregateState::Done => { let _ = self.reservation.try_resize(0); self.state = Some(state); @@ -369,6 +398,19 @@ impl Stream for PartialReduceHashAggregateStream { self.state = Some(next_state); continue; } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!( + next_state, + PartialReduceHashAggregateState::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(PartialReduceHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } ControlFlow::Break((poll, next_state)) => { self.state = Some(next_state); return poll; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 90bb55b0f0d47..48f8a1a6ac771 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -379,7 +379,7 @@ datafusion.execution.coalesce_batches true When set to true, record batches will datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Defaults to true. datafusion.execution.enable_ansi_mode false Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, or invalid casts raise runtime errors rather than returning `NULL` or adjusted values. - **Consistent ANSI behavior** for string concatenation, comparisons, and `NULL` handling. When `enable_ansi_mode` is `false` (the default), the engine uses a more permissive, non-ANSI mode designed for user convenience and backward compatibility. In this mode: - Implicit casts between types are allowed (e.g., string to integer when possible). - Arithmetic operations are more lenient — for example, `abs()` on the minimum representable integer value returns the input value instead of raising overflow. - Division by zero or invalid casts may return `NULL` instead of failing. # Default `false` — ANSI SQL mode is disabled by default. datafusion.execution.enable_file_stream_work_stealing true When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. -datafusion.execution.enable_migration_aggregate true Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, DataFusion tries the migrated implementations when their preconditions are satisfied. When set to false, grouped aggregation falls back to `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See for details. +datafusion.execution.enable_migration_aggregate true Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, grouped hash aggregation executes only migrated implementations. When set to false, grouped hash aggregation uses `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See for details. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 2811fb4df2900..30b6403ddabe5 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -1188,3 +1188,9 @@ See [PR #24030](https://github.com/apache/datafusion/pull/24030) for details. [`bufwriter`]: https://docs.rs/tokio/latest/tokio/io/struct.BufWriter.html [`asyncarrowwriter`]: https://docs.rs/parquet/59.1.0/parquet/arrow/async_writer/struct.AsyncArrowWriter.html [`parquet/examples/object_store.rs`]: https://github.com/apache/arrow-rs/blob/main/parquet/examples/object_store.rs + +### `PartialReduce` aggregation no longer spills + +`AggregateMode::PartialReduce` previously used a spill-capable aggregation path +when execution memory was bounded. It now returns an execution error when it runs out +of memory space. Larger-than-memory execution for `PartialReduce` is left for future work. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 860884e11fbf1..2b3ebe0fe84df 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -121,7 +121,7 @@ The following configuration settings are available: | datafusion.execution.parquet.content_defined_chunking.norm_level | 0 | Normalization level. Increasing this improves deduplication ratio but increases fragmentation. Recommended range is [-3, 3], default is 0. | | datafusion.execution.planning_concurrency | 0 | Fan-out during initial physical planning. This is mostly use to plan `UNION` children in parallel. Defaults to the number of CPU cores on the system | | datafusion.execution.skip_physical_aggregate_schema_check | false | When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. | -| datafusion.execution.enable_migration_aggregate | true | Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, DataFusion tries the migrated implementations when their preconditions are satisfied. When set to false, grouped aggregation falls back to `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See for details. | +| datafusion.execution.enable_migration_aggregate | true | Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, grouped hash aggregation executes only migrated implementations. When set to false, grouped hash aggregation uses `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See for details. | | datafusion.execution.spill_compression | uncompressed | Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. | | datafusion.execution.sort_spill_reservation_bytes | 10485760 | Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). | | datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. |