Skip to content

refactor: Simplify LimitedDistinctAggregation optimizer rule - #26069

Open
2010YOUY01 wants to merge 3 commits into
apache:mainfrom
2010YOUY01:cleanup-limit-aggr
Open

2010YOUY01 wants to merge 3 commits into
apache:mainfrom
2010YOUY01:cleanup-limit-aggr

Conversation

@2010YOUY01

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

This is primarily a simplifying refactor, with an optimizer regression fix piggy-backed on top (an inefficient plan shape, not a correctness issue):

The optimizer fix itself is only a one-line change. If this simplification is not desirable, I'll open a replacement PR containing just the direct fix.

Rationale for this change

Part 1: Optimizer fix

See sqllogictest diff for the reproducer
I'll mark the fix in comments

Part 2: Optimizer simplification

During different stages in physical optimization, a logical aggregation can be either

AggregateExec(mode=final)
--AggregateExec(mode=partial)

AggregateExec(mode=single)

AggregateExec(mode=final)
--RepartitionExec
----AggregateExec(mode=partial)

So the existing implementation is using a nested dfs (the closure has another dfs inside, to match non-consecutive aggregates) try to match all of the 3 cases.

However, after looking at related optimizer rules:

  • initial physical planning
  • CombinePartialFinalAggregate
  • EnsureDistributions
  • and this rule

We can find only case 1 is possible, so the implementation can be simplified into a naive pattern matching

Find the exact below shape, and try push limits

LimitExec
--AggregateExec(mode=final)
----AggregateExec(mode=partial)

What changes are included in this PR?

  • simplify an optimizer rule according to the above rationale
  • fix a small bug

What is the testing strategy for this PR?

  • for bug fix, changes in slt
  • To ensure refactor correctness: no e2e test need to be changed

Are there any user-facing changes?

@github-actions github-actions Bot added optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) labels Oct 6, 2026
if let Some(aggr) = plan.downcast_ref::<AggregateExec>() {
if found_match_aggr
&& let Some(parent_aggr) = match_aggr.downcast_ref::<AggregateExec>()
&& !parent_aggr.group_expr().eq(aggr.group_expr())

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bug fix: missing a as_final to adapt projection difference

(
AggregateMode::Final | AggregateMode::FinalPartitioned,
AggregateMode::Partial,
) if final_agg.group_expr() == &partial_agg.group_expr().as_final() => {}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👉🏼

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 88.57143% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.69%. Comparing base (c3ef346) to head (784d746).
⚠️ Report is 16 commits behind head on main.

Files with missing lines Patch % Lines
...ical-optimizer/src/limited_distinct_aggregation.rs 88.57% 3 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #26069      +/-   ##
==========================================
+ Coverage   82.66%   82.69%   +0.02%     
==========================================
  Files        1147     1147              
  Lines      446357   447312     +955     
  Branches   446357   447312     +955     
==========================================
+ Hits       368971   369892     +921     
- Misses      54997    55008      +11     
- Partials    22389    22412      +23     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants