diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 741010c595197..bd1c2c48d3667 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -165,6 +165,19 @@ 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 [`ExecutionPlan::group_contiguous_exprs`] for the full correctness + /// contract. Expressions must refer to the schema returned by + /// [`Self::eq_properties`]. A source rewrite such as [`Self::repartitioned`] + /// or [`Self::try_swapping_with_projection`] must preserve, remap, or clear + /// this assertion as appropriate for the rewritten output streams. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } + fn scheduling_type(&self) -> SchedulingType { SchedulingType::NonCooperative } @@ -384,7 +397,17 @@ impl DisplayAs for DataSourceExec { } DisplayFormatType::TreeRender => {} } - self.data_source.fmt_as(t, f) + self.data_source.fmt_as(t, f)?; + if matches!(t, DisplayFormatType::Default | DisplayFormatType::Verbose) + && !self.group_contiguous_exprs().is_empty() + { + write!( + f, + ", group_contiguous=[{}]", + self.group_contiguous_exprs().iter().join(", ") + )?; + } + Ok(()) } } @@ -397,6 +420,10 @@ impl ExecutionPlan for DataSourceExec { &self.cache } + fn group_contiguous_exprs(&self) -> &[Arc] { + self.data_source.group_contiguous_exprs() + } + fn children(&self) -> Vec<&Arc> { Vec::new() } @@ -705,3 +732,94 @@ 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::{DefaultDisplay, 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.group_contiguous_exprs().len(), 1); + assert!(exec.group_contiguous_exprs()[0].eq(&key)); + assert_eq!( + DefaultDisplay(exec).to_string(), + "DataSourceExec: GroupContiguousSource, group_contiguous=[key@0]" + ); + Ok(()) + } +} 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, diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 17166e287e6dc..945e2e8853898 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -261,6 +261,10 @@ impl ExecutionPlan for CooperativeExec { &self.properties } + fn group_contiguous_exprs(&self) -> &[Arc] { + self.input.group_contiguous_exprs() + } + fn maintains_input_order(&self) -> Vec { vec![true; self.children().len()] } @@ -477,7 +481,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 +503,19 @@ 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.group_contiguous_exprs().len(), 1); + assert!(cooperative.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..7cd5f11060d65 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -159,6 +159,32 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// trait, which is implemented for all `ExecutionPlan`s. fn properties(&self) -> &Arc; + /// Expressions whose complete tuple is contiguous within each output + /// partition. + /// + /// For every distinct tuple of expression values, all rows with that tuple + /// must occur in at most one contiguous range in each output stream. Once a + /// stream produces a different tuple, the previous tuple must never occur + /// again. Tuple equality follows `GROUP BY` semantics, and tuple values do + /// not need to be sorted. + /// + /// This property does not imply any particular output ordering or + /// distribution. It is also deliberately fail-closed: the default is no + /// guarantee, and operators must explicitly preserve it. Currently, + /// [`ProjectionExec`] preserves the complete tuple when every expression + /// can be projected, and [`crate::coop::CooperativeExec`] preserves it + /// because it does not change rows. Other operators use the default empty + /// value. + /// + /// # Correctness + /// + /// This is a correctness contract. An invalid declaration can cause a + /// streaming aggregate to emit a group before all of its rows have been + /// observed, producing incorrect results. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } + /// Returns an error if this individual node does not conform to its invariants. /// These invariants are typically only checked in debug mode. /// diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 2a24eb60e6fbc..086681b3fee58 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -83,6 +83,8 @@ pub struct ProjectionExec { metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. cache: Arc, + /// Complete group-contiguous tuple projected from the input, if any. + group_contiguous_exprs: Vec>, } impl ProjectionExec { @@ -208,6 +210,8 @@ 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, @@ -219,9 +223,26 @@ impl ProjectionExec { input, metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), + group_contiguous_exprs, }) } + /// 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() @@ -345,6 +366,10 @@ impl ExecutionPlan for ProjectionExec { &self.cache } + fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + fn maintains_input_order(&self) -> Vec { // Tell optimizer this operator doesn't reorder its input vec![true] @@ -386,11 +411,21 @@ impl ExecutionPlan for ProjectionExec { ) -> Result> { validate_child_count!(self, children); match options.children_properties { - ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })), + ChildrenPropertiesMode::Keep => { + let input = children.swap_remove(0); + let projection_mapping = self + .projector + .projection() + .projection_mapping(&input.schema())?; + let group_contiguous_exprs = + Self::project_group_contiguous_exprs(&input, &projection_mapping); + Ok(Arc::new(Self { + input, + metrics: ExecutionPlanMetricsSet::new(), + group_contiguous_exprs, + ..Self::clone(&*self) + })) + } ChildrenPropertiesMode::Recompute => { // `Keep` above requires the child's properties to be unchanged // outright. A rule that introduces a sort below this projection @@ -631,6 +666,8 @@ impl ExecutionPlan for ProjectionExec { metrics: _, // Derived plan properties, recomputed on decode. cache: _, + // Derived from the input assertion and projection expressions. + group_contiguous_exprs: _, } = self; let projection_exprs = projector.projection().as_ref(); let input = ctx.encode_child(input)?; @@ -1515,6 +1552,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 +1564,70 @@ 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.group_contiguous_exprs().len(), expected.len()); + assert!( + projection + .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.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.group_contiguous_exprs().is_empty()); + + // The assertion deliberately lives outside PlanProperties. Exercise + // the child-replacement fast path with identical cached properties and + // verify that ProjectionExec still recomputes it. + 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/test.rs b/datafusion/physical-plan/src/test.rs index b38a46d160755..e037276762b0c 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`, @@ -106,16 +108,27 @@ impl DisplayAs for TestMemoryExec { let limit = self .fetch .map_or(String::new(), |limit| format!(", fetch={limit}")); + let group_contiguous = if self.group_contiguous_exprs.is_empty() { + String::new() + } else { + let exprs = self + .group_contiguous_exprs + .iter() + .map(ToString::to_string) + .collect::>() + .join(", "); + format!(", group_contiguous=[{exprs}]") + }; if self.show_sizes { write!( f, - "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}{group_contiguous}", partition_sizes.len(), ) } else { write!( f, - "partitions={}{limit}{output_ordering}{constraints}", + "partitions={}{limit}{output_ordering}{constraints}{group_contiguous}", partition_sizes.len(), ) } @@ -137,6 +150,10 @@ impl ExecutionPlan for TestMemoryExec { &self.cache } + fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + fn children(&self) -> Vec<&Arc> { Vec::new() } @@ -268,6 +285,7 @@ impl TestMemoryExec { projected_schema, projection, sort_information: vec![], + group_contiguous_exprs: vec![], show_sizes: true, fetch: None, }) @@ -363,6 +381,49 @@ impl TestMemoryExec { Ok(self) } + /// Attach a composite key whose values occur in one contiguous range in + /// each output stream. See [`ExecutionPlan::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; + Ok(self) + } + /// Arc clone of ref to original schema pub fn original_schema(&self) -> SchemaRef { Arc::clone(&self.schema)