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..9bfbe08bd46de 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. @@ -876,6 +879,36 @@ pub struct AggregateExec { dynamic_filter: Option>, } +/// Returns whether `group_contiguous_exprs` and `groupby_exprs` describe the +/// same complete tuple, allowing for equivalent expressions and permutation. +fn group_contiguous_exprs_match( + eq_properties: &EquivalenceProperties, + groupby_exprs: &[Arc], + group_contiguous_exprs: &[Arc], +) -> bool { + if group_contiguous_exprs.is_empty() + || group_contiguous_exprs.len() != groupby_exprs.len() + { + return false; + } + + let mut matched_group_exprs = vec![false; groupby_exprs.len()]; + group_contiguous_exprs.iter().all(|contiguous_expr| { + let Some((idx, _)) = + groupby_exprs.iter().enumerate().find(|(idx, group_expr)| { + !matched_group_exprs[*idx] + && eq_properties + .eq_group() + .exprs_equal(contiguous_expr, group_expr) + }) + else { + return false; + }; + matched_group_exprs[idx] = true; + true + }) +} + impl AggregateExec { /// Function used in `OptimizeAggregateOrder` optimizer rule, /// where we need parts of the new value, others cloned from the old one @@ -890,6 +923,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 +944,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 +1067,17 @@ impl AggregateExec { input_order_mode = InputOrderMode::Linear; } + let mut group_completion_mode = GroupCompletionMode::from(&input_order_mode); + if group_by.is_single() + && group_contiguous_exprs_match( + input_eq_properties, + &groupby_exprs, + input.group_contiguous_exprs(), + ) + { + group_completion_mode = GroupCompletionMode::Full; + } + // 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())?; @@ -1049,6 +1095,11 @@ impl AggregateExec { aggr_expr.as_ref(), )? }; + let cache = if group_completion_mode != GroupCompletionMode::None { + cache.with_emission_type(input.pipeline_behavior()) + } else { + cache + }; let mut exec = AggregateExec { mode, @@ -1062,6 +1113,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, + group_completion_mode, cache: Arc::new(cache), dynamic_filter: None, }; @@ -1071,6 +1123,21 @@ impl AggregateExec { Ok(exec) } + fn try_with_new_input(&self, input: Arc) -> Result { + let mut me = AggregateExec::try_new_with_schema( + self.mode, + Arc::clone(&self.group_by), + self.aggr_expr.to_vec(), + Arc::clone(&self.filter_expr), + input, + Arc::clone(&self.input_schema), + Arc::clone(&self.schema), + )?; + me.limit_options = self.limit_options; + me.dynamic_filter.clone_from(&self.dynamic_filter); + Ok(me) + } + /// Aggregation mode (full, partial) pub fn mode(&self) -> &AggregateMode { &self.mode @@ -1270,7 +1337,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 +1348,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 +1359,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 +1372,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 +1382,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 +1392,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 +1402,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() } @@ -1963,6 +2030,15 @@ impl DisplayAs for AggregateExec { if self.input_order_mode != InputOrderMode::Linear { write!(f, ", ordering_mode={:?}", self.input_order_mode)?; } + if self.group_completion_mode + != GroupCompletionMode::from(&self.input_order_mode) + { + write!( + f, + ", group_completion_mode={:?}", + self.group_completion_mode + )?; + } } DisplayFormatType::TreeRender => { let format_expr_with_alias = @@ -2100,6 +2176,20 @@ impl ExecutionPlan for AggregateExec { vec![self.input_order_mode != InputOrderMode::Linear] } + fn benefits_from_input_partitioning(&self) -> Vec { + // A repartition would discard a group-contiguous assertion. Avoid + // introducing one when that assertion is what enables early emission. + if self.group_completion_mode != GroupCompletionMode::from(&self.input_order_mode) + { + return vec![false]; + } + + self.input_distribution_requirements() + .per_child_distributions() + .map(|dist| !matches!(dist, Distribution::SinglePartition)) + .collect() + } + fn children(&self) -> Vec<&Arc> { vec![&self.input] } @@ -2111,25 +2201,29 @@ impl ExecutionPlan for AggregateExec { ) -> 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::Recompute => { - let mut me = AggregateExec::try_new_with_schema( - self.mode, - Arc::clone(&self.group_by), - self.aggr_expr.to_vec(), - Arc::clone(&self.filter_expr), - Arc::clone(&children[0]), - Arc::clone(&self.input_schema), - Arc::clone(&self.schema), - )?; - me.limit_options = self.limit_options; - me.dynamic_filter.clone_from(&self.dynamic_filter); - Ok(Arc::new(me)) + ChildrenPropertiesMode::Keep => { + let input = children.swap_remove(0); + let old_exprs = self.input.group_contiguous_exprs(); + let new_exprs = input.group_contiguous_exprs(); + let same_group_contiguous_exprs = old_exprs.len() == new_exprs.len() + && old_exprs + .iter() + .zip(new_exprs) + .all(|(old, new)| old.eq(new)); + if same_group_contiguous_exprs { + Ok(Arc::new(Self { + input, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } else { + self.try_with_new_input(input) + .map(|plan| Arc::new(plan) as _) + } } + ChildrenPropertiesMode::Recompute => self + .try_with_new_input(children.swap_remove(0)) + .map(|plan| Arc::new(plan) as _), } } @@ -2343,6 +2437,8 @@ impl ExecutionPlan for AggregateExec { required_input_ordering: _, // Derived at construction from the input ordering and `group_by`. input_order_mode: _, + // Derived from input ordering or a group-contiguous input tuple. + group_completion_mode: _, // Derived at construction by `Self::compute_properties`. cache: _, dynamic_filter, @@ -3221,6 +3317,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::{ @@ -3229,10 +3326,11 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, StructArray, UInt32Array, UInt64Array, + Int64Array, StructArray, TimestampSecondArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; - use arrow::datatypes::Int32Type; + use arrow::datatypes::{Int32Type, TimeUnit}; + use datafusion_common::config::ConfigOptions; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{DataFusionError, internal_err}; use datafusion_execution::config::SessionConfig; @@ -3243,6 +3341,7 @@ mod tests { Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature, Volatility, }; + use datafusion_functions::datetime::date_bin; use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; use datafusion_functions_aggregate::array_agg::array_agg_udaf; use datafusion_functions_aggregate::average::avg_udaf; @@ -3251,11 +3350,12 @@ mod tests { use datafusion_functions_aggregate::median::median_udaf; use datafusion_functions_aggregate::min_max::min_udaf; use datafusion_functions_aggregate::sum::sum_udaf; - use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::{Literal, NotExpr}; + use datafusion_physical_expr::{Partitioning, ScalarFunctionExpr}; + use crate::execution_plan::replace_children_if_necessary; use crate::projection::ProjectionExec; use crate::repartition::RepartitionExec; use datafusion_physical_expr::projection::ProjectionExpr; @@ -3379,6 +3479,111 @@ 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, + group_contiguous_exprs: Vec>, + cache: Arc, + } + + impl OneBatchThenPendingExec { + fn new( + batch: RecordBatch, + group_contiguous_exprs: Vec>, + ) -> Self { + let cache = PlanProperties::new( + EquivalenceProperties::new(batch.schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ); + Self { + batch, + group_contiguous_exprs, + 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 group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + + 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 +4733,315 @@ mod tests { Ok(()) } + #[test] + fn group_contiguous_requires_complete_group_tuple() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + assert!(group_contiguous_exprs_match( + &eq_properties, + &[Arc::clone(&b), Arc::clone(&a)], + &[Arc::clone(&a), Arc::clone(&b)], + )); + assert!(!group_contiguous_exprs_match( + &eq_properties, + &[Arc::clone(&a), Arc::clone(&b)], + &[Arc::clone(&a)], + )); + eq_properties.add_equal_conditions(Arc::clone(&a), Arc::clone(&b))?; + assert!(group_contiguous_exprs_match( + &eq_properties, + &[Arc::clone(&a)], + &[Arc::clone(&b)], + )); + + let plain_source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?; + let source = plain_source + .clone() + .try_with_group_contiguous_exprs(vec![Arc::clone(&a), Arc::clone(&b)])?; + let group_by = PhysicalGroupBy::new_single(vec![ + (Arc::clone(&b), "b".to_string()), + (Arc::clone(&a), "a".to_string()), + ]); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![], + vec![], + Arc::new(source.clone()), + Arc::clone(&schema), + )?; + + assert_eq!(aggregate.input_order_mode, InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + assert_eq!(aggregate.benefits_from_input_partitioning(), vec![false]); + assert!(aggregate.group_contiguous_exprs().is_empty()); + assert!(aggregate.cache().output_ordering().is_none()); + + let streaming_partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + Arc::new(source.clone()), + Arc::clone(&schema), + )?; + let ordinary_partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + Arc::new(plain_source.clone()), + Arc::clone(&schema), + )?; + assert_eq!( + streaming_partial.benefits_from_input_partitioning(), + vec![false] + ); + assert_eq!( + ordinary_partial.benefits_from_input_partitioning(), + vec![true] + ); + + let grouping_sets = PhysicalGroupBy::new( + vec![ + (col("a", &schema)?, "a".to_string()), + (col("b", &schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Int32(None)), "a".to_string()), + (lit(ScalarValue::Int32(None)), "b".to_string()), + ], + vec![vec![false, false], vec![false, true]], + true, + ); + let grouping_sets_aggregate = AggregateExec::try_new( + AggregateMode::Single, + grouping_sets, + vec![], + vec![], + Arc::new(source.clone()), + Arc::clone(&schema), + )?; + assert_eq!( + grouping_sets_aggregate.group_completion_mode, + GroupCompletionMode::None + ); + + // The assertion is outside PlanProperties. Replacing the child with a + // plan that shares the same property cache must still recompute the + // aggregate's completion mode. + assert!(Arc::ptr_eq(source.properties(), plain_source.properties())); + let aggregate: Arc = Arc::new(aggregate); + let replaced = replace_children_if_necessary( + aggregate, + vec![Arc::new(plain_source.clone())], + )?; + let replaced = replaced.downcast_ref::().unwrap(); + assert_eq!(replaced.group_completion_mode, GroupCompletionMode::None); + assert_eq!(replaced.cache().emission_type, EmissionType::Final); + assert_eq!(replaced.benefits_from_input_partitioning(), vec![false]); + + // A strict subset is deliberately unsupported in the initial design. + let subset_source = plain_source + .clone() + .try_with_group_contiguous_exprs(vec![a])?; + let subset_aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![], + vec![], + Arc::new(subset_source), + Arc::clone(&schema), + )?; + assert_eq!( + subset_aggregate.group_completion_mode, + GroupCompletionMode::None + ); + + Ok(()) + } + + #[tokio::test] + async fn group_contiguous_aggregate_emits_before_input_ends() -> 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 key = col("key", &schema)?; + let group_by = + PhysicalGroupBy::new_single(vec![(Arc::clone(&key), "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.clone(), + vec![Arc::clone(&key)], + )); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + input, + Arc::clone(&schema), + )?; + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + + let task_ctx = new_migrated_hash_ctx(1024); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + let mut stream: SendableRecordBatchStream = stream.into(); + let first = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("group-contiguous aggregate should emit before input ends") + .transpose()? + .expect("one group should be complete"); + assert_snapshot!(batches_to_sort_string(&[first]), @r" ++-----+------------+ +| key | SUM(value) | ++-----+------------+ +| 2 | 30 | ++-----+------------+ +"); + + let input: Arc = + Arc::new(OneBatchThenPendingExec::new(batch, vec![])); + 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 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(()) + } + + #[tokio::test] + async fn group_contiguous_date_bin_projects_to_aggregate() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time", DataType::Timestamp(TimeUnit::Second, None), false), + Field::new("value", DataType::Int64, false), + ])); + let batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1, 1])), + Arc::new(TimestampSecondArray::from(vec![20, 21, 30, 31])), + Arc::new(Int64Array::from(vec![1, 2, 3, 4])), + ], + )?, + // Tuple order resets between logical source runs, while no tuple + // appears in both runs. + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1, 1])), + Arc::new(TimestampSecondArray::from(vec![0, 1, 10, 11])), + Arc::new(Int64Array::from(vec![5, 6, 7, 8])), + ], + )?, + ]; + // The source certifies the derived tuple after validating its logical + // run boundaries. ProjectionExec only maps that complete assertion. + let time_bin = Arc::new(ScalarFunctionExpr::try_new( + date_bin(), + vec![ + lit(ScalarValue::new_interval_dt(0, 10_000)), + col("time", &schema)?, + ], + &schema, + Arc::new(ConfigOptions::default()), + )?) as Arc; + let source = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? + .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"), + ProjectionExpr::new(col("value", &schema)?, "value"), + ], + Arc::new(source), + )?; + + let projected_schema = projection.schema(); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (col("key", &projected_schema)?, "key".to_string()), + (col("time_bin", &projected_schema)?, "time_bin".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new( + sum_udaf(), + vec![col("value", &projected_schema)?], + ) + .schema(Arc::clone(&projected_schema)) + .alias("SUM(value)") + .build()?, + )], + vec![None], + Arc::new(projection), + projected_schema, + )?; + assert_eq!(aggregate.input_order_mode, InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + + let stream = aggregate.execute_typed(0, &new_migrated_hash_ctx(2))?; + assert!(matches!(stream, StreamType::OrderedSingleAggregate(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+---------------------+------------+ +| key | time_bin | SUM(value) | ++-----+---------------------+------------+ +| 1 | 1970-01-01T00:00:00 | 11 | +| 1 | 1970-01-01T00:00:10 | 15 | +| 1 | 1970-01-01T00:00:20 | 3 | +| 1 | 1970-01-01T00:00:30 | 7 | ++-----+---------------------+------------+ +"); + + 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)