diff --git a/datafusion/core/tests/execution/coop.rs b/datafusion/core/tests/execution/coop.rs index e02364a0530cc..38791580a9efe 100644 --- a/datafusion/core/tests/execution/coop.rs +++ b/datafusion/core/tests/execution/coop.rs @@ -243,25 +243,22 @@ async fn agg_grouped_topk_yields( let group = binary(value_col.clone(), Divide, lit(1000000i64), &inf.schema())?; let aggr = Arc::new( - AggregateExec::try_new( - AggregateMode::Single, - PhysicalGroupBy::new( + AggregateExec::builder(AggregateMode::Single, inf.clone()) + .with_group_by(PhysicalGroupBy::new( vec![(group, "group".to_string())], vec![], vec![vec![false]], false, - ), - vec![Arc::new( + )) + .with_aggr_exprs(vec![Arc::new( AggregateExprBuilder::new(min_max::max_udaf(), vec![value_col.clone()]) .schema(inf.schema()) .alias("max") .build()?, - )], - vec![None], - inf.clone(), - inf.schema(), - )? - .with_limit_options(Some(LimitOptions::new(100))), + )]) + .with_input_schema(inf.schema()) + .with_limit_options(LimitOptions::new(100)) + .build()?, ); query_yields(aggr, session_ctx.task_ctx()).await diff --git a/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs b/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs index 9e63c341c92d9..49fb7f2cf3a69 100644 --- a/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs +++ b/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs @@ -251,16 +251,12 @@ fn aggregations_with_limit_combined() -> datafusion_common::Result<()> { let schema = partial_agg.schema(); let final_agg = Arc::new( - AggregateExec::try_new( - AggregateMode::Final, - final_group_by, - aggr_expr, - vec![], - partial_agg, - schema, - ) - .unwrap() - .with_limit_options(Some(LimitOptions::new(5))), + AggregateExec::builder(AggregateMode::Final, partial_agg) + .with_group_by(final_group_by) + .with_aggr_exprs(aggr_expr) + .with_input_schema(schema) + .with_limit_options(LimitOptions::new(5)) + .build()?, ); let plan: Arc = final_agg; // should combine the Partial/Final AggregateExecs to a Single AggregateExec diff --git a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs index 297a92c45a16d..de8535a86e7dd 100644 --- a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs +++ b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs @@ -90,19 +90,17 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate { } else { AggregateMode::SinglePartitioned }; - AggregateExec::try_new( - mode, - input_agg_exec.group_expr().clone(), - input_agg_exec.aggr_expr().to_vec(), - input_agg_exec.filter_expr().to_vec(), - Arc::clone(input_agg_exec.input()), - input_agg_exec.input_schema(), - ) - .map(|combined_agg| { - combined_agg.with_limit_options(agg_exec.limit_options()) - }) - .ok() - .map(Arc::new) + // the partial aggregate, re-planned in the combined mode and + // carrying the limit that was pushed into the final aggregate + input_agg_exec + .to_builder() + .with_mode(mode) + .with_limit_options(agg_exec.limit_options()) + .build() + // the combined aggregate cannot execute the limit: leave + // the two aggregates alone + .ok() + .map(Arc::new) } else { None }; diff --git a/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs b/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs index 192a139f36021..68b3a37627897 100644 --- a/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs +++ b/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs @@ -54,7 +54,12 @@ impl LimitedDistinctAggregation { } // We found what we want: clone, copy the limit down, and return modified node - let new_aggr = aggr.with_new_limit_options(Some(LimitOptions::new(limit))); + let new_aggr = aggr + .to_builder() + .with_limit_options(LimitOptions::new(limit)) + .build() + // the aggregate cannot execute the limit: leave the plan alone + .ok()?; Some(Arc::new(new_aggr)) } diff --git a/datafusion/physical-optimizer/src/topk_aggregation.rs b/datafusion/physical-optimizer/src/topk_aggregation.rs index 0eddb5d5507e4..3650c6698626c 100644 --- a/datafusion/physical-optimizer/src/topk_aggregation.rs +++ b/datafusion/physical-optimizer/src/topk_aggregation.rs @@ -106,10 +106,12 @@ impl TopKAggregation { } // We found what we want: clone, copy the limit down, and return modified node - let new_aggr = AggregateExec::with_new_limit_options( - aggr, - Some(LimitOptions::new_with_order(limit, order_desc)), - ); + let new_aggr = aggr + .to_builder() + .with_limit_options(LimitOptions::new_with_order(limit, order_desc)) + .build() + // the aggregate cannot execute the limit: leave the plan alone + .ok()?; Some(Arc::new(new_aggr)) } diff --git a/datafusion/physical-optimizer/src/update_aggr_exprs.rs b/datafusion/physical-optimizer/src/update_aggr_exprs.rs index a047ab11f423a..3937d75b63da6 100644 --- a/datafusion/physical-optimizer/src/update_aggr_exprs.rs +++ b/datafusion/physical-optimizer/src/update_aggr_exprs.rs @@ -125,7 +125,8 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder { input.equivalence_properties(), )?; - let aggr_exec = aggr_exec.with_new_aggr_exprs(aggr_exprs); + let aggr_exec = + aggr_exec.to_builder().with_aggr_exprs(aggr_exprs).build()?; Ok(Transformed::yes(Arc::new(aggr_exec) as _)) } else { diff --git a/datafusion/physical-plan/src/aggregates/builder.rs b/datafusion/physical-plan/src/aggregates/builder.rs new file mode 100644 index 0000000000000..676264fa44ef6 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/builder.rs @@ -0,0 +1,606 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`AggregateExecBuilder`]: build and rewrite [`AggregateExec`] nodes + +use std::sync::Arc; + +use super::{ + AggrDynFilter, AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, + create_schema, get_finer_aggregate_exprs_requirement, +}; +use crate::metrics::ExecutionPlanMetricsSet; +use crate::{ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PlanProperties}; + +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::aggregate::AggregateFunctionExpr; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::{ + LexRequirement, OrderingRequirements, PhysicalSortRequirement, +}; + +/// The `FILTER` expression of each aggregate expression, `None` where an +/// aggregate has no filter. +type FilterExprs = Arc<[Option>]>; + +/// Builds an [`AggregateExec`], and is the single place one is constructed. +/// +/// Reached through [`AggregateExec::builder`] for a new node and +/// [`AggregateExec::to_builder`] for a rewrite of an existing one; a rewrite +/// keeps the derived state (output schema, plan properties, ordering +/// requirements, dynamic filter) of the node it came from unless a field it is +/// computed from changes, so it costs no more than the hand-written +/// clone-with-one-change methods it replaces. +/// +/// `build` is fallible so that the node can be checked here, in the one place +/// it is built, rather than at execution time. It currently only checks that +/// the aggregate and `FILTER` expressions have the same length. +/// +/// Public for internal use only: this is how DataFusion's own physical +/// optimizer rules build and rewrite aggregates, and it may change without +/// notice. +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow::datatypes::{DataType, Field, Schema}; +/// # use datafusion_physical_plan::aggregates::{ +/// # AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, +/// # }; +/// # use datafusion_physical_plan::{ExecutionPlan, empty::EmptyExec}; +/// # use datafusion_physical_expr::expressions::col; +/// # fn main() -> datafusion_common::Result<()> { +/// # let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); +/// # let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); +/// # let group_by = +/// # PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); +/// let exec = AggregateExec::builder(AggregateMode::Single, input) +/// .with_group_by(group_by) +/// .build()?; +/// +/// // push a limit into it, keeping its schema and plan properties +/// let limited = exec.to_builder().with_limit_options(LimitOptions::new(10)).build()?; +/// assert_eq!(limited.schema(), exec.schema()); +/// # Ok(()) +/// # } +/// ``` +#[doc(hidden)] +#[derive(Debug, Clone)] +pub struct AggregateExecBuilder { + mode: AggregateMode, + group_by: Arc, + aggr_expr: Arc<[Arc]>, + /// `None` means "no filter for any aggregate expression" + filter_expr: Option, + input: Arc, + /// `None` means "the schema of `input`" + input_schema: Option, + limit_options: Option, + /// Output schema explicitly supplied by the caller, see + /// [`AggregateExecBuilder::with_output_schema`]. Always honored. + output_schema: Option, + /// State carried over from the [`AggregateExec`] this builder was derived + /// from, dropped as soon as a field it is computed from changes. + derived: Option, +} + +/// State of an [`AggregateExec`] that is computed from its inputs, and which is +/// preserved verbatim when a node is rewritten without touching what it is +/// computed from. +#[derive(Debug, Clone)] +struct DerivedState { + schema: SchemaRef, + cache: Arc, + required_input_ordering: Option, + input_order_mode: InputOrderMode, + dynamic_filter: Option>, +} + +impl AggregateExecBuilder { + /// Create a builder for an aggregate over `input`. + /// + /// Unless overridden the aggregate has no group by expressions, no + /// aggregate expressions, no filters, no limit, and uses the schema of + /// `input` as its [input schema](AggregateExec::input_schema). + pub fn new(mode: AggregateMode, input: Arc) -> Self { + Self { + mode, + group_by: Arc::new(PhysicalGroupBy::default()), + aggr_expr: Arc::from([]), + filter_expr: None, + input, + input_schema: None, + limit_options: None, + output_schema: None, + derived: None, + } + } + + /// Create a builder pre-populated from `exec`. + /// + /// Takes `&AggregateExec` rather than ownership because every caller holds + /// a borrow from `downcast_ref` on an `Arc`; nothing is + /// deep-copied, the fields are `Arc`s. + pub(crate) fn from_exec(exec: &AggregateExec) -> Self { + Self { + mode: exec.mode, + group_by: Arc::clone(&exec.group_by), + aggr_expr: Arc::clone(&exec.aggr_expr), + filter_expr: Some(Arc::clone(&exec.filter_expr)), + input: Arc::clone(&exec.input), + input_schema: Some(Arc::clone(&exec.input_schema)), + limit_options: exec.limit_options, + output_schema: None, + derived: Some(DerivedState { + schema: Arc::clone(&exec.schema), + cache: Arc::clone(&exec.cache), + required_input_ordering: exec.required_input_ordering.clone(), + input_order_mode: exec.input_order_mode.clone(), + dynamic_filter: exec.dynamic_filter.clone(), + }), + } + } + + /// Set the [`AggregateMode`]. + pub fn with_mode(mut self, mode: AggregateMode) -> Self { + if mode == self.mode { + return self; + } + self.mode = mode; + self.invalidate_derived() + } + + /// Set the group by expressions. + pub fn with_group_by(mut self, group_by: impl Into>) -> Self { + let group_by = group_by.into(); + if Arc::ptr_eq(&self.group_by, &group_by) || *self.group_by == *group_by { + return self; + } + self.group_by = group_by; + self.invalidate_derived() + } + + /// Set the aggregate expressions. + /// + /// A builder derived from an existing node keeps that node's output schema, + /// so rewriting the aggregate expressions (for example reversing them in + /// `OptimizeAggregateOrder`) cannot change output field names. This matches + /// the `AggregateExec::with_new_aggr_exprs` it replaces; nothing checks that + /// the new expressions still describe that schema. + pub fn with_aggr_exprs( + mut self, + aggr_expr: impl Into]>>, + ) -> Self { + self.aggr_expr = aggr_expr.into(); + self + } + + /// Set the `FILTER` expression of each aggregate expression. + /// + /// Must have the same length as the aggregate expressions; `build` returns + /// an error otherwise. If never called, no aggregate is filtered. + pub fn with_filter_exprs(mut self, filter_expr: impl Into) -> Self { + let filter_expr = filter_expr.into(); + if self.filter_expr.as_ref().is_some_and(|existing| { + Arc::ptr_eq(existing, &filter_expr) || **existing == *filter_expr + }) { + return self; + } + self.filter_expr = Some(filter_expr); + self.invalidate_derived() + } + + /// Set the input plan. + pub fn with_input(mut self, input: Arc) -> Self { + if Arc::ptr_eq(&self.input, &input) { + return self; + } + self.input = input; + self.invalidate_derived() + } + + /// Set the [input schema](AggregateExec::input_schema): the schema of the + /// data *before* any aggregation is applied. + /// + /// For `Partial` and `Single` aggregates this is the schema of the input + /// plan (the default). For `Final` and `FinalPartitioned` aggregates it is + /// the input schema of the matching partial aggregate, which is *not* the + /// schema of the input plan. + pub fn with_input_schema(mut self, input_schema: SchemaRef) -> Self { + self.input_schema = Some(input_schema); + self + } + + /// Set the limit pushed down into this aggregate, or `None` to remove it. + /// + /// The limit is a hint: operators above the aggregate still enforce it. + /// Accepts both `LimitOptions` and `Option`. + /// + /// Note that not every aggregate can execute every limit. `build` does not + /// check that yet, so the caller still owns it, exactly as it did before + /// this builder existed. + pub fn with_limit_options( + mut self, + limit_options: impl Into>, + ) -> Self { + self.limit_options = limit_options.into(); + self + } + + /// Use `schema` as the output schema instead of computing it. + /// + /// For callers that must preserve a schema exactly, such as decoding a + /// serialized plan. The caller owns the schema being correct. + pub(crate) fn with_output_schema(mut self, schema: SchemaRef) -> Self { + self.output_schema = Some(schema); + self + } + + /// Drop state derived from the node this builder came from, because a field + /// it is computed from was replaced. + fn invalidate_derived(mut self) -> Self { + self.derived = None; + self + } + + /// Build the [`AggregateExec`]. + pub fn build(self) -> Result { + let Self { + mode, + group_by, + aggr_expr, + filter_expr, + input, + input_schema, + limit_options, + output_schema, + derived, + } = self; + + let input_schema = input_schema.unwrap_or_else(|| input.schema()); + let filter_expr = filter_expr + .unwrap_or_else(|| std::iter::repeat_n(None, aggr_expr.len()).collect()); + + assert_eq_or_internal_err!( + aggr_expr.len(), + filter_expr.len(), + "Inconsistent aggregate expr: {:?} and filter expr: {:?} for AggregateExec, their size should match", + aggr_expr, + filter_expr + ); + + let mut exec = match derived { + // Nothing the derived state is computed from changed: clone the + // node this builder came from with the new values rather than + // recomputing. In particular its output schema is kept, so a + // rewrite of the aggregate expressions cannot rename output fields. + Some(derived) if output_schema.is_none() => AggregateExec { + mode, + group_by, + aggr_expr, + filter_expr, + input, + schema: derived.schema, + input_schema, + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering: derived.required_input_ordering, + input_order_mode: derived.input_order_mode, + cache: derived.cache, + limit_options: None, + dynamic_filter: derived.dynamic_filter, + }, + _ => build_from_scratch( + mode, + group_by, + &aggr_expr, + filter_expr, + input, + input_schema, + output_schema, + )?, + }; + + exec.limit_options = limit_options; + Ok(exec) + } +} + +/// Compute every derived part of an [`AggregateExec`] from its inputs: the +/// output schema (unless `output_schema` supplies one), the ordering the input +/// must have, how the input is ordered relative to the group by, the plan +/// properties, and the dynamic filter. +fn build_from_scratch( + mode: AggregateMode, + group_by: Arc, + aggr_expr: &[Arc], + filter_expr: FilterExprs, + input: Arc, + input_schema: SchemaRef, + output_schema: Option, +) -> Result { + // `get_finer_aggregate_exprs_requirement` may rewrite the aggregate + // expressions (e.g. reverse them), so it needs them owned. + let mut aggr_expr = aggr_expr.to_vec(); + + // The output schema is computed from the aggregate expressions *as given*, + // before the requirement analysis below may rewrite them: output field + // names come from those expressions and must not change as a side effect + // of a rewrite. This is why an explicitly supplied schema exists at all. + let schema = match output_schema { + Some(schema) => schema, + None => Arc::new(create_schema(&input.schema(), &group_by, &aggr_expr, mode)?), + }; + + let input_eq_properties = input.equivalence_properties(); + // Get GROUP BY expressions: + let groupby_exprs = group_by.input_exprs(); + // If existing ordering satisfies a prefix of the GROUP BY expressions, + // prefix requirements with this section. In this case, aggregation will + // work more efficiently. + // Copy the `PhysicalSortExpr`s to retain the sort options. + let (new_sort_exprs, indices) = + input_eq_properties.find_longest_permutation(&groupby_exprs)?; + + let mut new_requirements = new_sort_exprs + .into_iter() + .map(PhysicalSortRequirement::from) + .collect::>(); + + let req = get_finer_aggregate_exprs_requirement( + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + new_requirements.extend(req); + + let required_input_ordering = + LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft); + + // Constant expressions never change, so they cannot mark a completed group. + // Exclude them from both the ordering indices and the group expression count. + // If our aggregation has grouping sets then our base grouping exprs will + // be expanded based on the flags in `group_by.groups` where for each + // group we swap the grouping expr for `null` if the flag is `true` + // That means that each index in `indices` is valid if and only if + // it is not null in every group + let indices: Vec = indices + .into_iter() + .filter(|idx| group_by.groups.iter().all(|group| !group[*idx])) + .filter(|idx| { + input_eq_properties + .is_expr_constant(&groupby_exprs[*idx]) + .is_none() + }) + .collect(); + + let num_non_constant_groupby_exprs = groupby_exprs + .iter() + .filter(|expr| input_eq_properties.is_expr_constant(expr).is_none()) + .count(); + let mut input_order_mode = if indices.len() == num_non_constant_groupby_exprs + && !indices.is_empty() + && group_by.groups.len() == 1 + { + InputOrderMode::Sorted + } else if !indices.is_empty() { + InputOrderMode::PartiallySorted(indices) + } else { + InputOrderMode::Linear + }; + + // Input order mode is also used to advertise plan output ordering, grouping + // sets handling, and partial reduce aggregation can't promise that. + if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce { + input_order_mode = InputOrderMode::Linear; + } + + // 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())?; + + let cache = if group_by.has_grouping_set() { + AggregateExec::compute_grouping_set_properties(&input, Arc::clone(&schema)) + } else { + AggregateExec::compute_properties( + &input, + Arc::clone(&schema), + &group_expr_mapping, + group_by.is_true_no_grouping(), + &mode, + &input_order_mode, + aggr_expr.as_ref(), + )? + }; + + let mut exec = AggregateExec { + mode, + group_by, + aggr_expr: aggr_expr.into(), + filter_expr, + input, + schema, + input_schema, + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering, + limit_options: None, + input_order_mode, + cache: Arc::new(cache), + dynamic_filter: None, + }; + + exec.init_dynamic_filter(); + + Ok(exec) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::empty::EmptyExec; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::min_max::min_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + ])) + } + + fn test_input(schema: &SchemaRef) -> Arc { + Arc::new(EmptyExec::new(Arc::clone(schema))) + } + + fn group_by_a(schema: &SchemaRef) -> Result { + Ok(PhysicalGroupBy::new_single(vec![( + col("a", schema)?, + "a".to_string(), + )])) + } + + fn min_b(schema: &SchemaRef) -> Result> { + Ok(Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("min_b") + .build()?, + )) + } + + fn count_b(schema: &SchemaRef) -> Result> { + Ok(Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("count_b") + .build()?, + )) + } + + /// `filter_expr` defaults to "no filter" instead of having to be a vector of + /// `None`s of exactly the right length. + #[test] + fn filter_exprs_default_to_none() -> Result<()> { + let schema = test_schema(); + let exec = AggregateExec::builder(AggregateMode::Single, test_input(&schema)) + .with_group_by(group_by_a(&schema)?) + .with_aggr_exprs(vec![min_b(&schema)?, count_b(&schema)?]) + .build()?; + assert_eq!(exec.filter_expr(), &[None, None]); + Ok(()) + } + + #[test] + fn mismatched_filter_exprs_are_rejected() -> Result<()> { + let schema = test_schema(); + let err = AggregateExec::builder(AggregateMode::Single, test_input(&schema)) + .with_group_by(group_by_a(&schema)?) + .with_aggr_exprs(vec![min_b(&schema)?]) + .with_filter_exprs(vec![]) + .build() + .unwrap_err(); + assert!( + err.message().contains("their size should match"), + "unexpected error: {err}" + ); + Ok(()) + } + + /// Rewriting a node keeps the output schema and plan properties of the node + /// it was derived from, and resets its metrics. + #[test] + fn rewriting_preserves_derived_state() -> Result<()> { + let schema = test_schema(); + let exec = AggregateExec::builder(AggregateMode::Single, test_input(&schema)) + .with_group_by(group_by_a(&schema)?) + .with_aggr_exprs(vec![min_b(&schema)?]) + .build()?; + + let limited = exec + .to_builder() + .with_limit_options(LimitOptions::new(10)) + .build()?; + + assert_eq!(limited.limit_options(), Some(LimitOptions::new(10))); + assert_eq!(limited.schema(), exec.schema()); + assert_eq!(limited.input_schema(), exec.input_schema()); + assert_eq!(limited.mode(), exec.mode()); + // the plan properties were reused rather than recomputed + assert!(Arc::ptr_eq(&limited.cache, &exec.cache)); + // but the metrics of the original node were not carried over + assert_eq!(limited.metrics().unwrap().iter().count(), 0); + Ok(()) + } + + /// Changing the mode is a structural change: the output schema of a + /// `Partial` aggregate holds intermediate state, so it must be recomputed. + #[test] + fn changing_the_mode_recomputes_the_schema() -> Result<()> { + let schema = test_schema(); + let partial = AggregateExec::builder(AggregateMode::Partial, test_input(&schema)) + .with_group_by(group_by_a(&schema)?) + .with_aggr_exprs(vec![count_b(&schema)?]) + .build()?; + let single = partial + .to_builder() + .with_mode(AggregateMode::Single) + .build()?; + + // `Partial` emits the accumulator state, `Single` the final count + assert_eq!(partial.schema().field(1).data_type(), &DataType::Int64); + assert_eq!(single.schema().field(1).data_type(), &DataType::Int64); + assert_ne!( + partial.schema().field(1).name(), + single.schema().field(1).name() + ); + Ok(()) + } + + /// Setting a field to the value it already has must not invalidate the + /// derived state: a rewrite that changes nothing should cost nothing. + #[test] + fn setting_a_field_to_its_current_value_keeps_derived_state() -> Result<()> { + let schema = test_schema(); + let exec = AggregateExec::builder(AggregateMode::Single, test_input(&schema)) + .with_group_by(group_by_a(&schema)?) + .with_aggr_exprs(vec![min_b(&schema)?]) + .build()?; + + let unchanged = exec + .to_builder() + .with_mode(AggregateMode::Single) + .with_group_by(group_by_a(&schema)?) + .with_input(Arc::clone(exec.input())) + .with_filter_exprs(vec![None]) + .with_aggr_exprs(exec.aggr_expr().to_vec()) + .build()?; + assert!(Arc::ptr_eq(&unchanged.cache, &exec.cache)); + + // and a real change still does invalidate it + let changed = exec + .to_builder() + .with_mode(AggregateMode::Partial) + .build()?; + assert!(!Arc::ptr_eq(&changed.cache, &exec.cache)); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 13ead739309d3..7d27d98c97cf7 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -365,15 +365,12 @@ mod tests { .build()?, ); let aggregate_exec = Arc::new( - AggregateExec::try_new( - AggregateMode::Single, - group_by, - vec![aggregate], - vec![None], - input, - schema, - )? - .with_limit_options(Some(LimitOptions::new(2))), + AggregateExec::builder(AggregateMode::Single, input) + .with_group_by(group_by) + .with_aggr_exprs(vec![aggregate]) + .with_input_schema(schema) + .with_limit_options(LimitOptions::new(2)) + .build()?, ); let context = Arc::new(TaskContext::default()); let result = collect(Arc::clone(&aggregate_exec) as _, context).await?; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index ffa93a387d87a..043755b73f987 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -187,8 +187,8 @@ use arrow_schema::FieldRef; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ - ColumnStatistics, Constraint, Constraints, Result, ScalarValue, - assert_eq_or_internal_err, internal_err, not_impl_err, + ColumnStatistics, Constraint, Constraints, Result, ScalarValue, internal_err, + not_impl_err, }; use datafusion_execution::TaskContext; use datafusion_expr::{Accumulator, Aggregate, AggregateMetrics}; @@ -200,7 +200,7 @@ use datafusion_physical_expr::{ }; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, fmt_sql}; use datafusion_physical_expr_common::sort_expr::{ - LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortRequirement, + LexOrdering, OrderingRequirements, PhysicalSortRequirement, }; use datafusion_expr::utils::AggregateOrderSensitivity; @@ -211,6 +211,7 @@ use topk::heap::is_supported_heap_type; mod aggregate_hash_table; mod aggregate_stream; +mod builder; pub mod group_values; mod grouped_hash_stream; mod grouped_topk_stream; @@ -224,6 +225,9 @@ mod single_stream; mod skip_partial; mod topk; +#[doc(hidden)] +pub use builder::AggregateExecBuilder; + /// Returns true if TopK aggregation data structures support the provided key and value types. /// /// This function checks whether both the key type (used for grouping) and value type @@ -905,9 +909,39 @@ pub struct AggregateExec { } impl AggregateExec { + /// Create a builder for a new [`AggregateExec`] over `input`, see + /// [`AggregateExecBuilder`]. + /// + /// Public for internal use only and not part of the public API. + #[doc(hidden)] + pub fn builder( + mode: AggregateMode, + input: Arc, + ) -> AggregateExecBuilder { + AggregateExecBuilder::new(mode, input) + } + + /// Create a builder pre-populated with the fields of this + /// [`AggregateExec`], to derive a new node from it. + /// + /// This is the supported way to rewrite an existing aggregate: the derived + /// output schema and plan properties are carried over, so a rewrite cannot + /// rename output fields. See [`AggregateExecBuilder`]. + /// + /// Public for internal use only and not part of the public API. + #[doc(hidden)] + pub fn to_builder(&self) -> AggregateExecBuilder { + AggregateExecBuilder::from_exec(self) + } + /// Function used in `OptimizeAggregateOrder` optimizer rule, /// where we need parts of the new value, others cloned from the old one /// Rewrites aggregate exec with new aggregate expressions. + #[doc(hidden)] + #[deprecated( + since = "56.0.0", + note = "use `AggregateExec::to_builder().with_aggr_exprs(..).build()` instead" + )] pub fn with_new_aggr_exprs( &self, aggr_expr: impl Into]>>, @@ -931,6 +965,11 @@ impl AggregateExec { } /// Clone this exec, overriding only the limit hint. + #[doc(hidden)] + #[deprecated( + since = "56.0.0", + note = "use `AggregateExec::to_builder().with_limit_options(..).build()` instead" + )] pub fn with_new_limit_options(&self, limit_options: Option) -> Self { Self { limit_options, @@ -955,6 +994,10 @@ impl AggregateExec { } /// Create a new hash aggregate execution plan + /// + /// Delegates to [`AggregateExecBuilder`], which is where an + /// `AggregateExec` is built. DataFusion's own optimizer rules use that + /// builder directly, so each argument is named. pub fn try_new( mode: AggregateMode, group_by: impl Into>, @@ -963,19 +1006,12 @@ impl AggregateExec { input: Arc, input_schema: SchemaRef, ) -> Result { - let group_by = group_by.into(); - let schema = create_schema(&input.schema(), &group_by, &aggr_expr, mode)?; - - let schema = Arc::new(schema); - AggregateExec::try_new_with_schema( - mode, - group_by, - aggr_expr, - filter_expr, - input, - input_schema, - schema, - ) + Self::builder(mode, input) + .with_group_by(group_by) + .with_aggr_exprs(aggr_expr) + .with_filter_exprs(filter_expr) + .with_input_schema(input_schema) + .build() } /// Create a new hash aggregate execution plan with the given schema. @@ -989,125 +1025,19 @@ impl AggregateExec { fn try_new_with_schema( mode: AggregateMode, group_by: impl Into>, - mut aggr_expr: Vec>, + aggr_expr: Vec>, filter_expr: impl Into>]>>, input: Arc, input_schema: SchemaRef, schema: SchemaRef, ) -> Result { - let group_by = group_by.into(); - let filter_expr = filter_expr.into(); - - // Make sure arguments are consistent in size - assert_eq_or_internal_err!( - aggr_expr.len(), - filter_expr.len(), - "Inconsistent aggregate expr: {:?} and filter expr: {:?} for AggregateExec, their size should match", - aggr_expr, - filter_expr - ); - - let input_eq_properties = input.equivalence_properties(); - // Get GROUP BY expressions: - let groupby_exprs = group_by.input_exprs(); - // If existing ordering satisfies a prefix of the GROUP BY expressions, - // prefix requirements with this section. In this case, aggregation will - // work more efficiently. - // Copy the `PhysicalSortExpr`s to retain the sort options. - let (new_sort_exprs, indices) = - input_eq_properties.find_longest_permutation(&groupby_exprs)?; - - let mut new_requirements = new_sort_exprs - .into_iter() - .map(PhysicalSortRequirement::from) - .collect::>(); - - let req = get_finer_aggregate_exprs_requirement( - &mut aggr_expr, - &group_by, - input_eq_properties, - &mode, - )?; - new_requirements.extend(req); - - let required_input_ordering = - LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft); - - // Constant expressions never change, so they cannot mark a completed group. - // Exclude them from both the ordering indices and the group expression count. - // If our aggregation has grouping sets then our base grouping exprs will - // be expanded based on the flags in `group_by.groups` where for each - // group we swap the grouping expr for `null` if the flag is `true` - // That means that each index in `indices` is valid if and only if - // it is not null in every group - let indices: Vec = indices - .into_iter() - .filter(|idx| group_by.groups.iter().all(|group| !group[*idx])) - .filter(|idx| { - input_eq_properties - .is_expr_constant(&groupby_exprs[*idx]) - .is_none() - }) - .collect(); - - let num_non_constant_groupby_exprs = groupby_exprs - .iter() - .filter(|expr| input_eq_properties.is_expr_constant(expr).is_none()) - .count(); - let mut input_order_mode = if indices.len() == num_non_constant_groupby_exprs - && !indices.is_empty() - && group_by.groups.len() == 1 - { - InputOrderMode::Sorted - } else if !indices.is_empty() { - InputOrderMode::PartiallySorted(indices) - } else { - InputOrderMode::Linear - }; - - // Input order mode is also used to advertise plan output ordering, grouping - // sets handling, and partial reduce aggregation can't promise that. - if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce { - input_order_mode = InputOrderMode::Linear; - } - - // 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())?; - - let cache = if group_by.has_grouping_set() { - Self::compute_grouping_set_properties(&input, Arc::clone(&schema)) - } else { - Self::compute_properties( - &input, - Arc::clone(&schema), - &group_expr_mapping, - group_by.is_true_no_grouping(), - &mode, - &input_order_mode, - aggr_expr.as_ref(), - )? - }; - - let mut exec = AggregateExec { - mode, - group_by, - aggr_expr: aggr_expr.into(), - filter_expr, - input, - schema, - input_schema, - metrics: ExecutionPlanMetricsSet::new(), - required_input_ordering, - limit_options: None, - input_order_mode, - cache: Arc::new(cache), - dynamic_filter: None, - }; - - exec.init_dynamic_filter(); - - Ok(exec) + Self::builder(mode, input) + .with_group_by(group_by) + .with_aggr_exprs(aggr_expr) + .with_filter_exprs(filter_expr) + .with_input_schema(input_schema) + .with_output_schema(schema) + .build() } /// Aggregation mode (full, partial) @@ -1116,12 +1046,25 @@ impl AggregateExec { } /// Set the limit options for this AggExec + #[doc(hidden)] + #[deprecated( + since = "56.0.0", + note = "use `AggregateExec::to_builder().with_limit_options(..).build()` instead" + )] pub fn with_limit_options(mut self, limit_options: Option) -> Self { self.limit_options = limit_options; self } /// Get the limit options (if set) + /// + /// Set them with + /// [`to_builder().with_limit_options(..)`](AggregateExec::to_builder). + /// + /// This is public for internal use only and is not part of the public API. + /// Unlike the setters it is not deprecated: it has no replacement, and + /// reading the limit of an aggregate is safe. + #[doc(hidden)] pub fn limit_options(&self) -> Option { self.limit_options } @@ -2677,37 +2620,26 @@ impl AggregateExec { .collect::>>()?; let group_by = PhysicalGroupBy::new(group_expr, null_expr, groups, *has_grouping_set); - let aggregate = if let Some(schema) = schema { - let schema = SchemaRef::new(schema.try_into()?); - AggregateExec::try_new_with_schema( - mode, - group_by, - aggr_expr, - filter_expr, - input, - Arc::clone(&input_schema), - schema, - ) - } else { - AggregateExec::try_new( - mode, - group_by, - aggr_expr, - filter_expr, - input, - Arc::clone(&input_schema), - ) - }?; - let aggregate = if let Some(limit) = limit { - let fetch = usize_from_wire(limit.limit, "AggregateExec", "limit")?; - let options = match limit.descending { - Some(descending) => LimitOptions::new_with_order(fetch, descending), - None => LimitOptions::new(fetch), - }; - aggregate.with_limit_options(Some(options)) - } else { - aggregate + let limit_options = match limit { + Some(limit) => { + let fetch = usize_from_wire(limit.limit, "AggregateExec", "limit")?; + Some(match limit.descending { + Some(descending) => LimitOptions::new_with_order(fetch, descending), + None => LimitOptions::new(fetch), + }) + } + None => None, }; + let mut builder = AggregateExec::builder(mode, input) + .with_group_by(group_by) + .with_aggr_exprs(aggr_expr) + .with_filter_exprs(filter_expr) + .with_input_schema(Arc::clone(&input_schema)) + .with_limit_options(limit_options); + if let Some(schema) = schema { + builder = builder.with_output_schema(SchemaRef::new(schema.try_into()?)); + } + let aggregate = builder.build()?; let aggregate = if let Some(dynamic_filter) = dynamic_filter { let dynamic_filter = ctx.decode_expr(dynamic_filter, input_schema.as_ref())?; @@ -4445,15 +4377,11 @@ mod tests { None, )?; let partial_aggregate = Arc::new( - AggregateExec::try_new( - AggregateMode::Partial, - group_by.clone(), - vec![], - vec![], - partial_input, - Arc::clone(&schema), - )? - .with_limit_options(Some(LimitOptions::new(2))), + AggregateExec::builder(AggregateMode::Partial, partial_input) + .with_group_by(group_by.clone()) + .with_input_schema(Arc::clone(&schema)) + .with_limit_options(LimitOptions::new(2)) + .build()?, ); let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?; @@ -4486,15 +4414,11 @@ mod tests { let final_input = TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; let final_aggregate = Arc::new( - AggregateExec::try_new( - AggregateMode::Final, - group_by.as_final(), - vec![], - vec![], - final_input, - Arc::clone(&schema), - )? - .with_limit_options(Some(LimitOptions::new(2))), + AggregateExec::builder(AggregateMode::Final, final_input) + .with_group_by(group_by.as_final()) + .with_input_schema(Arc::clone(&schema)) + .with_limit_options(LimitOptions::new(2)) + .build()?, ); let final_stream = final_aggregate.execute_typed(0, &task_ctx)?; @@ -6953,20 +6877,21 @@ mod tests { let input = Arc::new(StatisticsExec::new(stats, (**schema).clone())) as Arc; - let mut agg = AggregateExec::try_new( - mode, - group_by, - vec![count_a_aggregate(schema)?], - vec![None], - input, - Arc::clone(schema), - )?; - - if let Some(limit) = limit { - agg = agg.with_limit_options(Some(limit)); - } + // A limit is only ever pushed into an aggregate that can execute it. + // Without a MIN/MAX aggregate to order by, that means a `SELECT + // DISTINCT`-style aggregate with no aggregate expressions. + let aggr_exprs = if limit.is_some() { + vec![] + } else { + vec![count_a_aggregate(schema)?] + }; - Ok(agg) + AggregateExec::builder(mode, input) + .with_group_by(group_by) + .with_aggr_exprs(aggr_exprs) + .with_input_schema(Arc::clone(schema)) + .with_limit_options(limit) + .build() } fn simple_group_by(schema: &SchemaRef, cols: &[&str]) -> PhysicalGroupBy { diff --git a/datafusion/proto/tests/cases/plans/aggregates.rs b/datafusion/proto/tests/cases/plans/aggregates.rs index e57ac9fb5045b..675dcad7919ff 100644 --- a/datafusion/proto/tests/cases/plans/aggregates.rs +++ b/datafusion/proto/tests/cases/plans/aggregates.rs @@ -213,15 +213,15 @@ fn roundtrip_aggregate_with_limit() -> Result<()> { .map(Arc::new)?, ]; - let agg = AggregateExec::try_new( + let agg = AggregateExec::builder( AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], Arc::new(EmptyExec::new(schema.clone())), - schema, - )?; - let agg = agg.with_limit_options(Some(LimitOptions::new_with_order(12, false))); + ) + .with_group_by(PhysicalGroupBy::new_single(groups.clone())) + .with_aggr_exprs(aggregates) + .with_input_schema(schema) + .with_limit_options(LimitOptions::new_with_order(12, false)) + .build()?; roundtrip_test(Arc::new(agg)) } diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 7e6543cefc778..c94864c925403 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -466,3 +466,59 @@ let plan = DistributionContext::new_default(plan) Callers that only need the previous behaviour can keep using the deprecated form, or pass a freshly constructed `StatisticsContext` per call. + +### `AggregateExec` is built and rewritten with `AggregateExecBuilder` + +`AggregateExec::try_new` takes six positional arguments, two of which are +schemas that are easy to transpose, and the fields that optimizer rules change +afterwards (the limit hint, the aggregate expressions) were set with `with_*` +methods that copied the remaining twelve fields by hand. + +The new `AggregateExecBuilder` names every argument and defaults the `FILTER` +expressions to "no filter": + +```rust,ignore +let exec = AggregateExec::builder(AggregateMode::Single, input) + .with_group_by(group_by) + .with_aggr_exprs(aggr_exprs) + .with_limit_options(LimitOptions::new(10)) + .build()?; +``` + +An existing node is rewritten with `AggregateExec::to_builder`, which carries +over the output schema and plan properties of the original node, so a rewrite +cannot rename output fields and costs no more than the methods it replaces: + +```rust,ignore +let with_limit = exec + .to_builder() + .with_limit_options(LimitOptions::new_with_order(10, true)) + .build()?; +``` + +Building and rewriting an `AggregateExec` is how DataFusion's own physical +optimizer rules work, not a public API. The builder, `AggregateExec::builder`, +`AggregateExec::to_builder` and the methods it replaces are therefore all +`#[doc(hidden)]`, and may change without notice. + +This is a refactor: the node a builder produces is the node the method it +replaces produced. `build` returns a `Result` so that an `AggregateExec` can be +checked in the one place it is now built, but it currently checks only what +`try_new` already checked, that the aggregate and `FILTER` expressions have the +same length. + +**Who is affected:** + +- Callers of `AggregateExec::with_limit_options`, + `AggregateExec::with_new_limit_options` and + `AggregateExec::with_new_aggr_exprs`, which are deprecated and hidden. +- Callers of `AggregateExec::limit_options`, which is hidden along with the rest + of this API. It is not deprecated: it has no replacement, and reading the + limit of an aggregate is safe. + +**Migration guide:** + +- Replace `agg.with_limit_options(opts)` and `agg.with_new_limit_options(opts)` + with `agg.to_builder().with_limit_options(opts).build()?`. +- Replace `agg.with_new_aggr_exprs(exprs)` with + `agg.to_builder().with_aggr_exprs(exprs).build()?`.