diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs index 4726e7c4aca5c..f9e7f2e10e789 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs @@ -350,7 +350,12 @@ async fn run_aggregate_test(input1: Vec, group_by_columns: Vec<&str schema.clone(), ) .unwrap(), - ) as Arc; + ); + assert_ne!( + aggregate_exec_running.input_order_mode(), + &InputOrderMode::Linear, + "running aggregate should observe ordered input for group_by: {group_by:?}" + ); let aggregate_exec_usual = Arc::new( AggregateExec::try_new( @@ -362,7 +367,7 @@ async fn run_aggregate_test(input1: Vec, group_by_columns: Vec<&str schema.clone(), ) .unwrap(), - ) as Arc; + ); let task_ctx = ctx.task_ctx(); let collected_usual = collect(aggregate_exec_usual.clone(), task_ctx.clone()) 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 719fbe93e5416..532aad53920a8 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -246,6 +246,8 @@ pub(super) struct HashAggregateAccumulator { accumulator: Box, } +pub(super) type AggregateAccumulator = HashAggregateAccumulator; + /// Evaluated aggregate arguments and filter for one input batch. /// /// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` @@ -348,7 +350,7 @@ impl MaterializedFinalOutput { } impl HashAggregateAccumulator { - fn new( + pub(super) fn new( aggregate_expr: Arc, arguments: Vec>, filter: Option>, @@ -380,7 +382,10 @@ impl HashAggregateAccumulator { /// and `x > 0`. /// /// These arrays can be passed directly to [`GroupsAccumulator`] next. - fn evaluate_acc_args(&self, batch: &RecordBatch) -> Result { + pub(super) fn evaluate_acc_args( + &self, + batch: &RecordBatch, + ) -> Result { let arguments = self .arguments .iter() @@ -403,6 +408,10 @@ impl HashAggregateAccumulator { Ok(EvaluatedAccumulatorArgs { arguments, filter }) } + pub(super) fn size(&self) -> usize { + self.accumulator.size() + } + pub(super) fn update_batch( &mut self, values: &EvaluatedAccumulatorArgs, 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 new file mode 100644 index 0000000000000..a98cc92c2031f --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -0,0 +1,357 @@ +// 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. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::assert_or_internal_err; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::row_hash::create_group_accumulator; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; + +/// Aggregate table shared by the ordered partial and final paths. +/// +/// # Ordering optimization +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// 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 +/// +/// The partial and final aggregate tables implement the two stages of grouped +/// aggregation. See +/// [`OrderedPartialAggregateStream`](crate::aggregates::ordered_partial_stream::OrderedPartialAggregateStream) +/// for the high-level plan shape. +/// +/// Example: `AVG(v) FILTER (WHERE v>0) GROUP BY k` +/// +/// Partial table ([`AggregateMode::Partial`], with optional filter from query): +/// - Input rows: `k, v` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, sum(v), count(v)` +/// +/// Final table ([`AggregateMode::Final`], no filters): +/// - Input rows: `k, sum(v), count(v)` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, avg(v)` +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl`; 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. + pub(super) output_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the +/// input ordering proves those groups are complete. +/// +/// [`GroupOrdering`] tracks when and how to do early emit. +/// [`GroupValues`] stores the physical group-key layout, while +/// [`datafusion_expr::GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct OrderedAggregateTableBuffer { + /// GROUP BY expressions evaluated against input batches. + pub(super) group_by: Arc, + + /// Tracks how far ordered input allows this table to drain safely. + pub(super) group_ordering: GroupOrdering, + + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec, +} + +/// Methods shared by all aggregate modes +impl OrderedAggregateTable { + #[expect( + clippy::too_many_arguments, + reason = "keeps ordered partial and final table construction explicit" + )] + pub(super) fn new_for_mode( + agg: &AggregateExec, + partition: usize, + input_schema: &SchemaRef, + output_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + aggregate_mode: &AggregateMode, + filters: Vec>>, + ) -> Result { + assert_or_internal_err!( + batch_size > 0, + "OrderedAggregateTable requires config batch_size >= 1" + ); + + let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_schema = agg.group_by.group_schema(input_schema)?; + let group_values = new_group_values(group_schema, &group_ordering)?; + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + aggregate_mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(AggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::>()?; + + Ok(Self { + output_schema, + batch_size, + group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + buffer: OrderedAggregateTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_ordering, + group_values, + group_indices: vec![], + accumulators, + }, + _mode: PhantomData, + }) + } + + /// Evaluates all group by keys and accumulator args. + /// + /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function + /// evaluates `k+1`, `v*v`. + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result { + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + let grouping_set_args = evaluate_group_by(&self.buffer.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + let accumulator_args = self + .buffer + .accumulators + .iter() + .map(|acc| acc.evaluate_acc_args(batch)) + .collect::>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Called after the input stream is exhausted and the last batch has been + /// aggregated. + /// + /// Updates the internal `GroupOrdering` so it can continue emitting until + /// the buffer is empty. + pub(in crate::aggregates) fn input_done(&mut self) { + self.buffer.group_ordering.input_done(); + } + + /// Check if there is zero groups accumulated so far. + pub(in crate::aggregates) fn is_empty(&self) -> bool { + self.buffer.group_values.is_empty() + } + + /// All internal buffer's memory size. + pub(in crate::aggregates) fn memory_size(&self) -> usize { + self.buffer + .accumulators + .iter() + .map(|acc| acc.size()) + .sum::() + + self.buffer.group_values.size() + + self.buffer.group_ordering.size() + + self.buffer.group_indices.allocated_size() + } + + /// Returns the [`EmitTo`], clamped to the specified batch size + /// + /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number + /// of groups to emit from `GroupValues` / accumulators, and + /// `should_remove_groups` indicates whether `GroupOrdering` must also shift + /// its tracked indexes. + pub(super) fn clamp_emit_to( + &self, + group_count: usize, + emit_to: EmitTo, + ) -> (EmitTo, bool) { + match emit_to { + EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), + EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), + 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` + /// + /// - `true`: merge partial aggregate states for final aggregation. + /// - `false`: update aggregate states from raw input for partial aggregation. + pub(super) fn aggregate_evaluated_batch( + &mut self, + evaluated_batch: &EvaluatedAggregateBatch, + is_final: bool, + ) -> Result<()> { + for group_values in &evaluated_batch.grouping_set_args { + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); + if total_num_groups > starting_num_groups { + self.buffer.group_ordering.new_groups( + group_values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + + let timer = self.group_by_metrics.aggregation_time.timer(); + for (acc, values) in self + .buffer + .accumulators + .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, + )?; + } + } + drop(timer); + } + + Ok(()) + } + + /// 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( + &mut self, + is_final: bool, + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let Some(emit_to) = self.buffer.group_ordering.emit_to() else { + return Ok(None); + }; + let (emit_to, should_remove_groups) = + self.clamp_emit_to(self.buffer.group_values.len(), emit_to); + + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.buffer.group_values.emit(emit_to)?; + if should_remove_groups { + match emit_to { + EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), + // `EmitTo::All` is only used after `input_done`, when all + // buffered groups are known complete and the ordering state is + // no longer needed. + EmitTo::All => {} + } + } + + for acc in &mut self.buffer.accumulators { + if is_final { + output.push(acc.evaluate(emit_to)?); + } else { + output.extend(acc.state(emit_to)?); + } + } + drop(timer); + + let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + Ok(Some(batch)) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index c3e4f831c4bbf..c5c3d4c5493c5 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -29,7 +29,13 @@ use super::common::{ MaterializedFinalOutput, }; -/// Methods specific to the aggregate hash table used in the final aggregation stage. +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` impl AggregateHashTable { pub(in crate::aggregates) fn new( agg: &AggregateExec, 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 eb152f4128896..2bb1d119f0d61 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -16,9 +16,13 @@ // under the License. mod common; +mod common_ordered; mod final_table; +mod ordered_final_table; +mod ordered_partial_table; mod partial_table; pub(super) use common::{ AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, }; +pub(super) use common_ordered::OrderedAggregateTable; 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 new file mode 100644 index 0000000000000..b7e3fd38edf25 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -0,0 +1,81 @@ +// 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 final aggregation when partial-state 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::InputOrderMode; +use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new_with_input_order( + agg: &AggregateExec, + partition: usize, + input_schema: &SchemaRef, + output_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + ) -> Result { + Self::new_for_mode( + agg, + partition, + input_schema, + output_schema, + batch_size, + input_order_mode, + &AggregateMode::Final, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Merges one partial-state 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)?; + // `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) + } + + /// 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) + } +} 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 new file mode 100644 index 0000000000000..033c14056a419 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -0,0 +1,100 @@ +// 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 partial aggregation when input is ordered by group keys. +//! +//! See the [`super::common_ordered`] comments for the high-level ideas. +//! +//! This operator handles input that is ordered by group keys: +//! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b` +//! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a` +//! +//! When a group key combination is exhausted, this table eagerly flushes the +//! completed groups to improve memory efficiency. +//! +//! The implementation is separated from other aggregate tables because this +//! execution path is likely to be optimized further in the future. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::{ + AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker, +}; + +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(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, + batch_size: usize, + ) -> Result { + let input_schema = agg.input().schema(); + Self::new_for_mode( + agg, + partition, + &input_schema, + output_schema, + batch_size, + &agg.input_order_mode, + &AggregateMode::Partial, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// 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, false) + } + + /// Emits the next batch of partial state rows for groups proven complete by + /// the input ordering. + /// + /// For example, when the query is `GROUP BY a` and the input is ordered by + /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` + /// are complete and safe to emit. + /// + /// Key steps: + /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly. + /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and + /// all `GroupsAccumulator`s. + /// + /// This may output small batches. Avoiding tiny batches is left to future + /// ordered-aggregation optimizations. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_for_mode(false) + } +} 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 9d226aa28b35f..fb30d72dbd958 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 @@ -34,7 +34,13 @@ use super::common::{ emit_to_for_batch_size, }; -/// Methods specific to the aggregate hash table used in the partial aggregation stage. +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` impl AggregateHashTable { pub(in crate::aggregates) fn new( agg: &AggregateExec, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 4f5b893578d74..45d2861ec5244 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -24,6 +24,8 @@ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::aggregates::{ hash_aggregate::{FinalHashAggregateStream, PartialHashAggregateStream}, no_grouping::AggregateStream, + ordered_final_stream::OrderedFinalAggregateStream, + ordered_partial_stream::OrderedPartialAggregateStream, row_hash::GroupedHashAggregateStream, topk_stream::GroupedTopKAggregateStream, }; @@ -77,6 +79,8 @@ pub mod group_values; mod hash_aggregate; mod no_grouping; pub mod order; +mod ordered_final_stream; +mod ordered_partial_stream; mod row_hash; mod skip_partial; mod topk; @@ -530,10 +534,16 @@ enum StreamType { /// Final stage of the hash aggregation /// Input output scheme: partial state -> final result FinalHash(FinalHashAggregateStream), + /// Partial stage of aggregation for ordered input. + OrderedPartialAggregate(OrderedPartialAggregateStream), + /// Final stage of aggregation for ordered input. + OrderedFinalAggregate(OrderedFinalAggregateStream), /// Hash aggregation reused for multiple stages /// /// Note this is being incrementally migrated to dedicated streams like - /// [`StreamType::PartialHash`] and [`StreamType::FinalHash`] + /// [`StreamType::PartialHash`], [`StreamType::FinalHash`], + /// [`StreamType::OrderedPartialAggregate`], and + /// [`StreamType::OrderedFinalAggregate`] /// /// See issue for details: GroupedHash(GroupedHashAggregateStream), @@ -551,6 +561,8 @@ impl From for SendableRecordBatchStream { StreamType::AggregateStream(stream) => Box::pin(stream), StreamType::PartialHash(stream) => Box::pin(stream), StreamType::FinalHash(stream) => Box::pin(stream), + StreamType::OrderedPartialAggregate(stream) => Box::pin(stream), + StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } @@ -989,6 +1001,8 @@ impl AggregateExec { Arc::clone(&self.input_schema) } + /// Aggregation has multiple specialized implementations optimized for + /// different workloads. This function picks the best available path. fn execute_typed( &self, partition: usize, @@ -1022,12 +1036,24 @@ impl AggregateExec { .execution .enable_migration_aggregate { + if self.should_use_ordered_partial_aggregate_stream(context) { + return Ok(StreamType::OrderedPartialAggregate( + OrderedPartialAggregateStream::new(self, context, partition)?, + )); + } + if self.should_use_partial_hash_stream(context) { return Ok(StreamType::PartialHash(PartialHashAggregateStream::new( self, context, partition, )?)); } + if self.should_use_ordered_final_aggregate_stream(context) { + return Ok(StreamType::OrderedFinalAggregate( + OrderedFinalAggregateStream::new(self, context, partition)?, + )); + } + if self.should_use_final_hash_stream(context) { return Ok(StreamType::FinalHash(FinalHashAggregateStream::new( self, context, partition, @@ -1054,6 +1080,19 @@ impl AggregateExec { && self.limit_options_supported_by_hash_stream() } + fn should_use_ordered_partial_aggregate_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; + } + + 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_final_hash_stream(&self, context: &TaskContext) -> bool { // TODO: implement memory-limited path and remove this limitation if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { @@ -1069,6 +1108,21 @@ impl AggregateExec { && self.group_by.is_single() } + fn should_use_ordered_final_aggregate_stream(&self, context: &TaskContext) -> bool { + // TODO: implement memory-limited path and remove this limitation + if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { + return false; + } + + matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && self.limit_options_supported_by_hash_stream() + && self.input_order_mode != InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + /// 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() @@ -2461,7 +2515,7 @@ mod tests { use crate::projection::ProjectionExec; use datafusion_physical_expr::projection::ProjectionExpr; - use futures::{FutureExt, Stream}; + use futures::{FutureExt, Stream, StreamExt}; use insta::{allow_duplicates, assert_snapshot}; // Generate a schema which consists of 5 columns (a, b, c, d, e) @@ -2571,6 +2625,34 @@ mod tests { Arc::new(task_ctx) } + fn migrated_hash_session_config(batch_size: usize) -> SessionConfig { + SessionConfig::new() + .with_batch_size(batch_size) + .set_bool("datafusion.execution.enable_migration_aggregate", true) + } + + fn new_migrated_hash_ctx(batch_size: usize) -> Arc { + Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)), + ) + } + + fn new_finite_memory_migrated_hash_ctx( + batch_size: usize, + max_memory: usize, + ) -> Result> { + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(max_memory, 1.0) + .build_arc()?; + + Ok(Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(migrated_hash_session_config(batch_size)), + )) + } + async fn check_grouping_sets( input: Arc, spill: bool, @@ -3326,6 +3408,246 @@ mod tests { Ok(()) } + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. + #[tokio::test] + async fn ordered_partial_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, 1, 1])), + ], + )?, + 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![1, 1])), + ], + )?, + ]; + 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 group_by = PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value_col)") + .build()?, + )]; + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + 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::OrderedPartialAggregate(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++----------+-----------+-------------------------+ +| sort_col | group_col | COUNT(value_col)[count] | ++----------+-----------+-------------------------+ +| 1 | 10 | 2 | +| 1 | 11 | 1 | +| 2 | 20 | 1 | +| 2 | 21 | 1 | ++----------+-----------+-------------------------+ +"); + + // Ordered streams don't implement memory limits yet. + 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::GroupedHash(_))); + + Ok(()) + } + + /// Ensures for ordered input, `OrderedFinalAggregateStream` is used. + #[tokio::test] + async fn ordered_final_aggregate_planning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial_aggregate = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggr_expr.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial_aggregate.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 3])), + Arc::new(Int64Array::from(vec![2, 3, 5, 7])), + ], + )?; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("key", 0), + ))]) + .unwrap(); + let final_input = + TestMemoryExec::try_new(&[vec![partial_state_batch]], partial_schema, None)? + .try_with_sort_information(vec![ordering])?; + let final_input = Arc::new(TestMemoryExec::update_cache(&Arc::new(final_input))); + + let final_aggregate = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggr_expr, + vec![None], + final_input, + Arc::clone(&schema), + )?; + assert_eq!(final_aggregate.input_order_mode(), &InputOrderMode::Sorted); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedFinalAggregate(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+--------------+ +| key | COUNT(value) | ++-----+--------------+ +| 1 | 5 | +| 2 | 5 | +| 3 | 7 | ++-----+--------------+ +"); + + // Ordered streams don't implement memory limits yet. + let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + let stream = final_aggregate.execute_typed(0, &finite_memory_task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + Ok(()) + } + + #[tokio::test] + async fn ordered_partial_aggregate_partially_sorted_no_emit_panic() -> Result<()> { + // Reproducer for #20445: emitting from PartiallySorted input must not + // drain more groups than the completed sort boundary allows. + 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), + ])); + + // All rows share sort_col=1, so there is no completed sort boundary + // inside this batch even though there are many distinct groups. + let n = 256; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; n])), + Arc::new(Int32Array::from((0..n as i32).collect::>())), + Arc::new(Int64Array::from(vec![1; n])), + ], + )?; + + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[vec![batch]], 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::Partial, + 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(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )], + vec![None], + input, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(4096, 1.0) + .build_arc()?; + let session_config = SessionConfig::new().with_batch_size(128).set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::UInt64(Some(u64::MAX)), + ); + let task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(session_config), + ); + + let mut stream: SendableRecordBatchStream = Box::pin( + OrderedPartialAggregateStream::new(&aggregate, &task_ctx, 0)?, + ); + + while let Some(result) = stream.next().await { + if let Err(e) = result { + if e.to_string().contains("Resources exhausted") { + break; + } + return Err(e); + } + } + + Ok(()) + } + #[tokio::test] async fn test_drop_cancel_without_groups() -> Result<()> { let task_ctx = Arc::new(TaskContext::default()); diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs new file mode 100644 index 0000000000000..89653e05ab4c7 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -0,0 +1,350 @@ +// 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. + +//! Final aggregate stream for ordered partial-state 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::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable}; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Final aggregate stream for `InputOrderMode::Sorted` and +/// `InputOrderMode::PartiallySorted`. +/// +/// See comments at [`super::ordered_partial_stream`] for details. +pub(crate) struct OrderedFinalAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + state: Option, +} + +/// See comments at `poll_next()` for details. +enum OrderedFinalAggregateState { + ReadingInput { + table: OrderedAggregateTable, + }, + DrainingFinal { + table: OrderedAggregateTable, + }, + Done, +} + +type OrderedFinalAggregatePoll = Poll>>; +type OrderedFinalAggregateStateTransition = ControlFlow< + (OrderedFinalAggregatePoll, OrderedFinalAggregateState), + OrderedFinalAggregateState, +>; + +impl OrderedFinalAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + )); + debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + + let input = agg.input.execute(partition, Arc::clone(context))?; + Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) + } + + pub(in crate::aggregates) fn new_with_input( + agg: &AggregateExec, + context: &Arc, + partition: usize, + input: SendableRecordBatchStream, + input_order_mode: &InputOrderMode, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + )); + debug_assert_ne!(*input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + + let table = OrderedAggregateTable::::new_with_input_order( + agg, + partition, + &input_schema, + Arc::clone(&schema), + batch_size, + input_order_mode, + )?; + let reservation = + MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + state: Some(OrderedFinalAggregateState::ReadingInput { table }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + /// Consumes one ordered partial-state 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: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::ReadingInput { mut table } = original_state + else { + unreachable!("expected reading input state") + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedFinalAggregateState::ReadingInput { table }, + )), + 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 ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { table }, + )); + } + + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + + match result { + // Some finalized groups can be emitted. Yield them, then + // continue aggregating input in the current state. + Ok(Some(batch)) => { + let next_state = + OrderedFinalAggregateState::ReadingInput { table }; + self.resize_reservation_for_state(&next_state); + + ControlFlow::Break(( + Poll::Ready(Some(Ok( + batch.record_output(&self.baseline_metrics) + ))), + next_state, + )) + } + Ok(None) => { + // Ordered variant doesn't support memory-limited + // execution, so it errors when memory reservation fails. + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { table }, + )); + } + + // Can't do early emit, continue aggregating. + ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + table, + }) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { table }, + )), + } + } + Poll::Ready(Some(Err(e))) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { table }, + )), + Poll::Ready(None) => { + self.close_input(); + table.input_done(); + ControlFlow::Continue(OrderedFinalAggregateState::DrainingFinal { table }) + } + } + } + + /// 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_draining_final( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::DrainingFinal { table } = original_state else { + unreachable!("expected draining final 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() { + OrderedFinalAggregateState::Done + } else { + OrderedFinalAggregateState::DrainingFinal { table } + }; + self.resize_reservation_for_state(&next_state); + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::DrainingFinal { table }, + )), + Ok(None) => { + let next_state = OrderedFinalAggregateState::Done; + self.resize_reservation_for_state(&next_state); + ControlFlow::Continue(next_state) + } + } + } + + fn resize_reservation_for_state(&mut self, state: &OrderedFinalAggregateState) { + let new_size = match state { + OrderedFinalAggregateState::ReadingInput { table } + | OrderedFinalAggregateState::DrainingFinal { table } => table.memory_size(), + OrderedFinalAggregateState::Done => 0, + }; + let _ = self.reservation.try_resize(new_size); + } +} + +impl Stream for OrderedFinalAggregateStream { + type Item = Result; + + /// Entry point for the ordered final aggregate state machine. + /// + /// See comments in [`OrderedFinalAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered partial-state input and merging + /// those states into the ordered final aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Merge one input batch. If the ordering proves some groups are + /// complete, yield one final aggregate batch immediately, then continue + /// reading input. Otherwise continue directly with the next input batch. + /// -> DrainingFinal + /// Input was exhausted. Mark the table input as done so every remaining + /// group is safe to emit. + /// + /// DrainingFinal + /// -> DrainingFinal + /// One remaining final aggregate batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("OrderedFinalAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ OrderedFinalAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ OrderedFinalAggregateState::DrainingFinal { .. } => { + self.handle_draining_final(state) + } + state @ OrderedFinalAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for OrderedFinalAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs new file mode 100644 index 0000000000000..b4b7fa073aee0 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -0,0 +1,382 @@ +// 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. + +//! Partial aggregate stream for ordered group 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::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; + +/// Partial 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`, the aggregate can use ordered partial and +/// final stages: +/// +/// ## Plan +/// AggregateExec(stage=final, ordered) +/// -- RepartitionExec(hash(k), preserves_order=true) +/// ---- AggregateExec(stage=partial, ordered) +/// +/// ## Partial Stage Behavior +/// Input: raw rows +/// Output: partial states for all groups (for example, `AVG(x)` emits `SUM(x)` +/// and `COUNT(x)`) +/// +/// ## Final Stage Behavior +/// Input: partial states +/// Output: results for all groups (for example, `AVG(x)` calculated from the +/// state) +/// +/// # 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. +/// +/// ## Implementation Note +/// +/// This is intentionally kept simple and closely maps to +/// `GroupedHashAggregateStream` to finish the refactor sooner. +/// +/// See issue for details: +/// +/// More applicable optimizations are left to future work. +pub(crate) struct OrderedPartialAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + reduction_factor: metrics::RatioMetrics, + state: Option, +} + +/// See comments at `poll_next()` for details. +enum OrderedPartialAggregateState { + ReadingInput { + table: OrderedAggregateTable, + }, + DrainingFinal { + table: OrderedAggregateTable, + }, + Done, +} + +type OrderedPartialAggregatePoll = Poll>>; +type OrderedPartialAggregateStateTransition = ControlFlow< + (OrderedPartialAggregatePoll, OrderedPartialAggregateState), + OrderedPartialAggregateState, +>; + +impl OrderedPartialAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert_eq!(agg.mode, AggregateMode::Partial); + 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 batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reduction_factor = MetricBuilder::new(&agg.metrics) + .with_type(metrics::MetricType::Summary) + .ratio_metrics("reduction_factor", partition); + + let table = OrderedAggregateTable::::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + let reservation = + MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + reduction_factor, + state: Some(OrderedPartialAggregateState::ReadingInput { table }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + /// Consumes one ordered input batch, then immediately emits completed 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: OrderedPartialAggregateState, + ) -> OrderedPartialAggregateStateTransition { + let OrderedPartialAggregateState::ReadingInput { mut table } = original_state + else { + unreachable!("expected reading input state") + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedPartialAggregateState::ReadingInput { table }, + )), + Poll::Ready(Some(Ok(batch))) => { + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + + 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 ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::ReadingInput { table }, + )); + } + + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + + match result { + // There is some previous group results can be emitted: emit + // them, and next continuing aggreagting input (loop in the + // current state) + Ok(Some(batch)) => { + self.reduction_factor.add_part(batch.num_rows()); + let next_state = + OrderedPartialAggregateState::ReadingInput { table }; + self.resize_reservation_for_state(&next_state); + + ControlFlow::Break(( + Poll::Ready(Some(Ok( + batch.record_output(&self.baseline_metrics) + ))), + next_state, + )) + } + Ok(None) => { + // Ordered variant don't support memory-limited execution, + // it have to error when OOM + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::ReadingInput { table }, + )); + } + + // Can't do early emit, continue aggregating. + ControlFlow::Continue( + OrderedPartialAggregateState::ReadingInput { table }, + ) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::ReadingInput { table }, + )), + } + } + Poll::Ready(Some(Err(e))) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::ReadingInput { table }, + )), + // Input has exhausted, move to the final draining stage. + Poll::Ready(None) => { + self.close_input(); + table.input_done(); + ControlFlow::Continue(OrderedPartialAggregateState::DrainingFinal { + table, + }) + } + } + } + + /// 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_draining_final( + &mut self, + original_state: OrderedPartialAggregateState, + ) -> OrderedPartialAggregateStateTransition { + let OrderedPartialAggregateState::DrainingFinal { table } = original_state else { + unreachable!("expected draining final 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)) => { + self.reduction_factor.add_part(batch.num_rows()); + let next_state = if table.is_empty() { + OrderedPartialAggregateState::Done + } else { + OrderedPartialAggregateState::DrainingFinal { table } + }; + self.resize_reservation_for_state(&next_state); + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::DrainingFinal { table }, + )), + Ok(None) => { + let next_state = OrderedPartialAggregateState::Done; + self.resize_reservation_for_state(&next_state); + ControlFlow::Continue(next_state) + } + } + } + + fn resize_reservation_for_state(&mut self, state: &OrderedPartialAggregateState) { + let new_size = match state { + OrderedPartialAggregateState::ReadingInput { table } + | OrderedPartialAggregateState::DrainingFinal { table } => { + table.memory_size() + } + OrderedPartialAggregateState::Done => 0, + }; + let _ = self.reservation.try_resize(new_size); + } +} + +impl Stream for OrderedPartialAggregateStream { + type Item = Result; + + /// Entry point for the ordered partial aggregate state machine. + /// + /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered input and aggregating batches + /// into the ordered partial aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one input batch. If the ordering proves some groups are + /// complete, yield one partial-state batch immediately, then continue + /// reading input. Otherwise continue directly with the next input batch. + /// -> DrainingFinal + /// Input was exhausted. Mark the table input as done so every remaining + /// group is safe to emit. + /// + /// DrainingFinal + /// -> DrainingFinal + /// One remaining partial-state batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("OrderedPartialAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ OrderedPartialAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ OrderedPartialAggregateState::DrainingFinal { .. } => { + self.handle_draining_final(state) + } + state @ OrderedPartialAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for OrderedPartialAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs index a4d19b0f7d18a..37b2473f4f92f 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/row_hash.rs @@ -1521,9 +1521,8 @@ mod tests { Ok(()) } - // TODO: migrate to PartialHashAggregateStream when it supports - // InputOrderMode::PartiallySorted; kept here for the legacy - // GroupedHashAggregateStream implementation. + // Migrated to OrderedPartialAggregateStream coverage in aggregates/mod.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_emit_early_with_partially_sorted() -> Result<()> { // Reproducer for #20445: EmitEarly with PartiallySorted panics in diff --git a/tmp/window_kernel_refactor.md b/tmp/window_kernel_refactor.md new file mode 100644 index 0000000000000..69d4f2f438331 --- /dev/null +++ b/tmp/window_kernel_refactor.md @@ -0,0 +1,213 @@ +The proposed refactor makes window function execution simpler and more extensible. I think it is a necessary step if we want to invest further in better vectorization or more parallel execution paradigms. + +The existing structure is not ideal: if we keep evolving the current shape, new optimization work will likely add more special cases and make the system harder to reason about. + +To sanity-check whether this refactor makes sense, we can use the potential optimizations mentioned in: + +- https://github.com/apache/datafusion/issues/23197 + +The examples include better parallelism and vectorization for fixed frames, parallel execution for prefix frames, and segment-tree-based parallelism. These optimizations are natural extensions of the ideal architecture introduced by this issue, but they are hard to add cleanly with the existing structure. + +This issue explains, in order: + +- How an ideal structure should look +- The issues in the existing implementation +- A possible implementation plan + +### Ideal Architecture + +The gist is that we should fully separate the logical and physical layers of window execution. + +- Logical layer: `WindowCall` purely describes what we want to calculate. It contains the expressions for arguments, partitioning, ordering, and frame bounds. +- Physical layer: `WindowKernel` purely provides the methods needed for execution. It represents the selected execution algorithm for a specific window call. + +This design brings below benefits: +- Simplicity: the control flow is one directional, `WindowCall` decides what window kernel to use, and window kernel purely provide methods for execution. +- Extensibility: adding new parallelism scheme/or improve vectorized fast path means adding one window kernel, no deep structural changes needed. + +#### Workflow + +```text +SQL / logical physical planning + -> WindowCall // pure description: function, args, partition/order/frame + -> WindowKernel selection // physical execution protocol chosen from shape + capabilities + -> WindowExec // execution routing: choose stream based on selected kernel + -> NaiveAccumulatorStream + -> SlidingAccumulatorStream + -> other specialized streams +``` + +In rough terms: + +```rust +/// pure description: function, args, partition/order/frame +struct WindowCall { + name: String, + field: FieldRef, + function: WindowFunctionKind, + args: Vec>, + filter: Option>, + partition_by: Vec>, + order_by: Vec, + frame: Arc, + options: WindowOptions, +} + +/// pure execution: provided methods needed for a specific path +enum WindowKernel { + /// Derived from existing Accumulator without `retract_batch` + /// A nested-loop algorithm will be used. + NaiveAccumulator(Box), + /// Derived from existing Accumulator with `retract_batch` + /// A sliding window algorithm will be. + SlidingAccumulator(Box), +} +``` + +DataFusion's existing `Accumulator` API already contains the primitives for two useful aggregate window algorithms: + +- `update_batch()` plus `evaluate()` can recompute a result for any frame. This supports a naive nested-loop fallback for all accumulators. +- `retract_batch()` plus `supports_retract_batch()` allow incremental sliding-window execution when rows leave the frame. + +If the accumulator does not support `retract_batch()`, a naive nested-loop evaluation can be used. If `retract_batch()` is supported and the window frame is a fixed sliding frame, a sliding-window algorithm can be used for optimization. + +Then the implication for newly added user-defined window function is, it should only support the naive method to make it work universally (for aggregate function in window cases, it requires only `update_batch()` for the above naive path), but it can optionally support more fast paths (`retract_batch` for sliding window, or even vectorized API in the future), then the optimizer/execution will route that into the fast path if the query expression shape allows. + +Here is a simple example to walk through the above workflow. + +#### Workload 1: Sliding Aggregate + +Example query: + +```sql +SELECT + avg(x) OVER ( + PARTITION BY k + ORDER BY ts + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ) AS avg_x +FROM t; +``` + +Planning: + +1. `WindowCall` holds the logical description: `avg(x)`, `PARTITION BY k`, `ORDER BY ts`, and `ROWS BETWEEN 2 PRECEDING AND CURRENT ROW`. +2. The planner sees that this is an aggregate window over a fixed moving frame. +3. The planner asks the aggregate accumulator whether it supports `retract_batch()`. `avg` does; +4. The planner chooses `SlidingAccumulatorWindowKernel`. +5. `WindowAggExec` routes execution to a dedicated `SlidingAccumulatorStream`, because the selected kernel has the sliding-window execution protocol. + +The kernel API can stay small because it only represents one physical protocol: + +```rust +trait SlidingAccumulatorWindowKernel { + fn evaluate_partition( + &mut self, + input: &PartitionWindowInput<'_>, + frame: &FrameIndex, + ) -> Result; +} + +struct PartitionWindowInput<'a> { + batch: &'a RecordBatch, + args: Vec, + filter: Option, +} +``` + +Very rough sliding-window algorithm sketch: + +```python +acc = create_avg_accumulator() +current_frame = range(0, 0) +output = [] + +for row_idx in partition_rows: + next_frame = frame_for(row_idx) + + # Rows that were in the previous frame but are not in the next frame. + leaving = current_frame.start .. next_frame.start + if leaving is not empty: + acc.retract_batch(values_for(leaving)) + + # Rows that are in the next frame but were not in the previous frame. + entering = current_frame.end .. next_frame.end + if entering is not empty: + acc.update_batch(values_for(entering)) + + output.append(acc.evaluate()) + current_frame = next_frame +``` + +This is the fast path: each input row is added and removed at most once, so the cost is linear in the partition size for row-based fixed frames. + +#### Workload 2: Naive Aggregate Fallback + +Example query: + +```sql +SELECT + my_udaf(x) OVER ( + PARTITION BY k + ORDER BY ts + ROWS BETWEEN t.n_gap PRECEDING AND CURRENT ROW + ) AS v +FROM t; +``` + +Assume `my_udaf` is a user-defined aggregate accumulator that supports `update_batch()` and `evaluate()`, but does not support `retract_batch()`. Also the window frame `t.n_gap` preceding can be arbitrary value, it's not supported by the sliding window algorithm. + +Planning: + +1. `WindowCall` holds the logical description: `my_udaf(x)`, `PARTITION BY k`, `ORDER BY ts`, and `ROWS BETWEEN 2 PRECEDING AND CURRENT ROW`. +2. The planner sees that this is an aggregate window (without `retract_batch()` capability), and also over a non-fixed moving frame. +3. The planner chooses `NaiveAccumulatorWindowKernel`. +4. `WindowAggExec` routes execution to a dedicated `NaiveAccumulatorStream`. + +The kernel API can again stay small: + +```rust +trait NaiveAccumulatorWindowKernel { + fn evaluate_partition( + &self, + input: &PartitionWindowInput<'_>, + frame: &FrameIndex, + ) -> Result; +} +``` + +Naive nested-loop algorithm sketch: + +```python +output = [] + +for row_idx in partition_rows: + frame = frame_for(row_idx) + + # This is slower, but it only needs update_batch() and evaluate(). + acc = create_my_udaf_accumulator() + acc.update_batch(values_for(frame)) + + output.append(acc.evaluate()) +``` + +### Issue with existing implementation +The major issue is that the existing abstraction layers leak into adjacent layers. I think the original design goal was: + +- `WindowExpr` is supposed to be the logical layer. +- `PartitionEvaluator` is supposed to be the physical layer. + +Over time, however, these responsibilities have become mixed. The decision-making flow has become bidirectional, and the implementation now relies on special cases to work around abstraction leaks. + +My guess is that these are mostly hacks accumulated over the years. I cannot find a strong reason to preserve this design. + +### Implementation Plan + +I plan to do some prototyping to work out a practical refactoring plan. The known goals are: + +- Remove all three `WindowExpr` implementations and use `WindowCall` as the pure logical layer. +- Use `WindowKernel` to replace the `PartitionEvaluator` + - `PartitionEvaluator` is now a large trait that uses 3+ flags to decide behavior. I think it is hard to use and extend; small, focused traits inside `WindowKernel` enum variants should be better. + - Provide an adapter like `WindowKernel::LegacyPartitionEvaluator` to make the refactor practical. +- Evolve `WindowAggExec` in this direction and avoid changing `BoundedWindowAggExec` + - See https://github.com/apache/datafusion/issues/23197#issuecomment-4806401319