From 5dfc10baf2e6f46c63bd90c65d7d5a7e917840b3 Mon Sep 17 00:00:00 2001 From: "xavier.lee" Date: Thu, 27 Aug 2026 14:22:48 -0400 Subject: [PATCH 1/3] test: cover unsorted contiguous groups in one partition --- .../physical-plan/src/aggregates/mod.rs | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 69a24ce17e852..adda048c45c15 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4539,6 +4539,83 @@ mod tests { Ok(()) } + #[tokio::test] + async fn unsorted_contiguous_groups_use_final_emission() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time_bin", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ])); + // Two sorted logical runs are emitted as batches in one DataFusion + // partition. Every distinct grouping tuple occupies one contiguous range, + // but tuple order resets at the batch boundary, so (key, time_bin) is not + // globally sorted. + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(Int64Array::from(vec![20, 20, 20, 20])), + Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(Int64Array::from(vec![0, 0, 0, 0])), + Arc::new(Int64Array::from(vec![50, 60, 70, 80])), + ], + )?, + ]; + let group_by = PhysicalGroupBy::new_single(vec![ + (col("key", &schema)?, "key".to_string()), + (col("time_bin", &schema)?, "time_bin".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 = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + assert_eq!(input.output_partitioning().partition_count(), 1); + + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggr_expr], + vec![None], + input, + schema, + )?; + + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + // This captures the behavior before #24438. When the source can declare + // `(key, time_bin)` group-contiguous, the corresponding case can use + // `EmissionType::Incremental`. + 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 stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+----------+------------+ +| key | time_bin | SUM(value) | ++-----+----------+------------+ +| 1 | 0 | 110 | +| 1 | 20 | 30 | +| 2 | 0 | 150 | +| 2 | 20 | 70 | ++-----+----------+------------+ +"); + + Ok(()) + } + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. #[tokio::test] async fn ordered_partial_aggregate_planning() -> Result<()> { From 332b8c98f0a8ac365f2aee0bfb3f9ef87d625962 Mon Sep 17 00:00:00 2001 From: "xavier.lee" Date: Wed, 26 Aug 2026 00:08:31 -0400 Subject: [PATCH 2/3] refactor: separate aggregate group completion from input ordering --- .../aggregate_hash_table/common_ordered.rs | 10 +-- .../ordered_final_table.rs | 8 +-- .../ordered_partial_table.rs | 2 +- .../ordered_single_table.rs | 2 +- .../src/aggregates/grouped_hash_stream.rs | 3 +- .../src/aggregates/hash_stream.rs | 4 +- .../physical-plan/src/aggregates/mod.rs | 38 +++++++++--- .../physical-plan/src/aggregates/order/mod.rs | 61 +++++++++++++++++-- .../src/aggregates/ordered_final_stream.rs | 46 +++++++------- .../src/aggregates/ordered_partial_stream.rs | 32 +++++----- .../src/aggregates/ordered_single_stream.rs | 46 +++++++------- .../src/aggregates/single_stream.rs | 4 +- 12 files changed, 173 insertions(+), 83 deletions(-) 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..00b6ffc214976 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,8 @@ impl OrderedAggregateTable { "OrderedAggregateTable requires config batch_size >= 1" ); - let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_ordering = + GroupOrdering::try_new_for_group_completion(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 cb5f8ade61c0d..ca58c51c79eb3 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -518,7 +518,8 @@ 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_group_completion(&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 adda048c45c15..664dd26eb4213 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; @@ -875,8 +876,13 @@ pub struct AggregateExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, required_input_ordering: Option, - /// Describes how the input is ordered relative to the group by columns + /// Describes how the input is ordered relative to the group by columns. input_order_mode: InputOrderMode, + /// Describes when the executor can determine that groups are complete. + /// + /// Initially derived from [`Self::input_order_mode`], but represented + /// separately so other input guarantees can establish group completion. + 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. @@ -901,6 +907,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), @@ -921,6 +928,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), @@ -1043,6 +1051,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())?; @@ -1073,6 +1083,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, + group_completion_mode, cache: Arc::new(cache), dynamic_filter: None, }; @@ -1281,7 +1292,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() @@ -1292,7 +1303,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() @@ -1303,7 +1314,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() } @@ -1316,7 +1327,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() } @@ -1326,7 +1337,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() } @@ -1336,7 +1347,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() } @@ -1346,7 +1357,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() } @@ -2354,6 +2365,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, @@ -4422,6 +4435,10 @@ mod tests { aggregate.input_order_mode(), InputOrderMode::PartiallySorted(_) )); + assert_eq!( + aggregate.group_completion_mode, + GroupCompletionMode::Partial(vec![0]) + ); let task_ctx = new_migrated_hash_ctx(2); let stream = aggregate.execute_typed(0, &task_ctx)?; @@ -4592,6 +4609,7 @@ mod tests { )?; assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::None); // This captures the behavior before #24438. When the source can declare // `(key, time_bin)` group-contiguous, the corresponding case can use // `EmissionType::Incremental`. @@ -4752,6 +4770,10 @@ mod tests { Arc::clone(&schema), )?; assert_eq!(final_aggregate.input_order_mode(), &InputOrderMode::Sorted); + assert_eq!( + final_aggregate.group_completion_mode, + GroupCompletionMode::Full + ); let task_ctx = new_migrated_hash_ctx(2); let stream = final_aggregate.execute_typed(0, &task_ctx)?; diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 259411b00b697..341114aedf773 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -28,6 +28,50 @@ 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. +/// +/// For example, when grouping by `key`, both inputs have fully contiguous +/// groups within the input partition: +/// +/// ```text +/// sorted: A A B B C C +/// not sorted: C C A A B B +/// ``` +/// +/// In both cases, once the key changes, the previous key will not appear again, +/// so its group is complete and can be emitted. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum GroupCompletionMode { + /// No group can be known complete before the input ends. + None, + /// Rows with the same values at these grouping-expression indices form one + /// contiguous range. When those values change, every group in the previous + /// range is complete and can be emitted. + /// + /// For example, with `GROUP BY (a, b)`, `Partial(vec![0])` means all rows + /// for each value of `a` are contiguous, while an `(a, b)` tuple may recur + /// within that range. + Partial(Vec), + /// Rows with the same complete grouping tuple form one contiguous range. + /// When the tuple changes, the previous group can be emitted. + 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 +84,24 @@ 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_group_completion(&GroupCompletionMode::from(mode)) + } + + /// Create a `GroupOrdering` for the specified group-completion mode. + pub(crate) fn try_new_for_group_completion( + 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..8e3d6009f1d4b 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! Final aggregate stream for ordered partial-state input. +//! Final aggregate stream for partial-state input with group-completion +//! guarantees. use std::ops::ControlFlow; use std::sync::Arc; @@ -35,16 +36,17 @@ 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 [`GroupCompletionMode::Partial`] and +/// [`GroupCompletionMode::Full`]. /// /// See comments at [`super::ordered_partial_stream::OrderedPartialAggregateStream`] for details. /// @@ -52,7 +54,7 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// /// This section is only for implementation notes, for background, see [`super::ordered_partial_stream::OrderedPartialAggregateStream`] /// -/// For partially sorted input, spilling works as follows: +/// For partial group completion, spilling works as follows: /// /// - Reserve the table footprint plus one `u32` sort index per buffered group. The /// extra index array is used in later sorting before spilling. @@ -70,8 +72,8 @@ pub(crate) struct OrderedFinalAggregateStream { state: Option, } -/// Spill configuration and accumulated runs for partially ordered final -/// aggregation. +/// Spill configuration and accumulated runs for final aggregation with partial +/// group completion. /// /// Each spill event drains all currently buffered groups, sorts their intermediate /// states by the full group key, and writes them to one spill file. All files are @@ -132,14 +134,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 +253,7 @@ impl OrderedFinalSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -269,10 +273,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 +284,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 +303,7 @@ impl OrderedFinalAggregateStream { context, partition, input, - input_order_mode, + group_completion_mode, baseline_metrics, metrics, Some(spill_metrics), @@ -319,7 +323,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 +333,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 +350,7 @@ impl OrderedFinalAggregateStream { context, partition, batch_size, - input_order_mode, + group_completion_mode, &input_schema, spill_metrics, )?)) @@ -354,12 +358,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..7626c192dba24 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Partial aggregate stream for ordered group input. +//! Partial aggregate stream for input with group-completion guarantees. use std::sync::Arc; @@ -29,13 +29,13 @@ 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`. +/// Partial aggregate stream for [`GroupCompletionMode::Partial`] and +/// [`GroupCompletionMode::Full`]. /// /// # Example /// @@ -59,20 +59,21 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Output: results for all groups (for example, `AVG(x)` calculated from the /// state) /// -/// # Order-based Optimization +/// # Group Completion 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. +/// After each input batch, the group-completion mode determines whether any +/// groups can be emitted eagerly to improve memory efficiency. For example, if +/// the input is ordered by `k` and the last group key seen is `k = 100`, all +/// groups with keys less than 100 are complete. /// /// # Memory Pressure and Spilling /// -/// ## Fully ordered case +/// ## Full group completion /// -/// If the input is ordered by every group key, for example: +/// Every complete grouping tuple is contiguous. Ordering by every group key is +/// one way to establish this mode, for example: /// /// - Input order: `a, b` /// - `GROUP BY`: `a, b` @@ -84,9 +85,10 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// If a memory reservation nevertheless fails, the stream returns the error /// directly, indicating an unexpected behavior. /// -/// ## Partially ordered case +/// ## Partial group completion /// -/// If the input is ordered by only a subset of the group keys, for example: +/// Rows are contiguous for a subset of the group keys. Ordering by that subset +/// is one way to establish this mode, for example: /// /// - Input order: `a` /// - `GROUP BY`: `a, b` @@ -126,7 +128,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..bc1140b186e4f 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Single-stage aggregate stream for ordered raw input. +//! Single-stage aggregate stream for raw input with group-completion guarantees. use std::ops::ControlFlow; use std::sync::Arc; @@ -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,8 @@ 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 [`GroupCompletionMode::Partial`] and +/// [`GroupCompletionMode::Full`]. /// /// # Example /// @@ -62,20 +63,21 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// Input: raw rows /// Output: final results for all groups (for example, `AVG(x)`) /// -/// # Order-based Optimization +/// # Group Completion 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. +/// After each input batch, the group-completion mode determines whether any +/// groups can be emitted eagerly to improve memory efficiency. For example, if +/// the input is ordered by `k` and the last group key seen is `k = 100`, all +/// groups with keys less than 100 are complete. /// /// # Memory Pressure and Spilling /// -/// ## Fully ordered case +/// ## Full group completion /// -/// If the input is ordered by every group key, for example: +/// Every complete grouping tuple is contiguous. Ordering by every group key is +/// one way to establish this mode, for example: /// /// - Input order: `a, b` /// - `GROUP BY`: `a, b` @@ -87,9 +89,10 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// If a memory reservation nevertheless fails, the stream returns the error /// directly, indicating an unexpected behavior. /// -/// ## Partially ordered case +/// ## Partial group completion /// -/// If the input is ordered by only a subset of the group keys, for example: +/// Rows are contiguous for a subset of the group keys. Ordering by that subset +/// is one way to establish this mode, for example: /// /// - Input order: `a` /// - `GROUP BY`: `a, b` @@ -109,8 +112,8 @@ pub(crate) struct OrderedSingleAggregateStream { state: Option, } -/// Spill configuration and accumulated runs for partially ordered single -/// aggregation. +/// Spill configuration and accumulated runs for single aggregation with partial +/// group completion. /// /// Each spill event drains all currently buffered groups, sorts their intermediate /// states by the full group key, and writes them to one spill file. All files are @@ -175,15 +178,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 +225,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 +313,7 @@ impl OrderedSingleSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -329,7 +333,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 +357,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 +365,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, From eec105932a6bfbf5ac55e72716354fd3c8ff1420 Mon Sep 17 00:00:00 2001 From: "xavier.lee" Date: Wed, 26 Aug 2026 09:12:03 -0400 Subject: [PATCH 3/3] feat: add narrow group-contiguous source property --- datafusion/datasource/src/source.rs | 100 ++++++++++++++++++ datafusion/ffi/src/session/mod.rs | 2 +- datafusion/physical-plan/src/buffer.rs | 18 ++++ datafusion/physical-plan/src/coop.rs | 20 +++- .../physical-plan/src/execution_plan.rs | 50 +++++++++ datafusion/physical-plan/src/projection.rs | 95 ++++++++++++++++- .../physical-plan/src/scalar_subquery.rs | 18 ++++ datafusion/physical-plan/src/test.rs | 49 +++++++++ 8 files changed, 349 insertions(+), 3 deletions(-) diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 741010c595197..58897e50ab5cd 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -165,6 +165,18 @@ pub trait DataSource: Any + Send + Sync + Debug { fn output_partitioning(&self) -> Partitioning; fn eq_properties(&self) -> EquivalenceProperties; + + /// Expressions whose complete tuple is contiguous within each output + /// partition. + /// + /// See + /// [`ExecutionPlanProperties::group_contiguous_exprs`](datafusion_physical_plan::ExecutionPlanProperties::group_contiguous_exprs) + /// for the full correctness contract. Expressions must refer to the schema + /// returned by [`Self::eq_properties`]. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } + fn scheduling_type(&self) -> SchedulingType { SchedulingType::NonCooperative } @@ -674,6 +686,7 @@ impl DataSourceExec { EmissionType::Incremental, Boundedness::Bounded, ) + .with_group_contiguous_exprs(data_source.group_contiguous_exprs().to_vec()) .with_scheduling_type(data_source.scheduling_type()) } @@ -705,3 +718,90 @@ where Self::new(Arc::new(source)) } } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_plan::EmptyRecordBatchStream; + + #[derive(Debug)] + struct GroupContiguousSource { + schema: SchemaRef, + group_contiguous_exprs: Vec>, + } + + impl DataSource for GroupContiguousSource { + fn open( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone( + &self.schema, + )))) + } + + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "GroupContiguousSource") + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(1) + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new(Arc::clone(&self.schema)) + } + + fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + + fn with_fetch(&self, _limit: Option) -> Option> { + None + } + + fn fetch(&self) -> Option { + None + } + + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExprs, + ) -> Result>> { + Ok(None) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + } + + #[test] + fn data_source_exec_exposes_group_contiguous_exprs() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let key = Arc::new(Column::new("key", 0)) as Arc; + let exec = DataSourceExec::new(Arc::new(GroupContiguousSource { + schema, + group_contiguous_exprs: vec![Arc::clone(&key)], + })); + + assert_eq!(exec.properties().group_contiguous_exprs().len(), 1); + assert!(exec.properties().group_contiguous_exprs()[0].eq(&key)); + Ok(()) + } +} diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index f215a6ba5a568..d428741a34d00 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -985,7 +985,7 @@ mod tests { let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?; assert_eq!( format!("{physical_plan:?}"), - "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None } }" + "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None, group_contiguous_exprs: [] } }" ); assert_eq!( diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index e1452b75bc4c3..5e720056d18d7 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -524,15 +524,33 @@ impl Stream for MemoryBufferedStream { #[cfg(test)] mod tests { use super::*; + use crate::test::TestMemoryExec; + use arrow_schema::{DataType, Field, Schema}; use datafusion_common::{DataFusionError, assert_contains}; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; + use datafusion_physical_expr::expressions::col; use std::error::Error; use std::fmt::Debug; use std::time::Duration; use tokio::time::timeout; + #[test] + fn buffer_exec_preserves_group_contiguity() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?; + let buffer = BufferExec::new(Arc::new(source), 1024); + + assert_eq!(buffer.properties().group_contiguous_exprs().len(), 1); + assert!( + buffer.properties().group_contiguous_exprs()[0].eq(&col("key", &schema)?) + ); + Ok(()) + } + #[tokio::test] async fn buffers_only_some_messages() -> Result<(), Box> { let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 17166e287e6dc..f6287b98b3e34 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -477,7 +477,9 @@ pub fn make_cooperative(stream: SendableRecordBatchStream) -> SendableRecordBatc mod tests { use super::*; - use arrow_schema::SchemaRef; + use crate::test::TestMemoryExec; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use datafusion_physical_expr::expressions::col; use futures::stream; @@ -497,6 +499,22 @@ mod tests { Box::pin(RecordBatchStreamAdapter::new(schema, s)) } + #[test] + fn cooperative_exec_preserves_group_contiguity() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?; + let cooperative = CooperativeExec::new(Arc::new(source)); + + assert_eq!(cooperative.properties().group_contiguous_exprs().len(), 1); + assert!( + cooperative.properties().group_contiguous_exprs()[0] + .eq(&col("key", &schema)?) + ); + Ok(()) + } + #[tokio::test] async fn yield_less_than_threshold() -> Result<()> { let count = TASK_BUDGET - 10; diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index c0fa4c6bede41..d15c2566322df 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -1245,6 +1245,28 @@ pub trait ExecutionPlanProperties { /// /// [`FilterExec`]: crate::filter::FilterExec fn equivalence_properties(&self) -> &EquivalenceProperties; + + /// Expressions whose complete tuple is contiguous within each output + /// partition. + /// + /// For every distinct tuple formed by these expressions, all matching rows + /// must occur in at most one contiguous range. Once a different tuple appears, + /// the previous tuple must not appear again. Tuple equality follows `GROUP BY` + /// semantics, and tuples need not be sorted. + /// + /// The guarantee applies to the complete tuple, not to individual expressions + /// or subsets. Expressions must refer to this plan's output schema. This + /// property implies neither output ordering nor distribution. The default is + /// no guarantee. + /// + /// # Correctness + /// + /// An invalid declaration can cause a streaming aggregate to emit a group + /// before observing all of its rows, producing incorrect results. Use + /// [`PlanProperties::with_group_contiguous_exprs`] to set this property. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } } impl ExecutionPlanProperties for Arc { @@ -1267,6 +1289,10 @@ impl ExecutionPlanProperties for Arc { fn equivalence_properties(&self) -> &EquivalenceProperties { self.properties().equivalence_properties() } + + fn group_contiguous_exprs(&self) -> &[Arc] { + self.properties().group_contiguous_exprs() + } } impl ExecutionPlanProperties for &dyn ExecutionPlan { @@ -1289,6 +1315,10 @@ impl ExecutionPlanProperties for &dyn ExecutionPlan { fn equivalence_properties(&self) -> &EquivalenceProperties { self.properties().equivalence_properties() } + + fn group_contiguous_exprs(&self) -> &[Arc] { + self.properties().group_contiguous_exprs() + } } /// Represents whether a stream of data **generated** by an operator is bounded (finite) @@ -1501,6 +1531,8 @@ pub struct PlanProperties { pub scheduling_type: SchedulingType, /// See [ExecutionPlanProperties::output_ordering] output_ordering: Option, + /// See [`ExecutionPlanProperties::group_contiguous_exprs`] + group_contiguous_exprs: Vec>, } impl PlanProperties { @@ -1521,6 +1553,7 @@ impl PlanProperties { evaluation_type: EvaluationType::Lazy, scheduling_type: SchedulingType::NonCooperative, output_ordering, + group_contiguous_exprs: vec![], } } @@ -1572,6 +1605,18 @@ impl PlanProperties { self } + /// Overwrite the group-contiguous composite key. + /// + /// See [`ExecutionPlanProperties::group_contiguous_exprs`] for the correctness + /// contract. + pub fn with_group_contiguous_exprs( + mut self, + group_contiguous_exprs: Vec>, + ) -> Self { + self.group_contiguous_exprs = group_contiguous_exprs; + self + } + /// Set constraints having mut reference. pub fn set_constraints(&mut self, constraints: Constraints) { self.eq_properties.set_constraints(constraints); @@ -1595,6 +1640,11 @@ impl PlanProperties { self.output_ordering.as_ref() } + /// Components of the group-contiguous composite key, if any. + pub fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + /// Get schema of the node. pub(crate) fn schema(&self) -> &SchemaRef { self.eq_properties.schema() diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 2a24eb60e6fbc..108b5a79724f2 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -208,12 +208,15 @@ impl ProjectionExec { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = projector.projection().projection_mapping(&input.schema())?; + let group_contiguous_exprs = + Self::project_group_contiguous_exprs(&input, &projection_mapping); let cache = Self::compute_properties( &input, &projection_mapping, Arc::clone(projector.output_schema()), reuse_from, - )?; + )? + .with_group_contiguous_exprs(group_contiguous_exprs); Ok(Self { projector, input, @@ -222,6 +225,22 @@ impl ProjectionExec { }) } + /// Projects the complete group-contiguous tuple, dropping it when any + /// component cannot be mapped to the output schema. + fn project_group_contiguous_exprs( + input: &Arc, + projection_mapping: &ProjectionMapping, + ) -> Vec> { + input + .equivalence_properties() + .project_expressions( + input.group_contiguous_exprs().iter(), + projection_mapping, + ) + .collect::>>() + .unwrap_or_default() + } + /// The projection expressions stored as tuples of (expression, output column name) pub fn expr(&self) -> &[ProjectionExpr] { self.projector.projection().as_ref() @@ -1515,6 +1534,7 @@ mod tests { use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; + use crate::test::TestMemoryExec; use crate::test::exec::StatisticsExec; use arrow::datatypes::{DataType, Field, Schema}; @@ -1526,6 +1546,79 @@ mod tests { BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, }; + #[test] + fn group_contiguous_projection_is_all_or_nothing() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time", DataType::Int32, false), + ])); + let time_bin = binary( + col("time", &schema)?, + Operator::Divide, + lit(ScalarValue::Int32(Some(10))), + &schema, + )?; + let plain_source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?; + let source = plain_source.clone().try_with_group_contiguous_exprs(vec![ + col("key", &schema)?, + Arc::clone(&time_bin), + ])?; + + let projection = ProjectionExec::try_new( + [ + ProjectionExpr::new(col("key", &schema)?, "key"), + ProjectionExpr::new(time_bin, "time_bin"), + ], + Arc::new(source.clone()), + )?; + let projected_schema = projection.schema(); + let expected = [ + col("key", &projected_schema)?, + col("time_bin", &projected_schema)?, + ]; + assert_eq!( + projection.properties().group_contiguous_exprs().len(), + expected.len() + ); + assert!( + projection + .properties() + .group_contiguous_exprs() + .iter() + .zip(expected) + .all(|(actual, expected)| actual.eq(&expected)) + ); + + // A strict subset is not sufficient: contiguity of `(key, time_bin)` + // does not imply that `key` alone is contiguous. + let partial_projection = ProjectionExec::try_new( + [ProjectionExpr::new(col("key", &schema)?, "key")], + Arc::new(source.clone()), + )?; + assert!( + partial_projection + .properties() + .group_contiguous_exprs() + .is_empty() + ); + + // Row-preserving operators do not inherit the assertion unless they + // opt in explicitly. + let filter = FilterExec::try_new(lit(true), Arc::new(source.clone()))?; + assert!(filter.properties().group_contiguous_exprs().is_empty()); + + // Group contiguity participates in child-property identity, so + // replacing the source with one lacking the assertion recomputes the + // projection properties. + assert!(!Arc::ptr_eq(source.properties(), plain_source.properties())); + let projection: Arc = Arc::new(projection); + let replaced = + replace_children_if_necessary(projection, vec![Arc::new(plain_source)])?; + assert!(replaced.group_contiguous_exprs().is_empty()); + + Ok(()) + } + #[test] fn test_try_new_with_schema_metadata_only_replaces_metadata() -> Result<()> { let input_schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index cac8f925039ae..674930350e4f4 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -442,8 +442,26 @@ mod tests { use arrow::array::{Int32Array, Int64Array}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; + use datafusion_physical_expr::expressions::col; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + #[test] + fn scalar_subquery_exec_preserves_group_contiguity() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?; + let exec = ScalarSubqueryExec::new( + Arc::new(source), + vec![], + ScalarSubqueryResults::new(0), + ); + + assert_eq!(exec.properties().group_contiguous_exprs().len(), 1); + assert!(exec.properties().group_contiguous_exprs()[0].eq(&col("key", &schema)?)); + Ok(()) + } + enum ExpectedSubqueryResult { Value(ScalarValue), Error(&'static str), diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index b38a46d160755..6b304c0b50379 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -73,6 +73,8 @@ pub struct TestMemoryExec { projection: Option>, /// Sort information: one or more equivalent orderings sort_information: Vec, + /// Composite key whose values are contiguous within each output stream. + group_contiguous_exprs: Vec>, /// if partition sizes should be displayed show_sizes: bool, /// The maximum number of records to read from this plan. If `None`, @@ -226,6 +228,7 @@ impl TestMemoryExec { EmissionType::Incremental, Boundedness::Bounded, ) + .with_group_contiguous_exprs(self.group_contiguous_exprs.clone()) } fn output_partitioning(&self) -> Partitioning { @@ -268,6 +271,7 @@ impl TestMemoryExec { projected_schema, projection, sort_information: vec![], + group_contiguous_exprs: vec![], show_sizes: true, fetch: None, }) @@ -363,6 +367,51 @@ impl TestMemoryExec { Ok(self) } + /// Attach a composite key whose values occur in one contiguous range in + /// each output stream. See + /// [`ExecutionPlanProperties::group_contiguous_exprs`](crate::ExecutionPlanProperties::group_contiguous_exprs) + /// for the correctness contract. + pub fn try_with_group_contiguous_exprs( + mut self, + mut group_contiguous_exprs: Vec>, + ) -> Result { + // All expressions must refer to the original schema. + let fields = self.schema.fields(); + let ambiguous_column = group_contiguous_exprs + .iter() + .flat_map(collect_columns) + .find(|col| { + fields + .get(col.index()) + .map(|field| field.name() != col.name()) + .unwrap_or(true) + }); + assert_or_internal_err!( + ambiguous_column.is_none(), + "Column {:?} is not found in the original schema of the TestMemoryExec", + ambiguous_column.as_ref().unwrap() + ); + + if let Some(projection) = &self.projection { + let base_schema = self.original_schema(); + let proj_exprs = projection.iter().map(|idx| { + let name = base_schema.field(*idx).name(); + (Arc::new(Column::new(name, *idx)) as _, name.to_string()) + }); + let projection_mapping = + ProjectionMapping::try_new(proj_exprs, &base_schema)?; + let base_eqp = EquivalenceProperties::new(base_schema); + group_contiguous_exprs = base_eqp + .project_expressions(group_contiguous_exprs.iter(), &projection_mapping) + .collect::>>() + .unwrap_or_default(); + } + + self.group_contiguous_exprs = group_contiguous_exprs; + self.cache = Arc::new(self.compute_properties()); + Ok(self) + } + /// Arc clone of ref to original schema pub fn original_schema(&self) -> SchemaRef { Arc::clone(&self.schema)