diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 741010c595197..58897e50ab5cd 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -165,6 +165,18 @@ pub trait DataSource: Any + Send + Sync + Debug { fn output_partitioning(&self) -> Partitioning; fn eq_properties(&self) -> EquivalenceProperties; + + /// Expressions whose complete tuple is contiguous within each output + /// partition. + /// + /// See + /// [`ExecutionPlanProperties::group_contiguous_exprs`](datafusion_physical_plan::ExecutionPlanProperties::group_contiguous_exprs) + /// for the full correctness contract. Expressions must refer to the schema + /// returned by [`Self::eq_properties`]. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } + fn scheduling_type(&self) -> SchedulingType { SchedulingType::NonCooperative } @@ -674,6 +686,7 @@ impl DataSourceExec { EmissionType::Incremental, Boundedness::Bounded, ) + .with_group_contiguous_exprs(data_source.group_contiguous_exprs().to_vec()) .with_scheduling_type(data_source.scheduling_type()) } @@ -705,3 +718,90 @@ where Self::new(Arc::new(source)) } } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_plan::EmptyRecordBatchStream; + + #[derive(Debug)] + struct GroupContiguousSource { + schema: SchemaRef, + group_contiguous_exprs: Vec>, + } + + impl DataSource for GroupContiguousSource { + fn open( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone( + &self.schema, + )))) + } + + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "GroupContiguousSource") + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(1) + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new(Arc::clone(&self.schema)) + } + + fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + + fn with_fetch(&self, _limit: Option) -> Option> { + None + } + + fn fetch(&self) -> Option { + None + } + + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExprs, + ) -> Result>> { + Ok(None) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + } + + #[test] + fn data_source_exec_exposes_group_contiguous_exprs() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let key = Arc::new(Column::new("key", 0)) as Arc; + let exec = DataSourceExec::new(Arc::new(GroupContiguousSource { + schema, + group_contiguous_exprs: vec![Arc::clone(&key)], + })); + + assert_eq!(exec.properties().group_contiguous_exprs().len(), 1); + assert!(exec.properties().group_contiguous_exprs()[0].eq(&key)); + Ok(()) + } +} diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index f215a6ba5a568..d428741a34d00 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -985,7 +985,7 @@ mod tests { let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?; assert_eq!( format!("{physical_plan:?}"), - "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None } }" + "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None, group_contiguous_exprs: [] } }" ); assert_eq!( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 4ced967a0977b..00b6ffc214976 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -28,14 +28,13 @@ use datafusion_common::assert_or_internal_err; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::EmitTo; -use crate::InputOrderMode; use crate::PhysicalExpr; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, }; use crate::aggregates::grouped_hash_stream::create_group_accumulator; -use crate::aggregates::order::GroupOrdering; +use crate::aggregates::order::{GroupCompletionMode, GroupOrdering}; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, @@ -113,7 +112,7 @@ impl OrderedAggregateTableMetrics { /// `OrderedAggrMode` selects the aggregate semantics. For example, /// `OrderedAggregateTable::::new(...)` consumes raw rows /// and emits partial states, while -/// `OrderedAggregateTable::::new_with_input_order(...)` +/// `OrderedAggregateTable::::new_with_group_completion(...)` /// consumes partial states and emits final values. /// /// Shared methods live on `impl`; single/partial/final behavior lives on @@ -184,7 +183,7 @@ impl OrderedAggregateTable { output_schema: SchemaRef, state_schema: SchemaRef, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, aggregate_mode: &AggregateMode, filters: Vec>>, metrics: OrderedAggregateTableMetrics, @@ -194,7 +193,8 @@ impl OrderedAggregateTable { "OrderedAggregateTable requires config batch_size >= 1" ); - let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_ordering = + GroupOrdering::try_new_for_group_completion(group_completion_mode)?; let group_schema = agg.group_by.group_schema(input_schema)?; let group_values = new_group_values(group_schema, &group_ordering)?; let aggregate_arguments = aggregate_expressions( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index d0d0c99bb5bd8..35dbd8a7566e4 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -25,8 +25,8 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; -use crate::InputOrderMode; use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::order::GroupCompletionMode; use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase}; use super::common::HashAggregateAccumulator; @@ -42,12 +42,12 @@ use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics} /// /// See comments at [`OrderedAggregateTable`] for details. impl OrderedAggregateTable { - pub(in crate::aggregates) fn new_with_input_order( + pub(in crate::aggregates) fn new_with_group_completion( agg: &AggregateExec, input_schema: &SchemaRef, output_schema: SchemaRef, batch_size: usize, - input_order_mode: &InputOrderMode, + group_completion_mode: &GroupCompletionMode, metrics: OrderedAggregateTableMetrics, ) -> Result { Self::new_for_mode( @@ -56,7 +56,7 @@ impl OrderedAggregateTable { output_schema, Arc::clone(input_schema), batch_size, - input_order_mode, + group_completion_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], metrics, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index 6ed93e59f3296..80051e7d5f455 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 @@ -15,12 +15,14 @@ // specific language governing permissions and limitations // under the License. -//! Aggregate table for partial aggregation when input is ordered by group keys. +//! Aggregate table for partial aggregation when input groups can be completed +//! before end of input. //! //! See the [`super::common_ordered`] comments for the high-level ideas. //! -//! This operator handles input that is ordered by group keys: -//! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b` +//! This operator handles both complete-key contiguity and input ordered by a +//! subset of the group keys: +//! - Fully contiguous: every distinct `(a, b)` tuple occurs in one input range //! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a` //! //! When a group key combination is exhausted, this table eagerly flushes the @@ -68,7 +70,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, @@ -90,7 +92,7 @@ impl OrderedAggregateTable { } /// Emits the next batch of partial state rows for groups proven complete by - /// the input ordering. + /// the input's group-completion guarantee. /// /// For example, when the query is `GROUP BY a` and the input is ordered by /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs index ce1ce647b46fe..88a7bea6ed7a0 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs @@ -59,7 +59,7 @@ impl OrderedAggregateTable { output_schema, state_schema, batch_size, - &agg.input_order_mode, + &agg.group_completion_mode, &agg.mode, agg.filter_expr.iter().cloned().collect(), metrics, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index cb5f8ade61c0d..958511f7127fa 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -518,7 +518,8 @@ impl GroupedHashAggregateStream { .collect::>() .join(", "); let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})"); - let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?; + let group_ordering = + GroupOrdering::try_new_for_group_completion(&agg.group_completion_mode)?; let oom_mode = match (agg.mode, &group_ordering) { // In partial aggregation mode, always prefer to emit incomplete results early. (AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly, @@ -530,10 +531,10 @@ impl GroupedHashAggregateStream { { OutOfMemoryMode::Spill } - // For `GroupOrdering::Full`, the incoming stream is already sorted. This ensures the - // number of incomplete groups can be kept small at all times. If we still hit - // an out-of-memory condition, spilling to disk would not be beneficial since the same - // situation is likely to reoccur when reading back the spilled data. + // For `GroupOrdering::Full`, every complete group tuple is contiguous. This keeps the + // number of incomplete groups small at all times. If we still hit an out-of-memory + // condition, spilling to disk would not be beneficial since the same situation is + // likely to reoccur when reading back the spilled data. // Therefore, we fall back to simply reporting the error immediately. // This mode will also be used if the `DiskManager` is not configured to allow spilling // to disk. 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 69a24ce17e852..ab72af12816ed 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; @@ -877,6 +878,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. @@ -887,6 +890,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 @@ -901,6 +934,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + group_completion_mode: self.group_completion_mode.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -921,6 +955,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + group_completion_mode: self.group_completion_mode.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -1043,6 +1078,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())?; @@ -1060,6 +1106,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, @@ -1073,6 +1124,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, + group_completion_mode, cache: Arc::new(cache), dynamic_filter: None, }; @@ -1281,7 +1333,7 @@ impl AggregateExec { fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::Partial - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() @@ -1292,7 +1344,7 @@ impl AggregateExec { _context: &TaskContext, ) -> bool { self.mode == AggregateMode::Partial - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() @@ -1303,7 +1355,7 @@ impl AggregateExec { self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned ) && self.limit_options_supported_by_hash_stream() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1316,7 +1368,7 @@ impl AggregateExec { self.mode == AggregateMode::PartialReduce && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1326,7 +1378,7 @@ impl AggregateExec { self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned ) && self.limit_options.is_none() - && self.input_order_mode == InputOrderMode::Linear + && self.group_completion_mode == GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1336,7 +1388,7 @@ impl AggregateExec { self.mode, AggregateMode::Single | AggregateMode::SinglePartitioned ) && self.limit_options.is_none() - && self.input_order_mode != InputOrderMode::Linear + && self.group_completion_mode != GroupCompletionMode::None && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } @@ -1346,7 +1398,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() } @@ -1974,6 +2026,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 = @@ -2111,6 +2172,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] } @@ -2133,7 +2208,7 @@ impl ExecutionPlan for AggregateExec { Arc::clone(&self.group_by), self.aggr_expr.to_vec(), Arc::clone(&self.filter_expr), - Arc::clone(&children[0]), + children.swap_remove(0), Arc::clone(&self.input_schema), Arc::clone(&self.schema), )?; @@ -2354,6 +2429,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, @@ -3240,10 +3317,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; @@ -3254,6 +3332,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; @@ -3262,11 +3341,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; @@ -4539,6 +4619,314 @@ mod tests { Ok(()) } + #[test] + fn group_contiguous_exprs_require_complete_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)], + )); + + Ok(()) + } + + #[tokio::test] + async fn unsorted_contiguous_groups_use_incremental_emission() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time_bin", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ])); + // Two sorted logical runs are emitted as batches in one DataFusion + // partition. Every distinct grouping tuple occupies one contiguous range, + // but tuple order resets at the batch boundary, so (key, time_bin) is not + // globally sorted. + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(Int64Array::from(vec![20, 20, 20, 20])), + Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(Int64Array::from(vec![0, 0, 0, 0])), + Arc::new(Int64Array::from(vec![50, 60, 70, 80])), + ], + )?, + ]; + let key = col("key", &schema)?; + let time_bin = col("time_bin", &schema)?; + let group_by = PhysicalGroupBy::new_single(vec![ + (Arc::clone(&key), "key".to_string()), + (Arc::clone(&time_bin), "time_bin".to_string()), + ]); + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + ); + let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![key, time_bin])?; + let input: Arc = Arc::new(input); + assert_eq!(input.output_partitioning().partition_count(), 1); + + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggr_expr], + vec![None], + input, + schema, + )?; + + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + assert!(aggregate.properties().group_contiguous_exprs().is_empty()); + assert!(aggregate.cache().output_ordering().is_none()); + + let stream = aggregate.execute_typed(0, &new_migrated_hash_ctx(1024))?; + 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 | 0 | 110 | +| 1 | 20 | 30 | +| 2 | 0 | 150 | +| 2 | 20 | 70 | ++-----+----------+------------+ +"); + + Ok(()) + } + + #[test] + fn group_contiguous_aggregate_does_not_benefit_from_repartitioning() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let key = col("key", &schema)?; + let group_by = + PhysicalGroupBy::new_single(vec![(Arc::clone(&key), "key".to_string())]); + let plain_source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?; + let contiguous_source = plain_source + .clone() + .try_with_group_contiguous_exprs(vec![key])?; + + let contiguous = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + Arc::new(contiguous_source), + Arc::clone(&schema), + )?; + let ordinary = AggregateExec::try_new( + AggregateMode::Partial, + group_by, + vec![], + vec![], + Arc::new(plain_source), + schema, + )?; + + assert_eq!(contiguous.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(ordinary.group_completion_mode, GroupCompletionMode::None); + assert_eq!(contiguous.benefits_from_input_partitioning(), vec![false]); + assert_eq!(ordinary.benefits_from_input_partitioning(), vec![true]); + + Ok(()) + } + + #[test] + fn group_contiguous_exprs_do_not_apply_to_grouping_sets() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![ + col("a", &schema)?, + col("b", &schema)?, + ])?; + let group_by = 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 aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![], + vec![], + Arc::new(source), + schema, + )?; + + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::None); + assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + + Ok(()) + } + + #[test] + fn group_contiguous_aggregate_recomputes_after_child_replacement() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let key = col("key", &schema)?; + let plain_source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?; + let contiguous_source = plain_source + .clone() + .try_with_group_contiguous_exprs(vec![Arc::clone(&key)])?; + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(key, "key".to_string())]), + vec![], + vec![], + Arc::new(contiguous_source), + Arc::clone(&schema), + )?; + assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + + let aggregate: Arc = Arc::new(aggregate); + let replaced = + replace_children_if_necessary(aggregate, vec![Arc::new(plain_source)])?; + let replaced = replaced.downcast_ref::().unwrap(); + + assert_eq!(replaced.group_completion_mode, GroupCompletionMode::None); + assert_eq!(replaced.cache().emission_type, EmissionType::Final); + + 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/full.rs b/datafusion/physical-plan/src/aggregates/order/full.rs index ca818d6a2d598..26b6f3881ff76 100644 --- a/datafusion/physical-plan/src/aggregates/order/full.rs +++ b/datafusion/physical-plan/src/aggregates/order/full.rs @@ -18,12 +18,12 @@ use datafusion_expr::EmitTo; use std::mem::size_of; -/// Tracks grouping state when the data is ordered entirely by its -/// group keys +/// Tracks grouping state when every complete group key occurs in one contiguous +/// input range. /// -/// When the group values are sorted, as soon as we see group `n+1` we -/// know we will never see any rows for group `n` again and thus they -/// can be emitted. +/// A full sort by the group keys provides this guarantee, but the group values +/// do not need to be sorted. As soon as the next group begins, the preceding +/// group can be emitted because it cannot occur again. /// /// For example, given `SUM(amt) GROUP BY id` if the input is sorted /// by `id` as soon as a new `id` value is seen all previous values diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 259411b00b697..cbd1eb5e6aa49 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,24 @@ pub enum GroupOrdering { } impl GroupOrdering { - /// Create a `GroupOrdering` for the specified ordering + /// Create a `GroupOrdering` for the specified input order mode. pub fn try_new(mode: &InputOrderMode) -> Result { + Self::try_new_for_group_completion(&GroupCompletionMode::from(mode)) + } + + /// Create a `GroupOrdering` for the specified group-completion mode. + pub(crate) fn try_new_for_group_completion( + mode: &GroupCompletionMode, + ) -> Result { match mode { - InputOrderMode::Linear => Ok(GroupOrdering::None), - InputOrderMode::PartiallySorted(order_indices) => { + GroupCompletionMode::None => Ok(GroupOrdering::None), + GroupCompletionMode::Partial(order_indices) => { GroupOrderingPartial::try_new(order_indices.clone()) .map(GroupOrdering::Partial) } - InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())), + GroupCompletionMode::Full => { + Ok(GroupOrdering::Full(GroupOrderingFull::new())) + } } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 2c26b74da7748..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..506e71cb70f10 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`. @@ -59,27 +59,22 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Output: results for all groups (for example, `AVG(x)` calculated from the /// state) /// -/// # Order-based Optimization +/// # Group-completion optimization /// /// For the aggregation work, the hash aggregation implementation is reused. /// /// After each input batch, check whether any groups can be emitted eagerly to -/// improve memory efficiency. For example, if the last group key seen is -/// `k = 100`, it is safe to emit all groups with keys less than 100 because the -/// input is ordered. +/// improve memory efficiency. Input ordering can establish this guarantee, as +/// can complete-key contiguity. /// /// # Memory Pressure and Spilling /// -/// ## Fully ordered case +/// ## Full group-completion case /// -/// If the input is ordered by every group key, for example: -/// -/// - Input order: `a, b` -/// - `GROUP BY`: `a, b` -/// -/// Completed groups can be emitted as soon as the next group is observed. Thus, -/// only the current group remains active after completed groups are emitted, and -/// memory usage does not grow with the total number of groups. +/// If every complete group tuple is contiguous, completed groups can be emitted +/// as soon as the next group is observed. Thus, only the current group remains +/// active after completed groups are emitted, and memory usage does not grow +/// with the total number of groups. /// /// If a memory reservation nevertheless fails, the stream returns the error /// directly, indicating an unexpected behavior. @@ -126,7 +121,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..42ae64e322e82 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 /// @@ -62,27 +62,22 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// Input: raw rows /// Output: final results for all groups (for example, `AVG(x)`) /// -/// # Order-based Optimization +/// # Group-completion optimization /// /// For the aggregation work, the hash aggregation implementation is reused. /// /// After each input batch, check whether any groups can be emitted eagerly to -/// improve memory efficiency. For example, if the last group key seen is -/// `k = 100`, it is safe to emit all groups with keys less than 100 because the -/// input is ordered. +/// improve memory efficiency. Input ordering can establish this guarantee, as +/// can complete-key contiguity. /// /// # Memory Pressure and Spilling /// -/// ## Fully ordered case +/// ## Full group-completion case /// -/// If the input is ordered by every group key, for example: -/// -/// - Input order: `a, b` -/// - `GROUP BY`: `a, b` -/// -/// Completed groups can be emitted as soon as the next group is observed. Thus, -/// only the current group remains active after completed groups are emitted, and -/// memory usage does not grow with the total number of groups. +/// If every complete group tuple is contiguous, completed groups can be emitted +/// as soon as the next group is observed. Thus, only the current group remains +/// active after completed groups are emitted, and memory usage does not grow +/// with the total number of groups. /// /// If a memory reservation nevertheless fails, the stream returns the error /// directly, indicating an unexpected behavior. @@ -175,15 +170,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 +217,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 +305,7 @@ impl OrderedSingleSpillContext { &context, partition, merged, - &InputOrderMode::Sorted, + &GroupCompletionMode::Full, baseline_metrics.clone(), metrics, None, @@ -329,7 +325,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 +349,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 +357,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/buffer.rs b/datafusion/physical-plan/src/buffer.rs index e1452b75bc4c3..5e720056d18d7 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -524,15 +524,33 @@ impl Stream for MemoryBufferedStream { #[cfg(test)] mod tests { use super::*; + use crate::test::TestMemoryExec; + use arrow_schema::{DataType, Field, Schema}; use datafusion_common::{DataFusionError, assert_contains}; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; + use datafusion_physical_expr::expressions::col; use std::error::Error; use std::fmt::Debug; use std::time::Duration; use tokio::time::timeout; + #[test] + fn buffer_exec_preserves_group_contiguity() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?; + let buffer = BufferExec::new(Arc::new(source), 1024); + + assert_eq!(buffer.properties().group_contiguous_exprs().len(), 1); + assert!( + buffer.properties().group_contiguous_exprs()[0].eq(&col("key", &schema)?) + ); + Ok(()) + } + #[tokio::test] async fn buffers_only_some_messages() -> Result<(), Box> { let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 17166e287e6dc..f6287b98b3e34 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -477,7 +477,9 @@ pub fn make_cooperative(stream: SendableRecordBatchStream) -> SendableRecordBatc mod tests { use super::*; - use arrow_schema::SchemaRef; + use crate::test::TestMemoryExec; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use datafusion_physical_expr::expressions::col; use futures::stream; @@ -497,6 +499,22 @@ mod tests { Box::pin(RecordBatchStreamAdapter::new(schema, s)) } + #[test] + fn cooperative_exec_preserves_group_contiguity() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?; + let cooperative = CooperativeExec::new(Arc::new(source)); + + assert_eq!(cooperative.properties().group_contiguous_exprs().len(), 1); + assert!( + cooperative.properties().group_contiguous_exprs()[0] + .eq(&col("key", &schema)?) + ); + Ok(()) + } + #[tokio::test] async fn yield_less_than_threshold() -> Result<()> { let count = TASK_BUDGET - 10; diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index c0fa4c6bede41..d15c2566322df 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -1245,6 +1245,28 @@ pub trait ExecutionPlanProperties { /// /// [`FilterExec`]: crate::filter::FilterExec fn equivalence_properties(&self) -> &EquivalenceProperties; + + /// Expressions whose complete tuple is contiguous within each output + /// partition. + /// + /// For every distinct tuple formed by these expressions, all matching rows + /// must occur in at most one contiguous range. Once a different tuple appears, + /// the previous tuple must not appear again. Tuple equality follows `GROUP BY` + /// semantics, and tuples need not be sorted. + /// + /// The guarantee applies to the complete tuple, not to individual expressions + /// or subsets. Expressions must refer to this plan's output schema. This + /// property implies neither output ordering nor distribution. The default is + /// no guarantee. + /// + /// # Correctness + /// + /// An invalid declaration can cause a streaming aggregate to emit a group + /// before observing all of its rows, producing incorrect results. Use + /// [`PlanProperties::with_group_contiguous_exprs`] to set this property. + fn group_contiguous_exprs(&self) -> &[Arc] { + &[] + } } impl ExecutionPlanProperties for Arc { @@ -1267,6 +1289,10 @@ impl ExecutionPlanProperties for Arc { fn equivalence_properties(&self) -> &EquivalenceProperties { self.properties().equivalence_properties() } + + fn group_contiguous_exprs(&self) -> &[Arc] { + self.properties().group_contiguous_exprs() + } } impl ExecutionPlanProperties for &dyn ExecutionPlan { @@ -1289,6 +1315,10 @@ impl ExecutionPlanProperties for &dyn ExecutionPlan { fn equivalence_properties(&self) -> &EquivalenceProperties { self.properties().equivalence_properties() } + + fn group_contiguous_exprs(&self) -> &[Arc] { + self.properties().group_contiguous_exprs() + } } /// Represents whether a stream of data **generated** by an operator is bounded (finite) @@ -1501,6 +1531,8 @@ pub struct PlanProperties { pub scheduling_type: SchedulingType, /// See [ExecutionPlanProperties::output_ordering] output_ordering: Option, + /// See [`ExecutionPlanProperties::group_contiguous_exprs`] + group_contiguous_exprs: Vec>, } impl PlanProperties { @@ -1521,6 +1553,7 @@ impl PlanProperties { evaluation_type: EvaluationType::Lazy, scheduling_type: SchedulingType::NonCooperative, output_ordering, + group_contiguous_exprs: vec![], } } @@ -1572,6 +1605,18 @@ impl PlanProperties { self } + /// Overwrite the group-contiguous composite key. + /// + /// See [`ExecutionPlanProperties::group_contiguous_exprs`] for the correctness + /// contract. + pub fn with_group_contiguous_exprs( + mut self, + group_contiguous_exprs: Vec>, + ) -> Self { + self.group_contiguous_exprs = group_contiguous_exprs; + self + } + /// Set constraints having mut reference. pub fn set_constraints(&mut self, constraints: Constraints) { self.eq_properties.set_constraints(constraints); @@ -1595,6 +1640,11 @@ impl PlanProperties { self.output_ordering.as_ref() } + /// Components of the group-contiguous composite key, if any. + pub fn group_contiguous_exprs(&self) -> &[Arc] { + &self.group_contiguous_exprs + } + /// Get schema of the node. pub(crate) fn schema(&self) -> &SchemaRef { self.eq_properties.schema() diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 2a24eb60e6fbc..108b5a79724f2 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -208,12 +208,15 @@ impl ProjectionExec { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = projector.projection().projection_mapping(&input.schema())?; + let group_contiguous_exprs = + Self::project_group_contiguous_exprs(&input, &projection_mapping); let cache = Self::compute_properties( &input, &projection_mapping, Arc::clone(projector.output_schema()), reuse_from, - )?; + )? + .with_group_contiguous_exprs(group_contiguous_exprs); Ok(Self { projector, input, @@ -222,6 +225,22 @@ impl ProjectionExec { }) } + /// Projects the complete group-contiguous tuple, dropping it when any + /// component cannot be mapped to the output schema. + fn project_group_contiguous_exprs( + input: &Arc, + projection_mapping: &ProjectionMapping, + ) -> Vec> { + input + .equivalence_properties() + .project_expressions( + input.group_contiguous_exprs().iter(), + projection_mapping, + ) + .collect::>>() + .unwrap_or_default() + } + /// The projection expressions stored as tuples of (expression, output column name) pub fn expr(&self) -> &[ProjectionExpr] { self.projector.projection().as_ref() @@ -1515,6 +1534,7 @@ mod tests { use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; + use crate::test::TestMemoryExec; use crate::test::exec::StatisticsExec; use arrow::datatypes::{DataType, Field, Schema}; @@ -1526,6 +1546,79 @@ mod tests { BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, }; + #[test] + fn group_contiguous_projection_is_all_or_nothing() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time", DataType::Int32, false), + ])); + let time_bin = binary( + col("time", &schema)?, + Operator::Divide, + lit(ScalarValue::Int32(Some(10))), + &schema, + )?; + let plain_source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?; + let source = plain_source.clone().try_with_group_contiguous_exprs(vec![ + col("key", &schema)?, + Arc::clone(&time_bin), + ])?; + + let projection = ProjectionExec::try_new( + [ + ProjectionExpr::new(col("key", &schema)?, "key"), + ProjectionExpr::new(time_bin, "time_bin"), + ], + Arc::new(source.clone()), + )?; + let projected_schema = projection.schema(); + let expected = [ + col("key", &projected_schema)?, + col("time_bin", &projected_schema)?, + ]; + assert_eq!( + projection.properties().group_contiguous_exprs().len(), + expected.len() + ); + assert!( + projection + .properties() + .group_contiguous_exprs() + .iter() + .zip(expected) + .all(|(actual, expected)| actual.eq(&expected)) + ); + + // A strict subset is not sufficient: contiguity of `(key, time_bin)` + // does not imply that `key` alone is contiguous. + let partial_projection = ProjectionExec::try_new( + [ProjectionExpr::new(col("key", &schema)?, "key")], + Arc::new(source.clone()), + )?; + assert!( + partial_projection + .properties() + .group_contiguous_exprs() + .is_empty() + ); + + // Row-preserving operators do not inherit the assertion unless they + // opt in explicitly. + let filter = FilterExec::try_new(lit(true), Arc::new(source.clone()))?; + assert!(filter.properties().group_contiguous_exprs().is_empty()); + + // Group contiguity participates in child-property identity, so + // replacing the source with one lacking the assertion recomputes the + // projection properties. + assert!(!Arc::ptr_eq(source.properties(), plain_source.properties())); + let projection: Arc = Arc::new(projection); + let replaced = + replace_children_if_necessary(projection, vec![Arc::new(plain_source)])?; + assert!(replaced.group_contiguous_exprs().is_empty()); + + Ok(()) + } + #[test] fn test_try_new_with_schema_metadata_only_replaces_metadata() -> Result<()> { let input_schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index cac8f925039ae..674930350e4f4 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -442,8 +442,26 @@ mod tests { use arrow::array::{Int32Array, Int64Array}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; + use datafusion_physical_expr::expressions::col; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + #[test] + fn scalar_subquery_exec_preserves_group_contiguity() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let source = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_group_contiguous_exprs(vec![col("key", &schema)?])?; + let exec = ScalarSubqueryExec::new( + Arc::new(source), + vec![], + ScalarSubqueryResults::new(0), + ); + + assert_eq!(exec.properties().group_contiguous_exprs().len(), 1); + assert!(exec.properties().group_contiguous_exprs()[0].eq(&col("key", &schema)?)); + Ok(()) + } + enum ExpectedSubqueryResult { Value(ScalarValue), Error(&'static str), diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index b38a46d160755..6b304c0b50379 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -73,6 +73,8 @@ pub struct TestMemoryExec { projection: Option>, /// Sort information: one or more equivalent orderings sort_information: Vec, + /// Composite key whose values are contiguous within each output stream. + group_contiguous_exprs: Vec>, /// if partition sizes should be displayed show_sizes: bool, /// The maximum number of records to read from this plan. If `None`, @@ -226,6 +228,7 @@ impl TestMemoryExec { EmissionType::Incremental, Boundedness::Bounded, ) + .with_group_contiguous_exprs(self.group_contiguous_exprs.clone()) } fn output_partitioning(&self) -> Partitioning { @@ -268,6 +271,7 @@ impl TestMemoryExec { projected_schema, projection, sort_information: vec![], + group_contiguous_exprs: vec![], show_sizes: true, fetch: None, }) @@ -363,6 +367,51 @@ impl TestMemoryExec { Ok(self) } + /// Attach a composite key whose values occur in one contiguous range in + /// each output stream. See + /// [`ExecutionPlanProperties::group_contiguous_exprs`](crate::ExecutionPlanProperties::group_contiguous_exprs) + /// for the correctness contract. + pub fn try_with_group_contiguous_exprs( + mut self, + mut group_contiguous_exprs: Vec>, + ) -> Result { + // All expressions must refer to the original schema. + let fields = self.schema.fields(); + let ambiguous_column = group_contiguous_exprs + .iter() + .flat_map(collect_columns) + .find(|col| { + fields + .get(col.index()) + .map(|field| field.name() != col.name()) + .unwrap_or(true) + }); + assert_or_internal_err!( + ambiguous_column.is_none(), + "Column {:?} is not found in the original schema of the TestMemoryExec", + ambiguous_column.as_ref().unwrap() + ); + + if let Some(projection) = &self.projection { + let base_schema = self.original_schema(); + let proj_exprs = projection.iter().map(|idx| { + let name = base_schema.field(*idx).name(); + (Arc::new(Column::new(name, *idx)) as _, name.to_string()) + }); + let projection_mapping = + ProjectionMapping::try_new(proj_exprs, &base_schema)?; + let base_eqp = EquivalenceProperties::new(base_schema); + group_contiguous_exprs = base_eqp + .project_expressions(group_contiguous_exprs.iter(), &projection_mapping) + .collect::>>() + .unwrap_or_default(); + } + + self.group_contiguous_exprs = group_contiguous_exprs; + self.cache = Arc::new(self.compute_properties()); + Ok(self) + } + /// Arc clone of ref to original schema pub fn original_schema(&self) -> SchemaRef { Arc::clone(&self.schema)