You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/01 10:37:59 UTC
(datafusion) 01/01: refactor: Make AggregateExec state modeling and public updates safe (part2: topk) (#25696)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25696-33e45c721499e0dec0af90caefe3a7dcee4927a0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f065c794fd99ee335b833e94d5cf4379f3d23c44
Author: Yongting You <20...@gmail.com>
AuthorDate: Thu Oct 1 10:37:30 2026 +0000
refactor: Make AggregateExec state modeling and public updates safe (part2: topk) (#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 2 in the issue.
For backgrounds for aggregate topk optimization (this PR targets), see
-
https://github.com/apache/datafusion/pull/25696/changes#diff-4ac7ea0d6b4fd9ff25f1fcfaeda2b767a3d15d3ad5b6c1ba05fd7a5dfb805819R931
## 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`
3. Update related optimizer passes to use the safe alternative
## 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.
-->
---
datafusion/core/tests/execution/coop.rs | 6 +-
.../combine_partial_final_agg.rs | 15 +-
.../src/combine_partial_final_agg.rs | 47 ++-
.../physical-optimizer/src/topk_aggregation.rs | 111 ++---
.../src/aggregates/grouped_topk_stream.rs | 48 +--
datafusion/physical-plan/src/aggregates/mod.rs | 461 +++++++++++++++++----
datafusion/proto/tests/cases/plans/aggregates.rs | 12 +-
7 files changed, 492 insertions(+), 208 deletions(-)
diff --git a/datafusion/core/tests/execution/coop.rs b/datafusion/core/tests/execution/coop.rs
index e02364a053..86425fed37 100644
--- a/datafusion/core/tests/execution/coop.rs
+++ b/datafusion/core/tests/execution/coop.rs
@@ -24,7 +24,7 @@ use datafusion::physical_expr::aggregate::AggregateExprBuilder;
use datafusion::physical_plan;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::aggregates::{
- AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy,
+ AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion::physical_plan::execution_plan::Boundedness;
use datafusion::prelude::SessionContext;
@@ -261,7 +261,9 @@ async fn agg_grouped_topk_yields(
inf.clone(),
inf.schema(),
)?
- .with_limit_options(Some(LimitOptions::new(100))),
+ .try_optimize_topk(100, "max", SortOptions::new(true, false))
+ .unwrap()
+ .data,
);
query_yields(aggr, session_ctx.task_ctx()).await
diff --git a/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs b/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs
index 9e63c341c9..409032d874 100644
--- a/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs
+++ b/datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs
@@ -37,7 +37,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_optimizer::combine_partial_final_agg::CombinePartialFinalAggregate;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::aggregates::{
- AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy,
+ AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion_physical_plan::displayable;
use datafusion_physical_plan::repartition::RepartitionExec;
@@ -244,6 +244,15 @@ fn aggregations_with_limit_combined() -> datafusion_common::Result<()> {
let partial_group_by = PhysicalGroupBy::new_single(groups);
let partial_agg =
partial_aggregate_exec(parquet_exec(schema), partial_group_by, aggr_expr.clone());
+ let partial_agg = Arc::new(
+ partial_agg
+ .downcast_ref::<AggregateExec>()
+ .unwrap()
+ .clone()
+ .try_optimize_distinct_soft_limit(5)
+ .unwrap()
+ .data,
+ );
let groups: Vec<(Arc<dyn PhysicalExpr>, String)> =
vec![(col("c", &partial_agg.schema())?, "c".to_string())];
@@ -260,7 +269,9 @@ fn aggregations_with_limit_combined() -> datafusion_common::Result<()> {
schema,
)
.unwrap()
- .with_limit_options(Some(LimitOptions::new(5))),
+ .try_optimize_distinct_soft_limit(5)
+ .unwrap()
+ .data,
);
let plan: Arc<dyn ExecutionPlan> = final_agg;
// should combine the Partial/Final AggregateExecs to a Single AggregateExec
diff --git a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
index 297a92c45a..74a48f2a1b 100644
--- a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
+++ b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
@@ -21,9 +21,10 @@
use std::sync::Arc;
use datafusion_common::error::Result;
+use datafusion_common::internal_err;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::aggregates::{
- AggregateExec, AggregateMode, PhysicalGroupBy,
+ AggregateExec, AggregateKind, AggregateMode, PhysicalGroupBy,
};
use crate::PhysicalOptimizerRule;
@@ -90,19 +91,49 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
} else {
AggregateMode::SinglePartitioned
};
- AggregateExec::try_new(
+ let Ok(mut combined_agg) = AggregateExec::try_new(
mode,
input_agg_exec.group_expr().clone(),
input_agg_exec.aggr_expr().to_vec(),
input_agg_exec.filter_expr().to_vec(),
Arc::clone(input_agg_exec.input()),
input_agg_exec.input_schema(),
- )
- .map(|combined_agg| {
- combined_agg.with_limit_options(agg_exec.limit_options())
- })
- .ok()
- .map(Arc::new)
+ ) else {
+ return Ok(Transformed::no(plan));
+ };
+
+ // Re-apply distinct limit optimization.
+ //
+ // The optimizer related to aggregates are (in order):
+ // - 1. Initial planning: always final/partial two stage
+ // - 2. `LimitedDistinctAggregation`: push limit into `AggregateExec`,
+ // and build `AggregateKind::DistinctLimit`
+ // - 3. `CombinePartialFinalAggregate`: the current optimization
+ //
+ // Here it restores previously applied optimization in step 2
+ let combined_aggr_kind = match (&agg_exec.kind, &input_agg_exec.kind) {
+ (AggregateKind::General { .. }, AggregateKind::General { .. }) => {
+ combined_agg.kind
+ }
+ (
+ AggregateKind::DistinctLimit { limit, .. },
+ AggregateKind::DistinctLimit {
+ group_by,
+ limit: partial_limit,
+ },
+ ) if limit == partial_limit => AggregateKind::DistinctLimit {
+ // The combined aggregate groups raw input, like the partial.
+ group_by: Arc::clone(group_by),
+ limit: *limit,
+ },
+ _ => {
+ return internal_err!(
+ "The AggregateKind should stay either (a) General (b) DistinctLimit introduced with the previous optimizer pass `LimitedDistinctAggregation`, it's impossible to have other variant"
+ );
+ }
+ };
+ combined_agg.kind = combined_aggr_kind;
+ Some(Arc::new(combined_agg))
} else {
None
};
diff --git a/datafusion/physical-optimizer/src/topk_aggregation.rs b/datafusion/physical-optimizer/src/topk_aggregation.rs
index 0eddb5d550..c734e98449 100644
--- a/datafusion/physical-optimizer/src/topk_aggregation.rs
+++ b/datafusion/physical-optimizer/src/topk_aggregation.rs
@@ -25,8 +25,7 @@ use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_plan::ExecutionPlan;
-use datafusion_physical_plan::aggregates::LimitOptions;
-use datafusion_physical_plan::aggregates::{AggregateExec, topk_types_supported};
+use datafusion_physical_plan::aggregates::AggregateExec;
use datafusion_physical_plan::execution_plan::CardinalityEffect;
use datafusion_physical_plan::projection::ProjectionExec;
use datafusion_physical_plan::sorts::sort::SortExec;
@@ -42,77 +41,20 @@ impl TopKAggregation {
Self {}
}
- fn transform_agg(
- aggr: &AggregateExec,
- order_by: &str,
- order_desc: bool,
- nulls_first: bool,
- limit: usize,
- ) -> Option<Arc<dyn ExecutionPlan>> {
- // Current only support single group key
- let (group_key, group_key_alias) =
- aggr.group_expr().expr().iter().exactly_one().ok()?;
- let kt = group_key.data_type(&aggr.input().schema()).ok()?;
- let vt = if let Some((field, _)) = aggr.get_minmax_desc() {
- field.data_type().clone()
- } else {
- kt.clone()
- };
- if !topk_types_supported(&kt, &vt) {
- return None;
- }
- if aggr.filter_expr().iter().any(|e| e.is_some()) {
- return None;
- }
-
- // Check if this is ordering by an aggregate function (MIN/MAX)
- if let Some((field, desc)) = aggr.get_minmax_desc() {
- // A nullable MIN/MAX starts as NULL and becomes non-NULL when the
- // group sees its first value. With NULLS FIRST that transition
- // worsens the group's rank, so a bounded aggregation cannot safely
- // discard other NULL groups. Use regular aggregation for exact
- // results. Non-nullable inputs never take this transition and can
- // still use TopK.
- let input_nullable = aggr
- .aggr_expr()
- .iter()
- .exactly_one()
- .ok()?
- .expressions()
- .into_iter()
- .exactly_one()
- .ok()?
- .nullable(aggr.input_schema.as_ref())
- .ok()?;
- if nulls_first && input_nullable {
- return None;
- }
- // ensure the sort direction matches aggregate function
- if desc != order_desc {
- return None;
- }
- // ensure the sort is on the same field as the aggregate output
- if order_by != field.name() {
- return None;
- }
- } else if aggr.aggr_expr().is_empty() {
- // This is a GROUP BY without aggregates, check if ordering is on the group key itself
- if order_by != group_key_alias {
- return None;
- }
- } else {
- // Has aggregates but not MIN/MAX, or doesn't DISTINCT
- return None;
- }
-
- // We found what we want: clone, copy the limit down, and return modified node
- let new_aggr = AggregateExec::with_new_limit_options(
- aggr,
- Some(LimitOptions::new_with_order(limit, order_desc)),
- );
- Some(Arc::new(new_aggr))
- }
-
+ /// Push the bound into eligible aggregates to retain fewer groups; keep the
+ /// sort to enforce the final ordering and row count.
+ ///
+ /// ```txt
+ /// Before:
+ /// Sort(fetch=K)
+ /// Aggregate
+ /// input
+ ///
+ /// After:
+ /// Sort(fetch=K)
+ /// Aggregate(TopK(k))
+ /// input
+ /// ```
fn transform_sort(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
let sort = plan.downcast_ref::<SortExec>()?;
@@ -120,10 +62,9 @@ impl TopKAggregation {
let child = children.into_iter().exactly_one().ok()?;
let order = sort.properties().output_ordering()?;
let order = order.iter().exactly_one().ok()?;
- let order_desc = order.options.descending;
- let nulls_first = order.options.nulls_first;
+ let sort_options = order.options;
let order = order.expr.downcast_ref::<Column>()?;
- let mut cur_col_name = order.name().to_string();
+ let mut sort_col_name = order.name().to_string();
let limit = sort.fetch()?;
let mut cardinality_preserved = true;
@@ -132,16 +73,12 @@ 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)
+ {
None => cardinality_preserved = false,
- Some(plan) => return Ok(Transformed::yes(plan)),
+ Some(aggr) => return Ok(Transformed::yes(Arc::new(aggr.data))),
}
} else if let Some(proj) = plan.downcast_ref::<ProjectionExec>() {
// track renames due to successive projections
@@ -149,8 +86,8 @@ impl TopKAggregation {
let Some(src_col) = proj_expr.expr.downcast_ref::<Column>() else {
continue;
};
- if proj_expr.alias == cur_col_name {
- cur_col_name = src_col.name().to_string();
+ if proj_expr.alias == sort_col_name {
+ sort_col_name = src_col.name().to_string();
}
}
} else {
diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
index f4e719a9fc..2c44b68a94 100644
--- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
@@ -19,8 +19,6 @@
use crate::aggregates::group_values::{AggregateArgumentMetrics, GroupByMetrics};
use crate::aggregates::topk::priority_map::PriorityMap;
-#[cfg(debug_assertions)]
-use crate::aggregates::topk_types_supported;
use crate::aggregates::{
AggregateExec, PhysicalGroupBy, aggregate_expressions, aggregate_metric_label,
evaluate_group_by,
@@ -33,7 +31,6 @@ use arrow::compute::concat;
use arrow::datatypes::SchemaRef;
use arrow::util::pretty::print_batches;
use datafusion_common::Result;
-use datafusion_common::internal_datafusion_err;
use datafusion_execution::TaskContext;
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_expr_common::metrics::RecordOutput;
@@ -69,6 +66,7 @@ impl GroupedTopKAggregateStream {
context: &Arc<TaskContext>,
partition: usize,
limit: usize,
+ descending: bool,
) -> Result<Self> {
let agg_schema = Arc::clone(&aggr.schema);
let group_by = Arc::clone(aggr.group_by());
@@ -88,38 +86,17 @@ impl GroupedTopKAggregateStream {
let (expr, _) = &aggr.group_expr().expr()[0];
let kt = expr.data_type(&aggr.input().schema())?;
- // Check if this is a MIN/MAX aggregate or a DISTINCT-like operation
- let (vt, desc) = if let Some((val_field, desc)) = aggr.get_minmax_desc() {
- // MIN/MAX case: use the aggregate output type
- (val_field.data_type().clone(), desc)
- } else {
- // DISTINCT case: use the group key type and get ordering from limit_order_descending
- // The ordering direction is set by the optimizer when it pushes down the limit
- let desc = aggr
- .limit_options()
- .and_then(|config| config.descending)
- .ok_or_else(|| {
- internal_datafusion_err!(
- "Ordering direction required for DISTINCT with limit"
- )
- })?;
- (kt.clone(), desc)
- };
-
- // Type validation is performed by the optimizer and can_use_topk() check.
- // This debug assertion documents the contract without runtime overhead in release builds.
- #[cfg(debug_assertions)]
- {
- debug_assert!(
- topk_types_supported(&kt, &vt),
- "TopK type validation should have been performed by optimizer and can_use_topk(). \
- Found unsupported types: key={kt:?}, value={vt:?}"
- );
- }
+ // The TopK state has already validated the single key/aggregate shape,
+ // types, and direction. MIN/MAX use their value type; grouping-only
+ // aggregation uses its key as the priority value too.
+ let vt = aggr.aggr_expr().first().map_or_else(
+ || kt.clone(),
+ |aggregate| aggregate.field().data_type().clone(),
+ );
// Note: Null values in aggregate columns are filtered by the aggregation layer
// before reaching the heap, so the heap implementations don't need explicit null handling.
- let priority_map = PriorityMap::new(kt, vt, limit, desc)?;
+ let priority_map = PriorityMap::new(kt, vt, limit, descending)?;
Ok(GroupedTopKAggregateStream {
partition,
@@ -329,11 +306,12 @@ impl Stream for GroupedTopKAggregateStream {
mod tests {
use super::*;
use crate::ExecutionPlan;
- use crate::aggregates::{AggregateMode, LimitOptions};
+ use crate::aggregates::AggregateMode;
use crate::collect;
use crate::metrics::MetricValue;
use crate::test::TestMemoryExec;
use arrow::array::{Float64Array, UInt32Array};
+ use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion_common::assert_batches_eq;
@@ -373,7 +351,9 @@ mod tests {
input,
schema,
)?
- .with_limit_options(Some(LimitOptions::new(2))),
+ .try_optimize_topk(2, "MIN(a)", SortOptions::new(false, false))
+ .unwrap()
+ .data,
);
let context = Arc::new(TaskContext::default());
let result = collect(Arc::clone(&aggregate_exec) as _, context).await?;
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index aa05b6c5be..1bf7d0c2fb 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -181,6 +181,7 @@ use parking_lot::Mutex;
use std::collections::HashMap;
use arrow::array::{ArrayRef, UInt8Array, UInt16Array, UInt32Array, UInt64Array};
+use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use arrow_schema::FieldRef;
@@ -846,15 +847,19 @@ 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.
@@ -862,6 +867,20 @@ enum AggregateKind {
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
@@ -869,7 +888,14 @@ enum AggregateKind {
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.
+ /// This field is not part of the supported public API.
+ #[doc(hidden)]
+ pub kind: AggregateKind,
/// Input plan, could be a partial aggregate or the input to the aggregate
pub input: Arc<dyn ExecutionPlan>,
/// Schema after the aggregate is applied. Contains the group by columns followed by the
@@ -901,6 +927,201 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to use TopK (min/max heap) optimization in AggregateExec.
+ ///
+ /// The optimizer must keep the limit operator above because the aggregate may
+ /// still produce more than K rows.
+ /// Returns `None` if not applicable.
+ ///
+ /// # Public Only for Internal Use:
+ ///
+ /// This is not part of the supported public API, it's made public for internal
+ /// optimizer usage.
+ ///
+ /// # TopK Optimization Overview
+ ///
+ /// Retain candidate groups for a parent ORDER BY / LIMIT. There are two applicable
+ /// cases:
+ ///
+ /// ## TopK MIN/MAX
+ ///
+ /// ```txt
+ /// SELECT k, MAX(v) AS score FROM t GROUP BY k
+ /// ORDER BY score DESC NULLS LAST LIMIT 10;
+ ///
+ /// Before:
+ /// Sort: score DESC, fetch=10
+ /// Aggregate: k, MAX(v)
+ /// input
+ ///
+ /// After (Added info in `AggregateExec` to trigger heap optimization):
+ /// Sort: score DESC, fetch=10
+ /// Aggregate: k, MAX(v), TopKMinMax(10)
+ /// input
+ /// ```
+ ///
+ /// For an aggregate output named `score`, the example uses:
+ /// ```txt
+ /// try_optimize_topk(10, "score", SortOptions {
+ /// descending: true, // DESC
+ /// nulls_first: false, // NULLS LAST
+ /// })
+ /// ```
+ ///
+ /// A min-heap keeps the 10 largest non-NULL scores.
+ /// MIN with ASC uses a max-heap.
+ ///
+ /// ## TopK DISTINCT
+ ///
+ /// ```txt
+ /// SELECT DISTINCT k FROM t ORDER BY k ASC NULLS LAST LIMIT 10;
+ ///
+ /// Before:
+ /// Sort: k ASC, fetch=10
+ /// Aggregate: k, no aggregates
+ /// input
+ ///
+ /// After (Added info in `AggregateExec` to trigger heap optimization):
+ /// Sort: k ASC, fetch=10
+ /// Aggregate: k, TopKDistinct(10)
+ /// input
+ /// ```
+ ///
+ /// Here the grouping key is the output used for ordering:
+ /// ```txt
+ /// try_optimize_topk(10, "k", SortOptions {
+ /// descending: false, // ASC
+ /// nulls_first: false, // NULLS LAST
+ /// })
+ /// ```
+ ///
+ /// A max-heap keeps the 10 smallest non-NULL keys.
+ /// DESC uses a min-heap.
+ #[doc(hidden)]
+ pub fn try_optimize_topk(
+ mut self,
+ limit: usize,
+ order_by: &str,
+ options: SortOptions,
+ ) -> Option<Transformed<Self>> {
+ if limit == 0
+ || matches!(self.kind, AggregateKind::DistinctLimit { .. })
+ || self.group_by().has_grouping_set()
+ || self.filter_expr().iter().any(Option::is_some)
+ {
+ return None;
+ }
+ // The priority map needs one supported key and value type.
+ let (group_key, group_key_alias) =
+ self.group_expr().expr().iter().exactly_one().ok()?;
+ let input_schema = self.input.schema();
+ let key_type = group_key.data_type(&input_schema).ok()?;
+ let minmax = self.get_minmax_desc();
+ let value_type = minmax
+ .as_ref()
+ .map_or(&key_type, |(field, _)| field.data_type());
+ if !topk_types_supported(&key_type, value_type)
+ || self.schema.fields().len() != 1 + self.aggr_expr().len()
+ {
+ return None;
+ }
+
+ let kind = if let Some((field, descending)) = minmax {
+ let aggregate = self.aggr_expr().iter().exactly_one().ok()?;
+ if field.name() != order_by
+ || descending != options.descending
+ || !aggregate.order_bys().is_empty()
+ {
+ return None;
+ }
+ let argument = aggregate.expressions().into_iter().exactly_one().ok()?;
+ // NULLS FIRST can worsen a group's rank when its first value arrives.
+ if options.nulls_first && argument.nullable(&self.input_schema).ok()? {
+ return None;
+ }
+
+ // Raw input, partial state, and output must fit the same heap value.
+ let state_field = aggregate
+ .state_fields()
+ .ok()?
+ .into_iter()
+ .exactly_one()
+ .ok()?;
+ let input_argument = aggregate_expressions(self.aggr_expr(), &self.mode, 1)
+ .ok()?
+ .into_iter()
+ .flatten()
+ .exactly_one()
+ .ok()?;
+ let value_type = field.data_type();
+ if state_field.data_type() != value_type
+ || argument.data_type(&self.input_schema).ok()? != *value_type
+ || input_argument.data_type(&input_schema).ok()? != *value_type
+ {
+ return None;
+ }
+ AggregateKind::TopKMinMax {
+ group_by: Arc::clone(self.group_by()),
+ aggr_expr: Arc::clone(aggregate),
+ limit,
+ descending,
+ nulls_first: options.nulls_first,
+ }
+ } else if self.aggr_expr().is_empty() && group_key_alias == order_by {
+ AggregateKind::TopKDistinct {
+ group_by: Arc::clone(self.group_by()),
+ limit,
+ descending: options.descending,
+ }
+ } else {
+ return None;
+ };
+
+ // Preserve compatible tighter bounds; never switch an existing order.
+ match &self.kind {
+ AggregateKind::TopKMinMax { nulls_first, .. }
+ if *nulls_first != options.nulls_first =>
+ {
+ return None;
+ }
+ AggregateKind::TopKMinMax {
+ limit: existing,
+ descending,
+ ..
+ }
+ | AggregateKind::TopKDistinct {
+ limit: existing,
+ descending,
+ ..
+ } => {
+ if *descending != options.descending {
+ return None;
+ }
+ if *existing <= limit {
+ return Some(Transformed::no(self));
+ }
+ }
+ _ => {}
+ }
+
+ // Commit the kind and properties together: heap output is unordered and final.
+ self.kind = kind;
+ self.input_order_mode = InputOrderMode::Linear;
+ self.required_input_ordering = None;
+ // Keep unchanged properties so parent aggregates do not need rebuilding.
+ if !self.cache.eq_properties.oeq_class().is_empty()
+ || self.cache.emission_type != EmissionType::Final
+ {
+ let mut equivalence = self.cache.eq_properties.clone();
+ equivalence.clear_orderings();
+ let cache = Arc::make_mut(&mut self.cache);
+ cache.set_eq_properties(equivalence);
+ cache.emission_type = EmissionType::Final;
+ }
+ self.metrics = ExecutionPlanMetricsSet::new();
+ Some(Transformed::yes(self))
+ }
+
/// Try to stop after enough distinct grouping keys have been accumulated.
///
/// The caller must establish that discarding other groups is legal, for
@@ -941,9 +1162,10 @@ impl AggregateExec {
Some(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.
+ /// Clone with replacement aggregate expressions.
+ ///
+ /// The caller must preserve the schema, filters, and ordering requirements.
+ /// Changing expressions in a TopK aggregate drops its specialization.
pub fn with_new_aggr_exprs(
&self,
aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>,
@@ -952,15 +1174,23 @@ 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,
};
}
}
@@ -968,9 +1198,18 @@ impl AggregateExec {
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 = "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."
+ )]
pub fn with_new_limit_options(&self, limit_options: Option<LimitOptions>) -> Self {
- let mut new = self.clone().with_limit_options(limit_options);
+ let mut new = self.clone().without_optimization();
+ new = new
+ .clone()
+ .restore_limit_options(limit_options)
+ .unwrap_or(new);
new.metrics = ExecutionPlanMetricsSet::new();
new
}
@@ -1133,7 +1372,6 @@ impl AggregateExec {
group_by,
aggr_expr: aggr_expr.into(),
filter_expr,
- limit_options: None,
},
input,
schema,
@@ -1155,12 +1393,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 = "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."
+ )]
+ 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(
+ mut self,
+ limit_options: Option<LimitOptions>,
+ ) -> Option<Self> {
+ let Some(options) = limit_options else {
+ return Some(self);
+ };
+ if options.limit == 0 {
+ return None;
+ }
+ if options.descending.is_none()
&& !self.group_by().is_true_no_grouping()
&& self.aggr_expr().is_empty()
&& self.filter_expr().is_empty()
@@ -1169,29 +1439,42 @@ impl AggregateExec {
group_by: Arc::clone(self.group_by()),
limit: options.limit,
};
+ Some(self)
} else {
- match &mut self.kind {
- AggregateKind::General {
- limit_options: old, ..
- } => *old = limit_options,
- AggregateKind::DistinctLimit { group_by, .. } => {
- self.kind = AggregateKind::General {
- group_by: Arc::clone(group_by),
- aggr_expr: Arc::from([]),
- filter_expr: Arc::from([]),
- limit_options,
- };
- }
- }
+ // The legacy representation has no NULL placement. NULLS LAST
+ // accepts every valid serialized MIN/MAX configuration; grouping
+ // only TopK retains a NULL candidate for either placement.
+ let descending = options
+ .descending
+ .or_else(|| self.get_minmax_desc().map(|(_, desc)| desc))?;
+ let order_by = match self.aggr_expr() {
+ [] => self.group_by().expr.first()?.1.clone(),
+ [aggregate] => aggregate.name().to_owned(),
+ _ => return None,
+ };
+ self.try_optimize_topk(
+ options.limit,
+ &order_by,
+ SortOptions {
+ descending,
+ nulls_first: false,
+ },
+ )
+ .map(|result| result.data)
}
- self
}
/// Get the limit options (if set)
pub fn limit_options(&self) -> Option<LimitOptions> {
match &self.kind {
- AggregateKind::General { limit_options, .. } => *limit_options,
+ AggregateKind::General { .. } => None,
AggregateKind::DistinctLimit { limit, .. } => Some(LimitOptions::new(*limit)),
+ AggregateKind::TopKMinMax {
+ limit, descending, ..
+ }
+ | AggregateKind::TopKDistinct {
+ limit, descending, ..
+ } => Some(LimitOptions::new_with_order(*limit, *descending)),
}
}
@@ -1204,7 +1487,9 @@ impl AggregateExec {
fn group_by(&self) -> &Arc<PhysicalGroupBy> {
match &self.kind {
AggregateKind::General { group_by, .. }
- | AggregateKind::DistinctLimit { group_by, .. } => group_by,
+ | AggregateKind::DistinctLimit { group_by, .. }
+ | AggregateKind::TopKMinMax { group_by, .. }
+ | AggregateKind::TopKDistinct { group_by, .. } => group_by,
}
}
@@ -1212,7 +1497,9 @@ impl AggregateExec {
fn group_by_mut(&mut self) -> &mut Arc<PhysicalGroupBy> {
match &mut self.kind {
AggregateKind::General { group_by, .. }
- | AggregateKind::DistinctLimit { group_by, .. } => group_by,
+ | AggregateKind::DistinctLimit { group_by, .. }
+ | AggregateKind::TopKMinMax { group_by, .. }
+ | AggregateKind::TopKDistinct { group_by, .. } => group_by,
}
}
@@ -1225,7 +1512,12 @@ impl AggregateExec {
pub fn aggr_expr(&self) -> &[Arc<AggregateFunctionExpr>] {
match &self.kind {
AggregateKind::General { aggr_expr, .. } => aggr_expr,
- AggregateKind::DistinctLimit { .. } => &[],
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. } => {
+ &[]
+ }
+ AggregateKind::TopKMinMax { aggr_expr, .. } => {
+ std::slice::from_ref(aggr_expr)
+ }
}
}
@@ -1233,21 +1525,32 @@ impl AggregateExec {
pub fn filter_expr(&self) -> &[Option<Arc<dyn PhysicalExpr>>] {
match &self.kind {
AggregateKind::General { filter_expr, .. } => filter_expr,
- AggregateKind::DistinctLimit { .. } => &[],
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. } => {
+ &[]
+ }
+ AggregateKind::TopKMinMax { .. } => &[None],
}
}
fn clone_aggr_exprs(&self) -> Arc<[Arc<AggregateFunctionExpr>]> {
match &self.kind {
AggregateKind::General { aggr_expr, .. } => Arc::clone(aggr_expr),
- AggregateKind::DistinctLimit { .. } => Arc::from([]),
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. } => {
+ Arc::from([])
+ }
+ AggregateKind::TopKMinMax { aggr_expr, .. } => {
+ Arc::from(std::slice::from_ref(aggr_expr))
+ }
}
}
fn clone_filter_exprs(&self) -> Arc<[Option<Arc<dyn PhysicalExpr>>]> {
match &self.kind {
AggregateKind::General { filter_expr, .. } => Arc::clone(filter_expr),
- AggregateKind::DistinctLimit { .. } => Arc::from([]),
+ AggregateKind::DistinctLimit { .. } | AggregateKind::TopKDistinct { .. } => {
+ Arc::from([])
+ }
+ AggregateKind::TopKMinMax { .. } => Arc::from([None]),
}
}
@@ -1335,29 +1638,21 @@ impl AggregateExec {
)?));
}
- // Grouping by an expression that has a sort/limit upstream.
- //
- // `GroupedTopKAggregateStream` keeps a priority queue, so it only works
- // when an ordering direction is available: either this is a MIN/MAX
- // aggregate, whose direction is implied by the accumulator, or the
- // optimizer pushed a direction down next to the limit.
- //
- // A direction-less limit is a soft hint pushed by
- // `LimitedDistinctAggregation`, which only pushes it while
- // `is_unordered_unfiltered_group_by_distinct()` holds. That predicate is
- // not stable: a later rule such as `EnsureRequirements` can insert the
- // sort a window function requires below this aggregate, which gives the
- // rebuilt aggregate an output ordering and makes the predicate false
- // without ever supplying a direction. Falling back to the regular
- // grouped streams is correct in that case, because they treat the limit
- // as a soft limit and the `LIMIT` above this aggregate still truncates
- // the result.
- if let Some(config) = self.limit_options()
- && !self.is_unordered_unfiltered_group_by_distinct()
- && (config.descending.is_some() || self.get_minmax_desc().is_some())
+ if let AggregateKind::TopKMinMax {
+ limit, descending, ..
+ }
+ | AggregateKind::TopKDistinct {
+ limit, descending, ..
+ } = &self.kind
{
return Ok(StreamType::GroupedPriorityQueue(
- GroupedTopKAggregateStream::new(self, context, partition, config.limit)?,
+ GroupedTopKAggregateStream::new(
+ self,
+ context,
+ partition,
+ *limit,
+ *descending,
+ )?,
));
}
@@ -2225,7 +2520,13 @@ impl ExecutionPlan for AggregateExec {
Arc::clone(&self.input_schema),
Arc::clone(&self.schema),
)?;
- me = me.with_limit_options(self.limit_options());
+ // Reapply a DISTINCT limit only if the new input remains eligible.
+ if let AggregateKind::DistinctLimit { limit, .. } = &self.kind
+ && let Some(optimized) =
+ me.clone().try_optimize_distinct_soft_limit(*limit)
+ {
+ me = optimized.data;
+ }
me.dynamic_filter.clone_from(&self.dynamic_filter);
Ok(Arc::new(me))
}
@@ -2810,7 +3111,13 @@ impl AggregateExec {
Some(descending) => LimitOptions::new_with_order(fetch, descending),
None => LimitOptions::new(fetch),
};
- aggregate.with_limit_options(Some(options))
+ aggregate
+ .restore_limit_options(Some(options))
+ .ok_or_else(|| {
+ datafusion_common::internal_datafusion_err!(
+ "Invalid DISTINCT or TopK aggregate limit configuration"
+ )
+ })?
} else {
aggregate
};
@@ -5134,7 +5441,9 @@ mod tests {
partial_input,
Arc::clone(&schema),
)?
- .with_limit_options(Some(LimitOptions::new(2))),
+ .try_optimize_distinct_soft_limit(2)
+ .unwrap()
+ .data,
);
let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?;
@@ -5175,7 +5484,9 @@ mod tests {
final_input,
Arc::clone(&schema),
)?
- .with_limit_options(Some(LimitOptions::new(2))),
+ .try_optimize_distinct_soft_limit(2)
+ .unwrap()
+ .data,
);
let final_stream = final_aggregate.execute_typed(0, &task_ctx)?;
@@ -7725,17 +8036,25 @@ mod tests {
let input = Arc::new(StatisticsExec::new(stats, (**schema).clone()))
as Arc<dyn ExecutionPlan>;
+ let (aggregates, filters) = if limit.is_some() {
+ (vec![], vec![])
+ } else {
+ (vec![count_a_aggregate(schema)?], vec![None])
+ };
let mut agg = AggregateExec::try_new(
mode,
group_by,
- vec![count_a_aggregate(schema)?],
- vec![None],
+ aggregates,
+ filters,
input,
Arc::clone(schema),
)?;
if let Some(limit) = limit {
- agg = agg.with_limit_options(Some(limit));
+ agg = agg
+ .try_optimize_topk(limit.limit, "a", SortOptions::new(false, false))
+ .unwrap()
+ .data;
}
Ok(agg)
diff --git a/datafusion/proto/tests/cases/plans/aggregates.rs b/datafusion/proto/tests/cases/plans/aggregates.rs
index e57ac9fb50..993f6dc9a2 100644
--- a/datafusion/proto/tests/cases/plans/aggregates.rs
+++ b/datafusion/proto/tests/cases/plans/aggregates.rs
@@ -27,7 +27,7 @@ use datafusion::physical_expr::aggregate::AggregateExprBuilder;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::update_aggr_exprs::OptimizeAggregateOrder;
use datafusion::physical_plan::aggregates::{
- AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy,
+ AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::expressions::{PhysicalSortExpr, col, lit};
@@ -44,6 +44,7 @@ use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_co
use datafusion_functions_aggregate::array_agg::array_agg_udaf;
use datafusion_functions_aggregate::average::avg_udaf;
use datafusion_functions_aggregate::first_last::first_value_udaf;
+use datafusion_functions_aggregate::min_max::min_udaf;
use datafusion_functions_aggregate::nth_value::nth_value_udaf;
use datafusion_functions_aggregate::string_agg::string_agg_udaf;
use datafusion_functions_aggregate::sum::sum_udaf;
@@ -206,9 +207,9 @@ fn roundtrip_aggregate_with_limit() -> Result<()> {
vec![(col("a", &schema)?, "unused".to_string())];
let aggregates = vec![
- AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?])
+ AggregateExprBuilder::new(min_udaf(), vec![col("b", &schema)?])
.schema(Arc::clone(&schema))
- .alias("AVG(b)")
+ .alias("MIN(b)")
.build()
.map(Arc::new)?,
];
@@ -221,7 +222,10 @@ fn roundtrip_aggregate_with_limit() -> Result<()> {
Arc::new(EmptyExec::new(schema.clone())),
schema,
)?;
- let agg = agg.with_limit_options(Some(LimitOptions::new_with_order(12, false)));
+ let agg = agg
+ .try_optimize_topk(12, "MIN(b)", SortOptions::new(false, false))
+ .unwrap()
+ .data;
roundtrip_test(Arc::new(agg))
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org