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/24 08:05:56 UTC
[PR] refactor: Make AggregateExec state modeling and public updates safe (part2: topk) [datafusion]
2010YOUY01 opened a new pull request, #25696:
URL: https://github.com/apache/datafusion/pull/25696
## 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.
-->
- Closes #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
(TODO)
## 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. Represent topk path in `AggregateExec` with enum state
2. Deprecate unsafe APIs on `AggregateExec`, and replace them with a atomic and safe one:
- Deprecated: `with_new_limit_options` and `with_limit_options`
- Safe and atomic alternative: `try_optimize_topk`
## 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
3. 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 (adapted the optimizer tests to use the new safe API for the same test goal)
## Are there any user-facing changes?
Deprecating error-prone APIs on AggregateExec
<!--
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.
-->
--
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 (part2: topk) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4154454893
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1383,7 +1383,7 @@ impl AggregateExec {
/// Set a legacy limit hint. Unsupported requests leave ordinary aggregation.
#[deprecated(
since = "56.0.0",
- note = "This API is intended for internal use only and was inadvertently made public. Do not use this API."
+ note = "The replacement APIs are [`try_optimize_topk`](Self::try_optimize_topk) and [`try_optimize_distinct_soft_limit`](Self::try_optimize_distinct_soft_limit). They still require specific plan-shape invariants to be used correctly, so they remain internal for now. If you have a use case that requires these APIs to be public, please open an issue in DataFusion."
Review Comment:
👍
--
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 (part2: topk) [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#issuecomment-5810406070
<!-- semver-check-comment -->
Thank you for opening this pull request!
Reviewer note: [cargo-semver-checks](https://github.com/obi1kenobi/cargo-semver-checks) reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).
<details>
<summary>Details</summary>
```
Cloning apache/main
Building datafusion v55.1.0 (current)
Built [ 50.738s] (current)
Parsing datafusion v55.1.0 (current)
Parsed [ 0.031s] (current)
Building datafusion v55.1.0 (baseline)
Built [ 50.374s] (baseline)
Parsing datafusion v55.1.0 (baseline)
Parsed [ 0.033s] (baseline)
Checking datafusion v55.1.0 -> v55.1.0 (no change; assume patch)
Checked [ 0.696s] 223 checks: 223 pass, 31 skip
Summary no semver update required
Finished [ 104.333s] datafusion
Building datafusion-physical-optimizer v55.1.0 (current)
Built [ 35.243s] (current)
Parsing datafusion-physical-optimizer v55.1.0 (current)
Parsed [ 0.020s] (current)
Building datafusion-physical-optimizer v55.1.0 (baseline)
Built [ 35.219s] (baseline)
Parsing datafusion-physical-optimizer v55.1.0 (baseline)
Parsed [ 0.020s] (baseline)
Checking datafusion-physical-optimizer v55.1.0 -> v55.1.0 (no change; assume patch)
Checked [ 0.117s] 223 checks: 223 pass, 31 skip
Summary no semver update required
Finished [ 71.772s] datafusion-physical-optimizer
Building datafusion-physical-plan v55.1.0 (current)
Built [ 32.552s] (current)
Parsing datafusion-physical-plan v55.1.0 (current)
Parsed [ 0.147s] (current)
Building datafusion-physical-plan v55.1.0 (baseline)
Built [ 32.351s] (baseline)
Parsing datafusion-physical-plan v55.1.0 (baseline)
Parsed [ 0.147s] (baseline)
Checking datafusion-physical-plan v55.1.0 -> v55.1.0 (no change; assume patch)
Checked [ 0.753s] 223 checks: 222 pass, 1 fail, 0 warn, 31 skip
--- failure type_method_marked_deprecated: type method #[deprecated] added ---
Description:
A type method is now #[deprecated]. Downstream crates will get a compiler warning when using this method.
ref: https://doc.rust-lang.org/reference/attributes/diagnostics.html#the-deprecated-attribute
impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.50.0/src/lints/type_method_marked_deprecated.ron
Failed in:
method datafusion_physical_plan::aggregates::AggregateExec::with_new_limit_options in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/aggregates/mod.rs:1192
method datafusion_physical_plan::aggregates::AggregateExec::with_limit_options in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/aggregates/mod.rs:1373
Summary semver requires new minor version: 0 major and 1 minor checks failed
Finished [ 67.258s] datafusion-physical-plan
Building datafusion-proto v55.1.0 (current)
Built [ 46.683s] (current)
Parsing datafusion-proto v55.1.0 (current)
Parsed [ 0.015s] (current)
Building datafusion-proto v55.1.0 (baseline)
Built [ 47.870s] (baseline)
Parsing datafusion-proto v55.1.0 (baseline)
Parsed [ 0.017s] (baseline)
Checking datafusion-proto v55.1.0 -> v55.1.0 (no change; assume patch)
Checked [ 0.118s] 223 checks: 223 pass, 31 skip
Summary no semver update required
Finished [ 96.242s] datafusion-proto
```
</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 (part2: topk) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4154447037
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -901,6 +927,202 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to use TopK (min/max heap) optimization in AggregateExec.
+ ///
+ /// If applicable, an inner `AggregateKind` will be set, and later [`ExecutionPlan::execute`]
Review Comment:
Yeah, it is very common in LLM generated code as well, I find
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152771498
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -846,30 +847,55 @@ impl LimitOptions {
}
}
-/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+/// 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)]
-enum AggregateKind {
- /// Ordinary aggregation, including the existing Top-K configuration.
+pub enum AggregateKind {
+ /// Ordinary aggregation, with no limit on the groups retained.
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,
},
+ /// See [`AggregateExec::try_optimize_topk`] for details.
+ TopKMinMax {
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Arc<AggregateFunctionExpr>,
+ limit: usize,
+ descending: bool,
+ nulls_first: bool,
+ },
+ /// See [`AggregateExec::try_optimize_topk`] for details.
+ TopKDistinct {
+ group_by: Arc<PhysicalGroupBy>,
+ limit: usize,
+ descending: bool,
+ },
}
/// Hash aggregate execution plan
#[derive(Debug, Clone)]
pub struct AggregateExec {
/// Aggregation mode (full, partial)
mode: AggregateMode,
- kind: AggregateKind,
+ /// 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.
Review Comment:
Yes we can remove this `pub` after this cleanup: https://github.com/apache/datafusion/pull/25696/changes#r4100893886
--
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 (part2: topk) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#issuecomment-5929653691
Thanks again @2010YOUY01
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152896619
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -952,25 +1175,42 @@ impl AggregateExec {
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.
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. }
+ if aggr_expr.is_empty() => {}
+ // Ordering optimization can revisit an existing TopK with the
+ // same expressions. Preserve its specialization in that case.
+ AggregateKind::TopKMinMax {
+ aggr_expr: existing,
+ ..
+ } if matches!(aggr_expr.as_ref(), [new] if Arc::ptr_eq(existing, new)) => {}
+ AggregateKind::DistinctLimit { group_by, .. }
+ | AggregateKind::TopKMinMax { group_by, .. }
+ | AggregateKind::TopKDistinct { group_by, .. } => {
+ // A replacement expression cannot inherit a specialization
+ // validated for the previous aggregate.
new.kind = AggregateKind::General {
group_by: Arc::clone(group_by),
- // The previous `DistinctLimit` type doesn't include filter
filter_expr: vec![None; aggr_expr.len()].into(),
aggr_expr,
- limit_options: None,
};
}
}
new.metrics = ExecutionPlanMetricsSet::new();
new
}
- /// Clone this exec, overriding only the limit hint.
+ /// Clone with a validated legacy limit hint, falling back to ordinary
+ /// aggregation when the request is unsupported.
+ #[deprecated(
+ since = "56.0.0",
+ note = "This API is intended for internal use only and was inadvertently made public. Do not use this API."
Review Comment:
This API seems impossible to use it correctly. The updated `try_optimize_topk` API makes the mutation itself atomic and safe, but it still has preconditions on the surrounding plan shape. (In this case, the parent `Limit` operator must remain in the plan.)
For that reason, I would prefer to stay conservative and keep the replacement API internal for now. Since this change is part of a deprecation, we can revisit the decision if an external use case comes up.
I updated it in [0880e98](https://github.com/apache/datafusion/pull/25696/commits/0880e988132bc32ce227f627b6c9c4c55340c957) to leave a note, let me know if you have different ideas
> The replacement API is [`try_optimize_topk`]. It still requires specific plan-shape invariants to be used correctly, so it remains internal for now. If you have a use case that requires this API to be public, please open an issue in DataFusion.
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4100893886
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -846,30 +847,55 @@ impl LimitOptions {
}
}
-/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+/// Mutually exclusive aggregation implementations and their configuration.
+///
+/// # Public Only for Internal Use:
Review Comment:
This leakage is not ideal, to clean it up, we could make `CombinePartialFinalAggregate` optimization use a safe API like `AggregateExec::try_combine_partial_final`.
It's better to split it into another PR.
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152757846
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -846,30 +847,55 @@ impl LimitOptions {
}
}
-/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+/// Mutually exclusive aggregation implementations and their configuration.
+///
+/// # Public Only for Internal Use:
Review Comment:
Filed https://github.com/apache/datafusion/issues/25928
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4091329985
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1142,12 +1365,44 @@ impl AggregateExec {
&self.mode
}
- /// Set the limit options for this AggExec
- pub fn with_limit_options(mut self, limit_options: Option<LimitOptions>) -> Self {
- // Restoring an existing hint must not depend on input ordering: a later
- // optimizer may have sorted the child since the hint was introduced.
- if let Some(options) = limit_options
- && options.descending.is_none()
+ /// Set a legacy limit hint. Unsupported requests leave ordinary aggregation.
+ #[deprecated(
+ since = "56.0.0",
+ note = "Use try_optimize_distinct_soft_limit or try_optimize_topk"
+ )]
+ pub fn with_limit_options(self, limit_options: Option<LimitOptions>) -> Self {
+ let ordinary = self.without_optimization();
+ ordinary
+ .clone()
+ .restore_limit_options(limit_options)
+ .unwrap_or(ordinary)
+ }
+
+ /// Clear a specialization. Keeping conservative unordered properties is
+ /// valid for the ordinary hash implementation too.
+ fn without_optimization(mut self) -> Self {
+ self.kind = AggregateKind::General {
+ group_by: Arc::clone(self.group_by()),
+ aggr_expr: self.clone_aggr_exprs(),
+ filter_expr: self.clone_filter_exprs(),
+ };
+ self
+ }
+
+ /// Decode the legacy limit representation at the compatibility boundary.
+ /// Unlike optimizer eligibility, restoring DISTINCT does not require an
+ /// unordered input: a later rule may have sorted the child.
+ fn restore_limit_options(
Review Comment:
TODO 2: the only usage now is `ExecutionPlan::replace_children`, and this function is used to recompute `ExecutionPlan` properties.
This is a bit confusing, but it is how the existing implementation is doing, so probably acceptable.
I'll try to figure out if there is any better way to do it.
##########
datafusion/physical-optimizer/src/combine_partial_final_agg.rs:
##########
@@ -99,7 +99,13 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
input_agg_exec.input_schema(),
)
.map(|combined_agg| {
- combined_agg.with_limit_options(agg_exec.limit_options())
+ #[expect(
+ deprecated,
+ reason = "preserve legacy limit hints while combining aggregates"
+ )]
+ let combined_agg =
Review Comment:
TODO1: avoid deprecated API here
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152896619
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -952,25 +1175,42 @@ impl AggregateExec {
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.
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. }
+ if aggr_expr.is_empty() => {}
+ // Ordering optimization can revisit an existing TopK with the
+ // same expressions. Preserve its specialization in that case.
+ AggregateKind::TopKMinMax {
+ aggr_expr: existing,
+ ..
+ } if matches!(aggr_expr.as_ref(), [new] if Arc::ptr_eq(existing, new)) => {}
+ AggregateKind::DistinctLimit { group_by, .. }
+ | AggregateKind::TopKMinMax { group_by, .. }
+ | AggregateKind::TopKDistinct { group_by, .. } => {
+ // A replacement expression cannot inherit a specialization
+ // validated for the previous aggregate.
new.kind = AggregateKind::General {
group_by: Arc::clone(group_by),
- // The previous `DistinctLimit` type doesn't include filter
filter_expr: vec![None; aggr_expr.len()].into(),
aggr_expr,
- limit_options: None,
};
}
}
new.metrics = ExecutionPlanMetricsSet::new();
new
}
- /// Clone this exec, overriding only the limit hint.
+ /// Clone with a validated legacy limit hint, falling back to ordinary
+ /// aggregation when the request is unsupported.
+ #[deprecated(
+ since = "56.0.0",
+ note = "This API is intended for internal use only and was inadvertently made public. Do not use this API."
Review Comment:
This API seems impossible to use it correctly. The updated `try_optimize_topk` API makes the mutation itself atomic and safe, but it still has preconditions on the surrounding plan shape. (In this case, the parent `Limit` operator must remain in the plan.)
For that reason, I would prefer to stay conservative and keep the replacement API internal for now. Since this change is part of a deprecation, we can revisit the decision if an external use case comes up.
How about leave a note like:
> The replacement API is [`try_optimize_topk`]. It still requires specific plan-shape invariants to be used correctly, so it remains internal for now. If you have a use case that requires this API to be public, please open an issue in DataFusion.
--
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 (part2: topk) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4143878776
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -901,6 +927,202 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to use TopK (min/max heap) optimization in AggregateExec.
+ ///
+ /// If applicable, an inner `AggregateKind` will be set, and later [`ExecutionPlan::execute`]
Review Comment:
I think the note about AggregateKind is somewhat too detailed as it refers to an implementation detail -- I think just saying that "the optimizer must keep the limit operator above because the aggregate may still produce more than K rows" is what is important.
##########
datafusion/physical-optimizer/src/topk_aggregation.rs:
##########
@@ -132,25 +73,21 @@ impl TopKAggregation {
return Ok(Transformed::no(plan));
}
if let Some(aggr) = plan.downcast_ref::<AggregateExec>() {
- // either we run into an Aggregate and transform it
- match Self::transform_agg(
- aggr,
- &cur_col_name,
- order_desc,
- nulls_first,
- limit,
- ) {
+ match aggr
+ .clone()
+ .try_optimize_topk(limit, &sort_col_name, sort_options)
Review Comment:
it is a very nice idea to move this into the AggregateExec itself
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -952,25 +1175,42 @@ impl AggregateExec {
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.
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. }
+ if aggr_expr.is_empty() => {}
+ // Ordering optimization can revisit an existing TopK with the
+ // same expressions. Preserve its specialization in that case.
+ AggregateKind::TopKMinMax {
+ aggr_expr: existing,
+ ..
+ } if matches!(aggr_expr.as_ref(), [new] if Arc::ptr_eq(existing, new)) => {}
+ AggregateKind::DistinctLimit { group_by, .. }
+ | AggregateKind::TopKMinMax { group_by, .. }
+ | AggregateKind::TopKDistinct { group_by, .. } => {
+ // A replacement expression cannot inherit a specialization
+ // validated for the previous aggregate.
new.kind = AggregateKind::General {
group_by: Arc::clone(group_by),
- // The previous `DistinctLimit` type doesn't include filter
filter_expr: vec![None; aggr_expr.len()].into(),
aggr_expr,
- limit_options: None,
};
}
}
new.metrics = ExecutionPlanMetricsSet::new();
new
}
- /// Clone this exec, overriding only the limit hint.
+ /// Clone with a validated legacy limit hint, falling back to ordinary
+ /// aggregation when the request is unsupported.
+ #[deprecated(
+ since = "56.0.0",
+ note = "This API is intended for internal use only and was inadvertently made public. Do not use this API."
Review Comment:
is there an alternate API we should direct them to instead? Maybe `try_optimize_topk` should be public?
It seems strange to say it wasn't meant to be public...
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -846,30 +847,55 @@ impl LimitOptions {
}
}
-/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+/// Mutually exclusive aggregation implementations and their configuration.
+///
+/// # Public Only for Internal Use:
Review Comment:
Yes, I agree making a public API on AggregateExec would be better. Maybe it is worth filing a ticket to track that idea
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -846,30 +847,55 @@ impl LimitOptions {
}
}
-/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+/// 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)]
-enum AggregateKind {
- /// Ordinary aggregation, including the existing Top-K configuration.
+pub enum AggregateKind {
+ /// Ordinary aggregation, with no limit on the groups retained.
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,
},
+ /// See [`AggregateExec::try_optimize_topk`] for details.
+ TopKMinMax {
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Arc<AggregateFunctionExpr>,
+ limit: usize,
+ descending: bool,
+ nulls_first: bool,
+ },
+ /// See [`AggregateExec::try_optimize_topk`] for details.
+ TopKDistinct {
+ group_by: Arc<PhysicalGroupBy>,
+ limit: usize,
+ descending: bool,
+ },
}
/// Hash aggregate execution plan
#[derive(Debug, Clone)]
pub struct AggregateExec {
/// Aggregation mode (full, partial)
mode: AggregateMode,
- kind: AggregateKind,
+ /// 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.
Review Comment:
can we hide it behind an accessor?
Maybe that could be part of a follow on PR to better encapsulate this struct (aka remove `pub` from all the fields)
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152923192
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -901,6 +927,202 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to use TopK (min/max heap) optimization in AggregateExec.
+ ///
+ /// If applicable, an inner `AggregateKind` will be set, and later [`ExecutionPlan::execute`]
Review Comment:
I agree, updated.
I'm aware of this issue, but I still miss it sometimes. It's the kind of comment that is thinking out loud while implementing a feature, but can confuse readers approaching the code from a different angle.
--
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 (part2: topk) [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#issuecomment-5810707073
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/25696?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 `76.73469%` with `57 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 82.48%. Comparing base ([`95bb0a0`](https://app.codecov.io/gh/apache/datafusion/commit/95bb0a0dfa48ca20d017faa55b2fc96ee2f300fe?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`152fa40`](https://app.codecov.io/gh/apache/datafusion/commit/152fa403cba88bfd8b7791e9035387bd6efea32d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 4 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/25696?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/25696?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==) | 75.33% | [39 Missing and 17 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25696?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/25696?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=) | 85.71% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25696?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 #25696 +/- ##
==========================================
- Coverage 82.48% 82.48% -0.01%
==========================================
Files 1140 1140
Lines 437614 437801 +187
Branches 437614 437801 +187
==========================================
+ Hits 360986 361133 +147
- Misses 54834 54860 +26
- Partials 21794 21808 +14
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/25696?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 (part2: topk) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4154450122
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -952,25 +1175,42 @@ impl AggregateExec {
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.
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. }
+ if aggr_expr.is_empty() => {}
+ // Ordering optimization can revisit an existing TopK with the
+ // same expressions. Preserve its specialization in that case.
+ AggregateKind::TopKMinMax {
+ aggr_expr: existing,
+ ..
+ } if matches!(aggr_expr.as_ref(), [new] if Arc::ptr_eq(existing, new)) => {}
+ AggregateKind::DistinctLimit { group_by, .. }
+ | AggregateKind::TopKMinMax { group_by, .. }
+ | AggregateKind::TopKDistinct { group_by, .. } => {
+ // A replacement expression cannot inherit a specialization
+ // validated for the previous aggregate.
new.kind = AggregateKind::General {
group_by: Arc::clone(group_by),
- // The previous `DistinctLimit` type doesn't include filter
filter_expr: vec![None; aggr_expr.len()].into(),
aggr_expr,
- limit_options: None,
};
}
}
new.metrics = ExecutionPlanMetricsSet::new();
new
}
- /// Clone this exec, overriding only the limit hint.
+ /// Clone with a validated legacy limit hint, falling back to ordinary
+ /// aggregation when the request is unsupported.
+ #[deprecated(
+ since = "56.0.0",
+ note = "This API is intended for internal use only and was inadvertently made public. Do not use this API."
Review Comment:
Sounds ok to me
--
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 (part2: topk) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #25696:
URL: https://github.com/apache/datafusion/pull/25696
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152886322
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -901,6 +927,202 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to use TopK (min/max heap) optimization in AggregateExec.
+ ///
+ /// If applicable, an inner `AggregateKind` will be set, and later [`ExecutionPlan::execute`]
Review Comment:
This API seems impossible to use it correctly. The updated `try_optimize_topk` API makes the mutation itself atomic and safe, but it still has preconditions on the surrounding plan shape. (In this case, the parent `Limit` operator must remain in the plan.)
For that reason, I would prefer to stay conservative and keep the replacement API internal for now. Since this change is part of a deprecation, we can revisit it easily.
How about leave a note like:
> The replacement API is [`try_optimize_topk`]. It still requires specific plan-shape invariants to be used correctly, so it remains internal for now. If you have a use case that requires this API to be public, please open an issue in DataFusion.
--
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 (part2: topk) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #25696:
URL: https://github.com/apache/datafusion/pull/25696#discussion_r4152886322
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -901,6 +927,202 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to use TopK (min/max heap) optimization in AggregateExec.
+ ///
+ /// If applicable, an inner `AggregateKind` will be set, and later [`ExecutionPlan::execute`]
Review Comment:
This API seems impossible to use it correctly. The updated `try_optimize_topk` API makes the mutation itself atomic and safe, but it still has preconditions on the surrounding plan shape. (In this case, the parent `Limit` operator must remain in the plan.)
For that reason, I would prefer to stay conservative and keep the replacement API internal for now. Since this change is part of a deprecation, we can revisit it easily.
How about leave a note like:
> The replacement API is [`try_optimize_topk`]. It still requires specific plan-shape invariants to be used correctly, so it remains internal for now. If you have a use case that requires this API to be public, please open an issue in DataFusion.
--
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