diff --git a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs index 7097520061e63..ec416e06a5142 100644 --- a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs +++ b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs @@ -22,21 +22,55 @@ use std::sync::Arc; use datafusion_common::error::Result; use datafusion_physical_plan::ExecutionPlan; -use datafusion_physical_plan::aggregates::{ - AggregateExec, AggregateKind, AggregateMode, PhysicalGroupBy, -}; +use datafusion_physical_plan::aggregates::AggregateExec; use crate::PhysicalOptimizerRule; use datafusion_common::config::ConfigOptions; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_physical_expr::aggregate::AggregateFunctionExpr; -use datafusion_physical_expr::{PhysicalExpr, physical_exprs_equal}; +use datafusion_common::tree_node::{TransformedResult, TreeNode}; -/// CombinePartialFinalAggregate optimizer rule combines the adjacent Partial and Final AggregateExecs -/// into a Single AggregateExec if their grouping exprs and aggregate exprs equal. +/// Combines adjacent partial and final `AggregateExec`s into a single aggregate +/// when their grouping and aggregate expressions match. +/// +/// ```txt +/// Before: +/// AggregateExec(mode=Final) <-- current plan +/// AggregateExec(mode=Partial) +/// DataSourceExec +/// +/// After: +/// AggregateExec(mode=Single) <-- current plan +/// DataSourceExec +/// ``` /// /// This rule should be applied after the `EnsureRequirements` rule (which /// handles both distribution and sorting enforcement). +/// +/// # Background +/// +/// The relevant optimization steps, in order, are: +/// +/// 1. Initial physical planning: creates a partial/final pair so aggregation can +/// run in parallel. For grouped, partitioned aggregation: +/// +/// ```txt +/// AggregateExec(mode=FinalPartitioned) +/// AggregateExec(mode=Partial) +/// DataSourceExec +/// ``` +/// +/// 2. `EnsureRequirements` rule: inserts a `RepartitionExec` between the stages when +/// needed to satisfy the final aggregate's distribution requirements: +/// +/// ```txt +/// AggregateExec(mode=FinalPartitioned) +/// RepartitionExec <-- inserted here +/// AggregateExec(mode=Partial) +/// DataSourceExec +/// ``` +/// +/// 3. This rule: If the input already satisfies these requirements, e.g. because it +/// is partitioned by the grouping keys, no repartition is needed. This rule +/// then combines the adjacent, compatible aggregate stages into one. #[derive(Default, Debug)] pub struct CombinePartialFinalAggregate {} @@ -53,94 +87,8 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate { plan: Arc, _config: &ConfigOptions, ) -> Result> { - plan.transform_down(|plan| { - // Check if the plan is AggregateExec - let Some(agg_exec) = plan.downcast_ref::() else { - return Ok(Transformed::no(plan)); - }; - - if !matches!( - agg_exec.mode(), - AggregateMode::Final | AggregateMode::FinalPartitioned - ) { - return Ok(Transformed::no(plan)); - } - - // Check if the input is AggregateExec - let Some(input_agg_exec) = agg_exec.input().downcast_ref::() - else { - return Ok(Transformed::no(plan)); - }; - - let transformed = if *input_agg_exec.mode() == AggregateMode::Partial - && can_combine( - ( - agg_exec.group_expr(), - agg_exec.aggr_expr(), - agg_exec.filter_expr(), - ), - ( - input_agg_exec.group_expr(), - input_agg_exec.aggr_expr(), - input_agg_exec.filter_expr(), - ), - ) { - let mode = if agg_exec.mode() == &AggregateMode::Final { - AggregateMode::Single - } else { - AggregateMode::SinglePartitioned - }; - let Ok(mut combined_agg) = 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(), - ) else { - return Ok(Transformed::no(plan)); - }; - - // Re-apply distinct limit optimization. - // - // The optimizer related to aggregates are (in order): - // - 1. Initial planning: always final/partial two stage - // - 2. `LimitedDistinctAggregation`: push limit into `AggregateExec`, - // and build `AggregateKind::DistinctLimit` - // - 3. `CombinePartialFinalAggregate`: the current optimization - // - // Here it restores previously applied optimization in step 2 - let combined_aggr_kind = match (&agg_exec.kind, &input_agg_exec.kind) { - (AggregateKind::General { .. }, AggregateKind::General { .. }) => { - combined_agg.kind - } - ( - AggregateKind::DistinctLimit { limit, .. }, - AggregateKind::DistinctLimit { - group_by, - limit: partial_limit, - }, - ) if limit == partial_limit => AggregateKind::DistinctLimit { - // The combined aggregate groups raw input, like the partial. - group_by: Arc::clone(group_by), - limit: *limit, - }, - // Step 2 may limit only the final aggregate, e.g. when its - // grouping differs from the partial's. Keep both stages. - _ => return Ok(Transformed::no(plan)), - }; - combined_agg.kind = combined_aggr_kind; - Some(Arc::new(combined_agg)) - } else { - None - }; - Ok(if let Some(transformed) = transformed { - Transformed::yes(transformed) - } else { - Transformed::no(plan) - }) - }) - .data() + plan.transform_down(AggregateExec::try_combine_partial_final) + .data() } fn name(&self) -> &str { @@ -152,42 +100,4 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate { } } -type GroupExprsRef<'a> = ( - &'a PhysicalGroupBy, - &'a [Arc], - &'a [Option>], -); - -fn can_combine(final_agg: GroupExprsRef, partial_agg: GroupExprsRef) -> bool { - let (final_group_by, final_aggr_expr, final_filter_expr) = final_agg; - let (input_group_by, input_aggr_expr, input_filter_expr) = partial_agg; - - // Compare output expressions of the partial, and input expressions of the final operator. - physical_exprs_equal( - &input_group_by.output_exprs(), - &final_group_by.input_exprs(), - ) && input_group_by.groups() == final_group_by.groups() - && input_group_by.null_expr().len() == final_group_by.null_expr().len() - && input_group_by - .null_expr() - .iter() - .zip(final_group_by.null_expr().iter()) - .all(|((lhs_expr, lhs_str), (rhs_expr, rhs_str))| { - lhs_expr.eq(rhs_expr) && lhs_str == rhs_str - }) - && final_aggr_expr.len() == input_aggr_expr.len() - && final_aggr_expr - .iter() - .zip(input_aggr_expr.iter()) - .all(|(final_expr, partial_expr)| final_expr.eq(partial_expr)) - && final_filter_expr.len() == input_filter_expr.len() - && final_filter_expr.iter().zip(input_filter_expr.iter()).all( - |(final_expr, partial_expr)| match (final_expr, partial_expr) { - (Some(l), Some(r)) => l.eq(r), - (None, None) => true, - _ => false, - }, - ) -} - // See tests in datafusion/core/tests/physical_optimizer diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 226ce3d271973..6012f484a431d 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -197,7 +197,7 @@ use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; use datafusion_physical_expr::{ - ConstExpr, EquivalenceProperties, physical_exprs_contains, + ConstExpr, EquivalenceProperties, physical_exprs_contains, physical_exprs_equal, }; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, fmt_sql}; use datafusion_physical_expr_common::sort_expr::{ @@ -850,13 +850,8 @@ impl LimitOptions { } /// Mutually exclusive aggregation implementations and their configuration. -/// -/// # Public Only for Internal Use: -/// `datafusion-physical-optimizer` inspects and combines aggregate kinds. -/// This enum is not part of the supported public API. -#[doc(hidden)] #[derive(Debug, Clone)] -pub enum AggregateKind { +enum AggregateKind { /// Ordinary aggregation, with no limit on the groups retained. General { group_by: Arc, @@ -891,13 +886,7 @@ pub struct AggregateExec { /// Aggregation mode (full, partial) mode: AggregateMode, /// Aggregation implementation and its configuration. - /// - /// # Public Only for Internal Use: - /// `datafusion-physical-optimizer` updates this when combining aggregates. - /// Changes must preserve the expressions, schema, and plan properties. - /// This field is not part of the supported public API. - #[doc(hidden)] - pub kind: AggregateKind, + kind: AggregateKind, /// Input plan, could be a partial aggregate or the input to the aggregate pub input: Arc, /// Schema after the aggregate is applied. Contains the group by columns followed by the @@ -1172,6 +1161,130 @@ impl AggregateExec { Some(Transformed::yes(self)) } + /// Combines adjacent partial and final aggregates in the given plan if they are + /// compatible. See [`CombinePartialFinalAggregate`] for the invariants required + /// by this transformation. + /// + /// If input plan is not aggregate, this function will short-circuit with + /// `Transformed::No`. + /// + /// ```txt + /// Before Plan: + /// AggregateExec(mode=Final) <-- current plan + /// AggregateExec(mode=Partial) + /// DataSourceExec + /// + /// After Plan: + /// AggregateExec(mode=Single) <-- current plan + /// DataSourceExec + /// ``` + /// + /// # Public Only for Internal Use: + /// + /// This is not a public API and is for internal use only in physical optimizer; + /// see [API policy] for details. + /// + /// [`CombinePartialFinalAggregate`]: https://docs.rs/datafusion/latest/datafusion/physical_optimizer/combine_partial_final_agg/struct.CombinePartialFinalAggregate.html + /// [API policy]: https://datafusion.apache.org/contributor-guide/api-health.html#internal-public-apis + #[doc(hidden)] + pub fn try_combine_partial_final( + plan: Arc, + ) -> Result>> { + // Step 1: Ensure the below plan shape + // + // AggregateExec(mode=final) <-- current plan + // AggregateExec(mode=partial) + + let Some(final_agg) = plan.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + let mode = match final_agg.mode { + AggregateMode::Final => AggregateMode::Single, + AggregateMode::FinalPartitioned => AggregateMode::SinglePartitioned, + _ => return Ok(Transformed::no(plan)), + }; + let Some(partial) = final_agg.input.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + if partial.mode != AggregateMode::Partial { + return Ok(Transformed::no(plan)); + } + + // Step 2: Check compatibility between partial and final `AggregateExec` + + let final_group_by = final_agg.group_expr(); + let input_group_by = partial.group_expr(); + let final_aggr_expr = final_agg.aggr_expr(); + let input_aggr_expr = partial.aggr_expr(); + let final_filter_expr = final_agg.filter_expr(); + let input_filter_expr = partial.filter_expr(); + + // Compare the partial's output expressions with the final's input expressions. + let can_combine = physical_exprs_equal( + &input_group_by.output_exprs(), + &final_group_by.input_exprs(), + ) && input_group_by.groups() == final_group_by.groups() + && input_group_by.null_expr().len() == final_group_by.null_expr().len() + && input_group_by + .null_expr() + .iter() + .zip(final_group_by.null_expr().iter()) + .all(|((lhs_expr, lhs_str), (rhs_expr, rhs_str))| { + lhs_expr.eq(rhs_expr) && lhs_str == rhs_str + }) + && final_aggr_expr.len() == input_aggr_expr.len() + && final_aggr_expr + .iter() + .zip(input_aggr_expr.iter()) + .all(|(final_expr, partial_expr)| final_expr.eq(partial_expr)) + && final_filter_expr.len() == input_filter_expr.len() + && final_filter_expr.iter().zip(input_filter_expr.iter()).all( + |(final_expr, partial_expr)| match (final_expr, partial_expr) { + (Some(l), Some(r)) => l.eq(r), + (None, None) => true, + _ => false, + }, + ); + if !can_combine { + return Ok(Transformed::no(plan)); + } + + let Ok(mut combined) = Self::try_new( + mode, + input_group_by.clone(), + input_aggr_expr.to_vec(), + input_filter_expr.to_vec(), + Arc::clone(partial.input()), + partial.input_schema(), + ) else { + return Ok(Transformed::no(plan)); + }; + + // LimitedDistinctAggregation runs before this optimization. Preserve its + // limit while using the partial's grouping expressions over raw input. + // + // The following logic recognizes all possible patterns produced by `LimitedDistinctAggregation` + // See its comment for details. + combined.kind = match (&final_agg.kind, &partial.kind) { + (AggregateKind::General { .. }, AggregateKind::General { .. }) => { + combined.kind + } + ( + AggregateKind::DistinctLimit { limit, .. }, + AggregateKind::DistinctLimit { + group_by, + limit: partial_limit, + }, + ) if limit == partial_limit => AggregateKind::DistinctLimit { + group_by: Arc::clone(group_by), + limit: *limit, + }, + // Not possible in the default optimizer pipeline + _ => return Ok(Transformed::no(plan)), + }; + Ok(Transformed::yes(Arc::new(combined))) + } + /// Clone with replacement aggregate expressions. /// /// The caller must preserve the schema, filters, and ordering requirements. @@ -1689,7 +1802,7 @@ impl AggregateExec { .downcast_ref::() .is_some_and(|partial| { partial.mode == AggregateMode::Partial - && datafusion_physical_expr::physical_exprs_equal( + && physical_exprs_equal( &partial.output_group_expr(), &self.group_by().input_exprs(), )