You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "2010YOUY01 (via GitHub)" <gi...@apache.org> on 2026/09/21 04:22:10 UTC
[PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
2010YOUY01 opened a new pull request, #25564:
URL: https://github.com/apache/datafusion/pull/25564
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Part of https://github.com/apache/datafusion/issues/25563
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
Please explain the problem you are trying to solve in terms of the user-visible
behavior, rather than the implementation.
For example, "The code in `foo.rs` doesn't handle nulls" is a symptom of the
implementation. "COUNT(DISTINCT) returns wrong results when the column contains
nulls" is the user-visible problem.
-->
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?
<!--
There is no need to duplicate the description in the issue here, but it is sometimes worth providing a summary of the individual changes 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?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
4. Serve as another way to document the expected behavior of the code
Briefly describe how this PR is tested, and point to the specific tests you added. For example: 'This new feature is covered by the `sqllogictest` cases added in `foo.slt`'.
If this PR does not add tests, explain why. For example, if the change is already covered by existing tests, please mention it.
You should also check the `codecov` bot reply on this PR to confirm the changed code is exercised.
-->
Existing tests.
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
If there are any breaking changes to public APIs, please add the `api change` label.
-->
Yes, public fields inside `AggregateExec` are changed. (probably we should hide them in some way, since they're implementation details)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4059248057
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -20,7 +20,7 @@
use std::sync::Arc;
-use datafusion_physical_plan::aggregates::{AggregateExec, LimitOptions};
+use datafusion_physical_plan::aggregates::AggregateExec;
Review Comment:
Structural change 1
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -185,13 +185,13 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use arrow_schema::FieldRef;
use datafusion_common::stats::Precision;
-use datafusion_common::tree_node::TreeNodeRecursion;
+use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};
Review Comment:
This file includes structural change 2.
Remaning diffs are all mechanical access method change
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4068407107
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,49 +900,74 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to stop after enough distinct grouping keys have been accumulated.
+ ///
+ /// The caller must establish that discarding other groups is legal, for
+ /// example beneath `SELECT DISTINCT k FROM t LIMIT 10`. The parent limit
+ /// remains responsible for enforcing the exact number of rows.
+ ///
+ /// # Trigger conditions
+ ///
+ /// - Grouping with no aggregate expressions or aggregate filters.
+ /// - Ordering eligible under [`Self::is_unordered_unfiltered_group_by_distinct`],
+ /// including no existing Top-K direction.
+ /// - A positive limit that is tighter than any existing limit.
+ ///
+ /// # Consistency
+ ///
+ /// This is a safe, atomic optimization: it preserves the grouping keys and
+ /// leaves the aggregate in a consistent state. Inapplicable requests are
+ /// no-ops; the existing configuration is returned unchanged.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Result<Transformed<Self>> {
+ if limit == 0
Review Comment:
Updated in [2cf27d8](https://github.com/apache/datafusion/pull/25564/commits/2cf27d8e1022d765d46548ac4260ce1b16eb9ef1)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#issuecomment-5758346794
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/25564?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `82.74112%` with `34 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 82.41%. Comparing base ([`1e09a2a`](https://app.codecov.io/gh/apache/datafusion/commit/1e09a2a3d76c03f469e80bed712a6bada45e77c8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`d55fa20`](https://app.codecov.io/gh/apache/datafusion/commit/d55fa20c5de3a64fedde087afdcf25ba5aae84f3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 1 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/25564?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [datafusion/physical-plan/src/aggregates/mod.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fmod.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL21vZC5ycw==) | 83.33% | [22 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...ical-optimizer/src/limited\_distinct\_aggregation.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-optimizer%2Fsrc%2Flimited_distinct_aggregation.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1vcHRpbWl6ZXIvc3JjL2xpbWl0ZWRfZGlzdGluY3RfYWdncmVnYXRpb24ucnM=) | 77.77% | [0 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...plan/src/aggregates/aggregate\_hash\_table/common.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Faggregate_hash_table%2Fcommon.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL2FnZ3JlZ2F0ZV9oYXNoX3RhYmxlL2NvbW1vbi5ycw==) | 80.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [.../aggregates/aggregate\_hash\_table/common\_ordered.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Faggregate_hash_table%2Fcommon_ordered.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL2FnZ3JlZ2F0ZV9oYXNoX3RhYmxlL2NvbW1vbl9vcmRlcmVkLnJz) | 80.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...n/physical-plan/src/aggregates/aggregate\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Faggregate_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL2FnZ3JlZ2F0ZV9zdHJlYW0ucnM=) | 80.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...hysical-plan/src/aggregates/grouped\_topk\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fgrouped_topk_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL2dyb3VwZWRfdG9wa19zdHJlYW0ucnM=) | 66.66% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...fusion/physical-plan/src/aggregates/hash\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fhash_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL2hhc2hfc3RyZWFtLnJz) | 50.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...ysical-plan/src/aggregates/ordered\_final\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fordered_final_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL29yZGVyZWRfZmluYWxfc3RyZWFtLnJz) | 0.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...sical-plan/src/aggregates/ordered\_single\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fordered_single_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL29yZGVyZWRfc2luZ2xlX3N0cmVhbS5ycw==) | 75.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...sion/physical-plan/src/aggregates/single\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fsingle_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL3NpbmdsZV9zdHJlYW0ucnM=) | 75.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25564?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #25564 +/- ##
==========================================
- Coverage 82.42% 82.41% -0.01%
==========================================
Files 1138 1138
Lines 435429 435486 +57
Branches 435429 435486 +57
==========================================
+ Hits 358889 358914 +25
- Misses 54839 54866 +27
- Partials 21701 21706 +5
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/25564?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4070382283
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,49 +900,74 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to stop after enough distinct grouping keys have been accumulated.
+ ///
+ /// The caller must establish that discarding other groups is legal, for
+ /// example beneath `SELECT DISTINCT k FROM t LIMIT 10`. The parent limit
+ /// remains responsible for enforcing the exact number of rows.
+ ///
+ /// # Trigger conditions
+ ///
+ /// - Grouping with no aggregate expressions or aggregate filters.
+ /// - Ordering eligible under [`Self::is_unordered_unfiltered_group_by_distinct`],
+ /// including no existing Top-K direction.
+ /// - A positive limit that is tighter than any existing limit.
+ ///
+ /// # Consistency
+ ///
+ /// This is a safe, atomic optimization: it preserves the grouping keys and
+ /// leaves the aggregate in a consistent state. Inapplicable requests are
+ /// no-ops; the existing configuration is returned unchanged.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Result<Transformed<Self>> {
+ if limit == 0
+ || !self.is_unordered_unfiltered_group_by_distinct()
+ || self
+ .limit_options()
+ .is_some_and(|existing| existing.limit <= limit)
+ {
+ return Ok(Transformed::no(self));
+ }
+ self.kind = AggregateKind::DistinctLimit {
+ group_by: Arc::clone(self.group_by()),
+ limit,
+ };
+ Ok(Transformed::yes(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.
pub fn with_new_aggr_exprs(
&self,
aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>,
) -> Self {
- Self {
- aggr_expr: aggr_expr.into(),
- // clone the rest of the fields
- required_input_ordering: self.required_input_ordering.clone(),
- metrics: ExecutionPlanMetricsSet::new(),
- input_order_mode: self.input_order_mode.clone(),
- cache: Arc::clone(&self.cache),
- mode: self.mode,
- group_by: Arc::clone(&self.group_by),
- filter_expr: Arc::clone(&self.filter_expr),
- limit_options: self.limit_options,
- input: Arc::clone(&self.input),
- schema: Arc::clone(&self.schema),
- input_schema: Arc::clone(&self.input_schema),
- dynamic_filter: self.dynamic_filter.clone(),
+ let aggr_expr = aggr_expr.into();
+ let mut new = self.clone();
+ match &mut new.kind {
+ AggregateKind::General { aggr_expr: old, .. } => *old = aggr_expr,
+ AggregateKind::DistinctLimit { .. } if aggr_expr.is_empty() => {}
+ AggregateKind::DistinctLimit { group_by, .. } => {
+ // 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(),
Review Comment:
> 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4068406255
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -47,16 +47,15 @@ impl LimitedDistinctAggregation {
fn transform_agg(
aggr: &AggregateExec,
limit: usize,
- ) -> Option<Arc<dyn ExecutionPlan>> {
+ ) -> Option<Transformed<Arc<dyn ExecutionPlan>>> {
// rules for transforming this Aggregate are held in this method
if !aggr.is_unordered_unfiltered_group_by_distinct() {
return None;
}
- // 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)));
-
- Some(Arc::new(new_aggr))
+ let new_aggr = aggr.clone().try_optimize_distinct_soft_limit(limit).ok()?;
Review Comment:
- `Option` instead of `Result`: I agree, updated in [2cf27d8](https://github.com/apache/datafusion/pull/25564/commits/2cf27d8e1022d765d46548ac4260ce1b16eb9ef1)
- Avoiding deep clone: It seems impossible due to the optimizer interface restriction?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4065957123
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -47,16 +47,15 @@ impl LimitedDistinctAggregation {
fn transform_agg(
aggr: &AggregateExec,
limit: usize,
- ) -> Option<Arc<dyn ExecutionPlan>> {
+ ) -> Option<Transformed<Arc<dyn ExecutionPlan>>> {
// rules for transforming this Aggregate are held in this method
if !aggr.is_unordered_unfiltered_group_by_distinct() {
return None;
}
- // 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)));
-
- Some(Arc::new(new_aggr))
+ let new_aggr = aggr.clone().try_optimize_distinct_soft_limit(limit).ok()?;
Review Comment:
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)
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -119,14 +118,18 @@ impl LimitedDistinctAggregation {
Some(new_aggr) => {
match_aggr = plan;
found_match_aggr = true;
- return Ok(Transformed::yes(new_aggr));
+ return Ok(new_aggr);
}
}
}
rewrite_applicable = false;
Ok(Transformed::no(plan))
};
- let child = child.to_owned().transform_down(closure).data().ok()?;
Review Comment:
though here is an example of discarding an error in existing code too 😬
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -845,35 +845,30 @@ impl LimitOptions {
}
}
+/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+#[derive(Debug, Clone)]
+enum AggregateKind {
+ /// Ordinary aggregation, including the existing Top-K configuration.
+ General {
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Arc<[Arc<AggregateFunctionExpr>]>,
+ filter_expr: Arc<[Option<Arc<dyn PhysicalExpr>>]>,
+ limit_options: Option<LimitOptions>,
+ },
+ /// `SELECT DISTINCT k FROM t LIMIT n`: eligible streams may stop after n groups.
+ /// Other streams consume all input; the parent LIMIT enforces the row count.
+ DistinctLimit {
+ group_by: Arc<PhysicalGroupBy>,
+ limit: usize,
+ },
+}
+
/// Hash aggregate execution plan
#[derive(Debug, Clone)]
pub struct AggregateExec {
/// Aggregation mode (full, partial)
mode: AggregateMode,
- /// Group by expressions
- /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance.
- group_by: Arc<PhysicalGroupBy>,
- /// Aggregate expressions
- /// The same reason to [`Arc`] it as for [`Self::group_by`].
- aggr_expr: Arc<[Arc<AggregateFunctionExpr>]>,
- /// FILTER (WHERE clause) expression for each aggregate expression
- /// The same reason to [`Arc`] it as for [`Self::group_by`].
- filter_expr: Arc<[Option<Arc<dyn PhysicalExpr>>]>,
- /// Soft limit for best-effort, limit-based optimizations.
- ///
- /// `AggregateExec` has multiple stream implementations and not all of them
- /// support the optimization, so this is only a hint. The downstream limit
- /// operator enforces the exact row count either way.
- ///
- /// Supported by:
- /// - [`StreamType::GroupedPriorityQueue`]: retains only the best `limit`
- /// groups per partition (this stream is selected only when a limit is set)
- /// - [`StreamType::SingleHash`], [`StreamType::PartialHash`], [`StreamType::FinalHash`]
- /// and the legacy [`StreamType::GroupedHash`]: stop reading input once `limit` groups
- /// have been accumulated
- ///
- /// The remaining streams consume all input.
- limit_options: Option<LimitOptions>,
+ kind: AggregateKind,
Review Comment:
👍
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,49 +900,74 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to stop after enough distinct grouping keys have been accumulated.
+ ///
+ /// The caller must establish that discarding other groups is legal, for
+ /// example beneath `SELECT DISTINCT k FROM t LIMIT 10`. The parent limit
+ /// remains responsible for enforcing the exact number of rows.
+ ///
+ /// # Trigger conditions
+ ///
+ /// - Grouping with no aggregate expressions or aggregate filters.
+ /// - Ordering eligible under [`Self::is_unordered_unfiltered_group_by_distinct`],
+ /// including no existing Top-K direction.
+ /// - A positive limit that is tighter than any existing limit.
+ ///
+ /// # Consistency
+ ///
+ /// This is a safe, atomic optimization: it preserves the grouping keys and
+ /// leaves the aggregate in a consistent state. Inapplicable requests are
+ /// no-ops; the existing configuration is returned unchanged.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Result<Transformed<Self>> {
+ if limit == 0
+ || !self.is_unordered_unfiltered_group_by_distinct()
+ || self
+ .limit_options()
+ .is_some_and(|existing| existing.limit <= limit)
+ {
+ return Ok(Transformed::no(self));
+ }
+ self.kind = AggregateKind::DistinctLimit {
+ group_by: Arc::clone(self.group_by()),
+ limit,
+ };
+ Ok(Transformed::yes(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.
pub fn with_new_aggr_exprs(
&self,
aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>,
) -> Self {
- Self {
- aggr_expr: aggr_expr.into(),
- // clone the rest of the fields
- required_input_ordering: self.required_input_ordering.clone(),
- metrics: ExecutionPlanMetricsSet::new(),
- input_order_mode: self.input_order_mode.clone(),
- cache: Arc::clone(&self.cache),
- mode: self.mode,
- group_by: Arc::clone(&self.group_by),
- filter_expr: Arc::clone(&self.filter_expr),
- limit_options: self.limit_options,
- input: Arc::clone(&self.input),
- schema: Arc::clone(&self.schema),
- input_schema: Arc::clone(&self.input_schema),
- dynamic_filter: self.dynamic_filter.clone(),
+ let aggr_expr = aggr_expr.into();
+ let mut new = self.clone();
+ match &mut new.kind {
+ AggregateKind::General { aggr_expr: old, .. } => *old = aggr_expr,
+ AggregateKind::DistinctLimit { .. } if aggr_expr.is_empty() => {}
+ AggregateKind::DistinctLimit { group_by, .. } => {
+ // 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(),
Review Comment:
Why are the filter expressions reset?
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,49 +900,74 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to stop after enough distinct grouping keys have been accumulated.
+ ///
+ /// The caller must establish that discarding other groups is legal, for
+ /// example beneath `SELECT DISTINCT k FROM t LIMIT 10`. The parent limit
+ /// remains responsible for enforcing the exact number of rows.
+ ///
+ /// # Trigger conditions
+ ///
+ /// - Grouping with no aggregate expressions or aggregate filters.
+ /// - Ordering eligible under [`Self::is_unordered_unfiltered_group_by_distinct`],
+ /// including no existing Top-K direction.
+ /// - A positive limit that is tighter than any existing limit.
+ ///
+ /// # Consistency
+ ///
+ /// This is a safe, atomic optimization: it preserves the grouping keys and
+ /// leaves the aggregate in a consistent state. Inapplicable requests are
+ /// no-ops; the existing configuration is returned unchanged.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Result<Transformed<Self>> {
+ if limit == 0
Review Comment:
This seems to always return "Ok" -- why is it returning Result?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 merged PR #25564:
URL: https://github.com/apache/datafusion/pull/25564
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#issuecomment-5787286755
Thanks @2010YOUY01 this is great!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4068256404
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,49 +900,74 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to stop after enough distinct grouping keys have been accumulated.
+ ///
+ /// The caller must establish that discarding other groups is legal, for
+ /// example beneath `SELECT DISTINCT k FROM t LIMIT 10`. The parent limit
+ /// remains responsible for enforcing the exact number of rows.
+ ///
+ /// # Trigger conditions
+ ///
+ /// - Grouping with no aggregate expressions or aggregate filters.
+ /// - Ordering eligible under [`Self::is_unordered_unfiltered_group_by_distinct`],
+ /// including no existing Top-K direction.
+ /// - A positive limit that is tighter than any existing limit.
+ ///
+ /// # Consistency
+ ///
+ /// This is a safe, atomic optimization: it preserves the grouping keys and
+ /// leaves the aggregate in a consistent state. Inapplicable requests are
+ /// no-ops; the existing configuration is returned unchanged.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Result<Transformed<Self>> {
+ if limit == 0
+ || !self.is_unordered_unfiltered_group_by_distinct()
+ || self
+ .limit_options()
+ .is_some_and(|existing| existing.limit <= limit)
+ {
+ return Ok(Transformed::no(self));
+ }
+ self.kind = AggregateKind::DistinctLimit {
+ group_by: Arc::clone(self.group_by()),
+ limit,
+ };
+ Ok(Transformed::yes(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.
pub fn with_new_aggr_exprs(
&self,
aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>,
) -> Self {
- Self {
- aggr_expr: aggr_expr.into(),
- // clone the rest of the fields
- required_input_ordering: self.required_input_ordering.clone(),
- metrics: ExecutionPlanMetricsSet::new(),
- input_order_mode: self.input_order_mode.clone(),
- cache: Arc::clone(&self.cache),
- mode: self.mode,
- group_by: Arc::clone(&self.group_by),
- filter_expr: Arc::clone(&self.filter_expr),
- limit_options: self.limit_options,
- input: Arc::clone(&self.input),
- schema: Arc::clone(&self.schema),
- input_schema: Arc::clone(&self.input_schema),
- dynamic_filter: self.dynamic_filter.clone(),
+ let aggr_expr = aggr_expr.into();
+ let mut new = self.clone();
+ match &mut new.kind {
+ AggregateKind::General { aggr_expr: old, .. } => *old = aggr_expr,
+ AggregateKind::DistinctLimit { .. } if aggr_expr.is_empty() => {}
+ AggregateKind::DistinctLimit { group_by, .. } => {
+ // 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(),
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org