refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) by 2010YOUY01 · Pull Request #25564 · apache/datafusion

2010YOUY01 · 2026-09-21T04:22:07Z

Which issue does this PR close?

Part of Make AggregateExec state modeling and public updates safe #25563

Rationale for this change

See issue for the idea, this PR is the step 1 in the issue.

For backgrounds for distinct limit optimization (this PR targets), see

datafusion/datafusion/physical-plan/src/aggregates/hash_stream.rs

Line 77 in 1e09a2a

/// # Optimization: DISTINCT LIMIT Soft Limit

What changes are included in this PR?

Refactor AggregateExec with enum state
Add safe API for this optimization, and use it in the optimizer.
(unsafe one will be able to be deleted after step2, now it's reused by other path)

I'll mark the structural changes (only 2 files)

Most diff is mechanical changes like aggregate_exec.aggr_expr -> aggregate_exec.aggr_expr(), caused by the refactor.

What is the testing strategy for this PR?

Existing tests.

Are there any user-facing changes?

no

2010YOUY01 · 2026-09-21T04:22:46Z

 use std::sync::Arc;

-use datafusion_physical_plan::aggregates::{AggregateExec, LimitOptions};
+use datafusion_physical_plan::aggregates::AggregateExec;


Structural change 1

2010YOUY01 · 2026-09-21T04:23:38Z

 use arrow_schema::FieldRef;
 use datafusion_common::stats::Precision;
-use datafusion_common::tree_node::TreeNodeRecursion;
+use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};


This file includes structural change 2.

Remaning diffs are all mechanical access method change

codecov-commenter · 2026-09-21T09:32:36Z

Codecov Report

❌ Patch coverage is 83.75635% with 32 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.42%. Comparing base (2424dae) to head (e01af79).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/aggregates/mod.rs	82.99%	23 Missing and 2 partials ⚠️
datafusion/physical-plan/src/aggregates/spill.rs	33.33%	0 Missing and 2 partials ⚠️
...ical-optimizer/src/limited_distinct_aggregation.rs	88.88%	0 Missing and 1 partial ⚠️
...plan/src/aggregates/aggregate_hash_table/common.rs	80.00%	0 Missing and 1 partial ⚠️
.../aggregates/aggregate_hash_table/common_ordered.rs	80.00%	0 Missing and 1 partial ⚠️
...n/physical-plan/src/aggregates/aggregate_stream.rs	80.00%	0 Missing and 1 partial ⚠️
...hysical-plan/src/aggregates/grouped_topk_stream.rs	66.66%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25564      +/-   ##
==========================================
- Coverage   82.42%   82.42%   -0.01%     
==========================================
  Files        1139     1139              
  Lines      435410   435467      +57     
  Branches   435410   435467      +57     
==========================================
+ Hits       358884   358922      +38     
- Misses      54822    54838      +16     
- Partials    21704    21707       +3

☔ 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.

alamb

Thanks @2010YOUY01 this makes sense to me. I had a few concerns about some extra cloning and filter reset, but the basic pattern looks great to me. I like where this is heading

FYI @rluvaton @Rachelint and @jayzhan211 who I think are also working in this area

alamb · 2026-09-21T20:06:35Z

-        let new_aggr = aggr.with_new_limit_options(Some(LimitOptions::new(limit)));
-
-        Some(Arc::new(new_aggr))
+        let new_aggr = aggr.clone().try_optimize_distinct_soft_limit(limit).ok()?;


The pattern of ignoring Errs I think is bad in DataFusion because each Err contains an allocated string (so this is inefficient).

Perhaps we could just change try_optimize_distinct_soft_limit to return an Option rather than Result 🤔

Also, it would be great to avoid the need to clone the input(as this is a deep clone, not just a an arc clone)

Option instead of Result: I agree, updated in 2cf27d8

Avoiding deep clone: It seems impossible due to the optimizer interface restriction?

alamb · 2026-09-21T20:07:12Z

            rewrite_applicable = false;
            Ok(Transformed::no(plan))
        };
-        let child = child.to_owned().transform_down(closure).data().ok()?;


though here is an example of discarding an error in existing code too 😬

alamb · 2026-09-21T20:07:52Z

-    ///
-    /// The remaining streams consume all input.
-    limit_options: Option<LimitOptions>,
+    kind: AggregateKind,


alamb · 2026-09-21T20:09:38Z

+        mut self,
+        limit: usize,
+    ) -> Result<Transformed<Self>> {
+        if limit == 0


This seems to always return "Ok" -- why is it returning Result?

Updated in 2cf27d8

alamb · 2026-09-21T20:10:45Z

+                // An accumulator rewrite cannot inherit DISTINCT's early stop.
+                new.kind = AggregateKind::General {
+                    group_by: Arc::clone(group_by),
+                    filter_expr: vec![None; aggr_expr.len()].into(),


Why are the filter expressions reset?

The previous AggregateKind is DistinctLimit, which is for the query pattern select distinct k from t limit 10 after it gets optimized into an aggregate, so it must have no aggregate functions or filters.

This transformation converts DistinctLimit -> General, so there is no filter to init.

However, the deeper question seems to be: is it safe? I'm not 100% sure. Before this refactor, this conversion happened implicitly and was easy to miss. With the typed state, it now shows up obviously.

I'll add more comment for it.

However, the deeper question seems to be: is it safe?

Agreed

With the typed state, it now shows up obviously.

That is a neat result

adriangb · 2026-09-23T01:16:10Z

Thanks @2010YOUY01 this is great!

… (part 1: distinct limit) (apache#25564) ## Which issue does this PR close?  - Part of apache#25563 ## Rationale for this change  See issue for the idea, this PR is the step 1 in the issue. For backgrounds for distinct limit optimization (this PR targets), see https://github.com/apache/datafusion/blob/1e09a2a3d76c03f469e80bed712a6bada45e77c8/datafusion/physical-plan/src/aggregates/hash_stream.rs#L77 ## What changes are included in this PR?  1. Refactor AggregateExec with enum state 2. Add safe API for this optimization, and use it in the optimizer. (unsafe one will be able to be deleted after step2, now it's reused by other path) I'll mark the structural changes (only 2 files) Most diff is mechanical changes like `aggregate_exec.aggr_expr` -> `aggregate_exec.aggr_expr()`, caused by the refactor. ## What is the testing strategy for this PR?  Existing tests. ## Are there any user-facing changes?  no

refactor: clean up aggregate planning for 'distinct with limit'

20ae528

github-actions Bot added optimizer Optimizer rules physical-plan Changes to the physical-plan crate labels Sep 21, 2026

2010YOUY01 commented Sep 21, 2026

View reviewed changes

2010YOUY01 mentioned this pull request Sep 21, 2026

refactor: build and rewrite AggregateExec through an AggregateExecBuilder #25376

Closed

2010YOUY01 requested review from adriangb and alamb September 21, 2026 04:39

Merge branch 'main' into distinct-aggregate

d55fa20

alamb approved these changes Sep 21, 2026

View reviewed changes

2010YOUY01 added 3 commits September 22, 2026 11:50

review: let try_optimize_.. return Option instead of Result

2cf27d8

Merge remote-tracking branch 'upstream/main' into distinct-aggregate

8488d72

review: more comments

e01af79

2010YOUY01 added this pull request to the merge queue Sep 23, 2026

Merged via the queue into apache:main with commit ebdb657 Sep 23, 2026
41 checks passed

2010YOUY01 deleted the distinct-aggregate branch September 23, 2026 00:56

nuno-faria added the v56.0.0 label Oct 3, 2026

Conversation

2010YOUY01 commented Sep 21, 2026 • edited by alamb Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

adriangb commented Sep 23, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

2010YOUY01 commented Sep 21, 2026 •

edited by alamb

Loading

codecov-commenter commented Sep 21, 2026 •

edited

Loading