Skip to content

refactor: Implement optimizer hook for combining partial/final aggregate in AggregateExec #25962

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alamb merged 5 commits into from
Oct 24, 2026
Merged

refactor: Implement optimizer hook for combining partial/final aggregate in AggregateExec #25962

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
833a6c0
refactor: Implement optimizer hook for combining partial/final aggreg…
2010YOUY01 Oct 2, 2026
928c2b5
Merge remote-tracking branch 'upstream/main' into simplify-aggr-combine
2010YOUY01 Oct 24, 2026
f94db31
review: improve comments
2010YOUY01 Oct 24, 2026
f5b19f6
Merge branch 'main' into simplify-aggr-combine
alamb Oct 24, 2026
e6c53ff
fix
alamb Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 44 additions & 134 deletions datafusion/physical-optimizer/src/combine_partial_final_agg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}

Expand All @@ -53,94 +87,8 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
plan: Arc<dyn ExecutionPlan>,
_config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
plan.transform_down(|plan| {
// Check if the plan is AggregateExec
let Some(agg_exec) = plan.downcast_ref::<AggregateExec>() 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::<AggregateExec>()
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 {
Expand All @@ -152,42 +100,4 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
}
}

type GroupExprsRef<'a> = (
&'a PhysicalGroupBy,
&'a [Arc<AggregateFunctionExpr>],
&'a [Option<Arc<dyn PhysicalExpr>>],
);

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
143 changes: 128 additions & 15 deletions datafusion/physical-plan/src/aggregates/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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<PhysicalGroupBy>,
Expand Down Expand Up @@ -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<dyn ExecutionPlan>,
/// Schema after the aggregate is applied. Contains the group by columns followed by the
Expand Down Expand Up @@ -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<dyn ExecutionPlan>,
) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
// Step 1: Ensure the below plan shape
//
// AggregateExec(mode=final) <-- current plan
// AggregateExec(mode=partial)

let Some(final_agg) = plan.downcast_ref::<Self>() 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::<Self>() 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.
Expand Down Expand Up @@ -1689,7 +1802,7 @@ impl AggregateExec {
.downcast_ref::<AggregateExec>()
.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(),
)
Expand Down
Loading