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 4ced967a0977b..87a44ef3ebdc0 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 @@ -28,14 +28,13 @@ 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::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, }; use crate::aggregates::grouped_hash_stream::create_group_accumulator; -use crate::aggregates::order::GroupOrdering; +use crate::aggregates::order::{GroupCompletionMode, GroupOrdering}; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, @@ -113,7 +112,7 @@ impl OrderedAggregateTableMetrics { /// `OrderedAggrMode` selects the aggregate semantics. For example, /// `OrderedAggregateTable::::new(...)` consumes raw rows /// and emits partial states, while -/// `OrderedAggregateTable::::new_with_input_order(...)` +/// `OrderedAggregateTable::::new_with_group_completion(...)` /// consumes partial states and emits final values. /// /// Shared methods live on `impl`; single/partial/final behavior lives on @@ -184,7 +183,7 @@ impl OrderedAggregateTable { output_schema: SchemaRef, state_schema: SchemaRef, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, aggregate_mode: &AggregateMode, filters: Vec>>, metrics: OrderedAggregateTableMetrics, @@ -194,7 +193,7 @@ impl OrderedAggregateTable { "OrderedAggregateTable requires config batch_size >= 1" ); - let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_ordering = GroupOrdering::try_new_for_mode(group_completion_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( 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 d0d0c99bb5bd8..35dbd8a7566e4 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 @@ -25,8 +25,8 @@ 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::order::GroupCompletionMode; use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase}; use super::common::HashAggregateAccumulator; @@ -42,12 +42,12 @@ use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics} /// /// See comments at [`OrderedAggregateTable`] for details. impl OrderedAggregateTable { - pub(in crate::aggregates) fn new_with_input_order( + pub(in crate::aggregates) fn new_with_group_completion( agg: &AggregateExec, input_schema: &SchemaRef, output_schema: SchemaRef, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, metrics: OrderedAggregateTableMetrics, ) -> Result { Self::new_for_mode( @@ -56,7 +56,7 @@ impl OrderedAggregateTable { output_schema, Arc::clone(input_schema), batch_size, - input_order_mode, + group_completion_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], metrics, 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 6ed93e59f3296..39564c66863a9 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 @@ -68,7 +68,7 @@ impl OrderedAggregateTable { output_schema, state_schema, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &AggregateMode::Partial, agg.filter_expr.iter().cloned().collect(), metrics, 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 index ce1ce647b46fe..88a7bea6ed7a0 100644 --- 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 @@ -59,7 +59,7 @@ impl OrderedAggregateTable { output_schema, state_schema, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &agg.mode, agg.filter_expr.iter().cloned().collect(), metrics, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index c0253093c8a7b..a1004f4b8420e 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -504,7 +504,7 @@ impl GroupedHashAggregateStream { .collect::>() .join(", "); let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})"); - let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?; + let group_ordering = GroupOrdering::try_new_for_mode(&agg.group_completion_mode)?; let oom_mode = match (agg.mode, &group_ordering) { // In partial aggregation mode, always prefer to emit incomplete results early. (AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly, diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2df5960188a2b..8fce12ad6d6b0 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -45,6 +45,7 @@ use super::aggregate_hash_table::{ AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker, PartialSkipMarker, }; +use super::order::GroupCompletionMode; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::skip_partial::SkipAggregationProbe; use crate::metrics::{ @@ -326,6 +327,7 @@ impl FinalSpillContext { let mut final_agg = agg.clone(); final_agg.input_order_mode = InputOrderMode::Sorted; + final_agg.group_completion_mode = GroupCompletionMode::Full; Ok(Self { final_agg, @@ -414,7 +416,7 @@ impl FinalSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 59d382302abe6..64ad057f7128a 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -203,6 +203,7 @@ use datafusion_physical_expr_common::sort_expr::{ use datafusion_expr::utils::AggregateOrderSensitivity; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use itertools::Itertools; +use order::GroupCompletionMode; use topk::hash_table::is_supported_hash_key_type; use topk::heap::is_supported_heap_type; @@ -866,6 +867,8 @@ pub struct AggregateExec { required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns input_order_mode: InputOrderMode, + /// Describes how the executor can determine that groups are complete. + group_completion_mode: GroupCompletionMode, cache: Arc, /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]), /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node. @@ -890,6 +893,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + group_completion_mode: self.group_completion_mode.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -910,6 +914,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + group_completion_mode: self.group_completion_mode.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -1032,6 +1037,8 @@ impl AggregateExec { input_order_mode = InputOrderMode::Linear; } + let group_completion_mode = GroupCompletionMode::from(&input_order_mode); + // construct a map from the input expression to the output expression of the Aggregation group by let group_expr_mapping = ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?; @@ -1062,6 +1069,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, + group_completion_mode, cache: Arc::new(cache), dynamic_filter: None, }; @@ -1270,7 +1278,7 @@ impl AggregateExec { fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::Partial - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() @@ -1281,7 +1289,7 @@ impl AggregateExec { _context: &TaskContext, ) -> bool { self.mode == AggregateMode::Partial - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() @@ -1292,7 +1300,7 @@ impl AggregateExec { self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1305,7 +1313,7 @@ impl AggregateExec { self.mode == AggregateMode::PartialReduce && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1315,7 +1323,7 @@ impl AggregateExec { self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned ) && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1325,7 +1333,7 @@ impl AggregateExec { self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned ) && self.limit_options.is_none() - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1335,7 +1343,7 @@ impl AggregateExec { self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -2343,6 +2351,8 @@ impl ExecutionPlan for AggregateExec { required_input_ordering: _, // Derived at construction from the input ordering and `group_by`. input_order_mode: _, + // Derived at construction from `input_order_mode`. + group_completion_mode: _, // Derived at construction by `Self::compute_properties`. cache: _, dynamic_filter, @@ -3221,6 +3231,7 @@ mod tests { use crate::filter::FilterExecBuilder; use crate::metrics::MetricValue; use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::stream::RecordBatchStreamAdapter; use crate::test::TestMemoryExec; use crate::test::assert_is_pending; use crate::test::exec::{ @@ -3379,6 +3390,102 @@ mod tests { ) } + /// A one-partition source that yields one batch and then remains pending. + /// It distinguishes aggregate output produced before EOF from output + /// produced after EOF. + #[derive(Debug)] + struct OneBatchThenPendingExec { + batch: RecordBatch, + cache: Arc, + } + + impl OneBatchThenPendingExec { + fn new(batch: RecordBatch) -> Self { + let cache = PlanProperties::new( + EquivalenceProperties::new(batch.schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ); + Self { + batch, + cache: Arc::new(cache), + } + } + } + + impl DisplayAs for OneBatchThenPendingExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "OneBatchThenPendingExec") + } + } + + impl ExecutionPlan for OneBatchThenPendingExec { + fn name(&self) -> &'static str { + "OneBatchThenPendingExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn replace_children( + self: Arc, + _children: Vec>, + _options: ReplaceChildrenOptions, + ) -> Result> { + internal_err!("Children cannot be replaced in {self:?}") + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let stream = futures::stream::iter([Ok(self.batch.clone())]) + .chain(futures::stream::pending()); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.batch.schema(), + stream, + ))) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))) + } + } + fn new_spill_ctx(batch_size: usize, max_memory: usize) -> Arc { let session_config = SessionConfig::new().with_batch_size(batch_size); let runtime = RuntimeEnvBuilder::new() @@ -4528,6 +4635,51 @@ mod tests { Ok(()) } + #[tokio::test] + async fn linear_aggregate_waits_for_input_end() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2, 3, 3])), + Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + ], + )?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + ); + let input: Arc = Arc::new(OneBatchThenPendingExec::new(batch)); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggr_expr], + vec![None], + input, + schema, + )?; + + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::None); + assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + + let task_ctx = new_migrated_hash_ctx(1024); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let mut stream: SendableRecordBatchStream = stream.into(); + let mut next = stream.next().boxed(); + assert_is_pending(&mut next); + + Ok(()) + } + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. #[tokio::test] async fn ordered_partial_aggregate_planning() -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 259411b00b697..0dae640e77b5c 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -28,6 +28,33 @@ use crate::InputOrderMode; pub use full::GroupOrderingFull; pub use partial::GroupOrderingPartial; +/// Describes how an aggregate can determine that groups are complete. +/// +/// This is distinct from [`InputOrderMode`], which describes the ordering of +/// the input relative to the grouping expressions. Input ordering is one way +/// to establish a group-completion mode, but the execution machinery only +/// needs to know when it can safely emit completed groups. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum GroupCompletionMode { + /// Groups cannot be completed before the input ends. + None, + /// Groups sharing the values at these grouping-expression indices form a + /// contiguous range. + Partial(Vec), + /// Every complete grouping tuple forms a contiguous range. + Full, +} + +impl From<&InputOrderMode> for GroupCompletionMode { + fn from(value: &InputOrderMode) -> Self { + match value { + InputOrderMode::Linear => Self::None, + InputOrderMode::PartiallySorted(indices) => Self::Partial(indices.clone()), + InputOrderMode::Sorted => Self::Full, + } + } +} + /// Ordering information for each group in the hash table #[derive(Debug)] pub enum GroupOrdering { @@ -40,15 +67,22 @@ pub enum GroupOrdering { } impl GroupOrdering { - /// Create a `GroupOrdering` for the specified ordering + /// Create a `GroupOrdering` for the specified input order mode. pub fn try_new(mode: &InputOrderMode) -> Result { + Self::try_new_for_mode(&GroupCompletionMode::from(mode)) + } + + /// Create a `GroupOrdering` for the specified group-completion mode. + pub(crate) fn try_new_for_mode(mode: &GroupCompletionMode) -> Result { match mode { - InputOrderMode::Linear => Ok(GroupOrdering::None), - InputOrderMode::PartiallySorted(order_indices) => { + GroupCompletionMode::None => Ok(GroupOrdering::None), + GroupCompletionMode::Partial(order_indices) => { GroupOrderingPartial::try_new(order_indices.clone()) .map(GroupOrdering::Partial) } - InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())), + GroupCompletionMode::Full => { + Ok(GroupOrdering::Full(GroupOrderingFull::new())) + } } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 2c26b74da7748..31776f2550198 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -35,16 +35,16 @@ use super::AggregateExec; use super::aggregate_hash_table::{ FinalMarker, OrderedAggregateTable, OrderedAggregateTableMetrics, }; +use super::order::GroupCompletionMode; 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}; +use crate::{RecordBatchStream, SendableRecordBatchStream}; -/// Final aggregate stream for `InputOrderMode::Sorted` and -/// `InputOrderMode::PartiallySorted`. +/// Final aggregate stream for input whose completed group ranges can be identified. /// /// See comments at [`super::ordered_partial_stream::OrderedPartialAggregateStream`] for details. /// @@ -132,14 +132,16 @@ impl OrderedFinalSpillContext { context: &Arc, partition: usize, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, spill_schema: &SchemaRef, spill_metrics: SpillMetrics, ) -> Result { let group_schema = agg.group_by.group_schema(spill_schema)?; let output_ordering = agg.cache.output_ordering(); - let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { - return internal_err!("Ordered final spill requires partially ordered input"); + let GroupCompletionMode::Partial(order_indices) = group_completion_mode else { + return internal_err!( + "Ordered final spill requires partial group completion" + ); }; let spill_indices = order_indices.iter().copied().chain( (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), @@ -249,7 +251,7 @@ impl OrderedFinalSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -269,10 +271,10 @@ impl OrderedFinalAggregateStream { agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned )); - debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(agg.group_completion_mode, GroupCompletionMode::None); let input = agg.input.execute(partition, Arc::clone(context))?; - Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) + Self::new_with_input(agg, context, partition, input, &agg.group_completion_mode) } pub(in crate::aggregates) fn new_with_input( @@ -280,7 +282,7 @@ impl OrderedFinalAggregateStream { context: &Arc, partition: usize, input: SendableRecordBatchStream, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, ) -> Result { let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let metrics = OrderedAggregateTableMetrics::new(agg, partition); @@ -299,7 +301,7 @@ impl OrderedFinalAggregateStream { context, partition, input, - input_order_mode, + group_completion_mode, baseline_metrics, metrics, Some(spill_metrics), @@ -319,7 +321,7 @@ impl OrderedFinalAggregateStream { context: &Arc, partition: usize, input: SendableRecordBatchStream, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, baseline_metrics: BaselineMetrics, metrics: OrderedAggregateTableMetrics, spill_metrics: Option, @@ -329,13 +331,13 @@ impl OrderedFinalAggregateStream { agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned )); - debug_assert_ne!(*input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(*group_completion_mode, GroupCompletionMode::None); let schema = Arc::clone(&agg.schema); let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); - let can_spill = matches!(input_order_mode, InputOrderMode::PartiallySorted(_)) + let can_spill = matches!(group_completion_mode, GroupCompletionMode::Partial(_)) && context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { let Some(spill_metrics) = spill_metrics else { @@ -346,7 +348,7 @@ impl OrderedFinalAggregateStream { context, partition, batch_size, - input_order_mode, + group_completion_mode, &input_schema, spill_metrics, )?)) @@ -354,12 +356,12 @@ impl OrderedFinalAggregateStream { None }; - let table = OrderedAggregateTable::::new_with_input_order( + let table = OrderedAggregateTable::::new_with_group_completion( agg, &input_schema, Arc::clone(&schema), batch_size, - input_order_mode, + group_completion_mode, metrics, )?; Ok(Self { diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 9e93a111a6466..61586ccd1d964 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -29,10 +29,10 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; -use crate::aggregates::order::GroupOrdering; +use crate::aggregates::order::{GroupCompletionMode, GroupOrdering}; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; -use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; +use crate::{SendableRecordBatchStream, metrics}; /// Partial aggregate stream for `InputOrderMode::Sorted` and /// `InputOrderMode::PartiallySorted`. @@ -126,7 +126,7 @@ impl OrderedPartialAggregateStream { partition: usize, ) -> Result { debug_assert_eq!(agg.mode, AggregateMode::Partial); - debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(agg.group_completion_mode, GroupCompletionMode::None); let schema = Arc::clone(&agg.schema); let input = agg.input.execute(partition, Arc::clone(context))?; diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index da00b42e5c3ed..e21a27102ba1c 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -34,6 +34,7 @@ use futures::stream::{Stream, StreamExt}; use super::aggregate_hash_table::{ OrderedAggregateTable, OrderedAggregateTableMetrics, SingleMarker, }; +use super::order::GroupCompletionMode; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -44,8 +45,7 @@ use crate::spill::spill_manager::SpillManager; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; -/// Single aggregate stream for `InputOrderMode::Sorted` and -/// `InputOrderMode::PartiallySorted`. +/// Single aggregate stream for input whose completed group ranges can be identified. /// /// # Example /// @@ -175,15 +175,15 @@ impl OrderedSingleSpillContext { context: &Arc, partition: usize, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, 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 { + let GroupCompletionMode::Partial(order_indices) = group_completion_mode else { return internal_err!( - "Ordered single spill requires partially ordered input" + "Ordered single spill requires partial group completion" ); }; let spill_indices = order_indices.iter().copied().chain( @@ -222,6 +222,7 @@ impl OrderedSingleSpillContext { }; final_agg.group_by = Arc::new(agg.group_by.as_final()); final_agg.input_order_mode = InputOrderMode::Sorted; + final_agg.group_completion_mode = GroupCompletionMode::Full; Ok(Self { final_agg, @@ -309,7 +310,7 @@ impl OrderedSingleSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -329,7 +330,7 @@ impl OrderedSingleAggregateStream { agg.mode, AggregateMode::Single | AggregateMode::SinglePartitioned )); - debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + debug_assert_ne!(agg.group_completion_mode, GroupCompletionMode::None); let schema = Arc::clone(&agg.schema); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -353,7 +354,7 @@ impl OrderedSingleAggregateStream { )?; let can_spill = - matches!(agg.input_order_mode, InputOrderMode::PartiallySorted(_)) + matches!(agg.group_completion_mode, GroupCompletionMode::Partial(_)) && context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { Some(Box::new(OrderedSingleSpillContext::new( @@ -361,7 +362,7 @@ impl OrderedSingleAggregateStream { context, partition, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &state_schema, spill_metrics, )?)) diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 3e306d72a7e82..6e3667bbe273a 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -39,6 +39,7 @@ use futures::stream::{Stream, StreamExt}; use super::aggregate_hash_table::{ AggregateHashTable, OrderedAggregateTableMetrics, SingleMarker, }; +use super::order::GroupCompletionMode; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -212,6 +213,7 @@ impl SingleSpillContext { }; final_agg.group_by = Arc::new(agg.group_by.as_final()); final_agg.input_order_mode = InputOrderMode::Sorted; + final_agg.group_completion_mode = GroupCompletionMode::Full; Ok(Self { final_agg, @@ -300,7 +302,7 @@ impl SingleSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None,