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/05 08:38:30 UTC
(datafusion) branch main updated: fix: keep aggregate limits when combining or rebuilding AggregateExec (#25983)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 982fca6907 fix: keep aggregate limits when combining or rebuilding AggregateExec (#25983)
982fca6907 is described below
commit 982fca6907fcfc12051fe64b989f3e708459a133
Author: Jay Zhan <ja...@gmail.com>
AuthorDate: Mon Oct 5 08:28:59 2026 +0000
fix: keep aggregate limits when combining or rebuilding AggregateExec (#25983)
## Which issue does this PR close?
- No separate issue. Both problems are regressions from #25696.
## Rationale for this change
1. Some `GROUP BY ... LIMIT` queries fail to plan. With one target
partition:
```sql
SET datafusion.execution.target_partitions = 1;
CREATE TABLE t AS SELECT * FROM (VALUES (1, 10), (2, 20), (3, 30), (4,
40), (5, 50), (6, 60)) AS v(a, b);
SELECT b, a FROM t GROUP BY b, a LIMIT 2;
SELECT a % 3 FROM t GROUP BY a % 3 LIMIT 2;
```
both fail with:
```
DataFusion error: CombinePartialFinalAggregate
caused by
Internal error: The AggregateKind should stay either (a) General (b)
DistinctLimit introduced with the previous optimizer pass
`LimitedDistinctAggregation`, it's impossible to have other variant.
```
When the final aggregate's grouping differs from the partial's
(expression keys, or columns not at their output positions),
`LimitedDistinctAggregation` limits only the final aggregate.
`CombinePartialFinalAggregate` treats that pair as impossible and
returns an error.
2. On Parquet tables the TopK aggregate limit is dropped. `EXPLAIN
SELECT DISTINCT k FROM pt ORDER BY k LIMIT 10` keeps `lim=[10]` on the
aggregates for a memory table but not for a Parquet table, so the query
falls back to a full hash aggregation. MIN/MAX TopK aggregates (`ORDER
BY max(x) DESC NULLS LAST LIMIT n`) lose their limit the same way.
Post-optimization filter pushdown rebuilds the Parquet scan when it
accepts a dynamic filter, and rebuilding the aggregates above it
re-applies only the DISTINCT limit, not the TopK ones.
Before #25696 both places carried the limit over with
`with_limit_options(...limit_options())`.
## What changes are included in this PR?
- `CombinePartialFinalAggregate`: if the final aggregate has a DISTINCT
limit, re-apply it to the combined aggregate with
`try_optimize_distinct_soft_limit`, which checks eligibility again.
Otherwise keep the combined aggregate as built.
- `AggregateExec::replace_children`
(`ChildrenPropertiesMode::Recompute`): also re-apply `TopKMinMax` and
`TopKDistinct` through `try_optimize_topk`, using the stored limit,
direction and NULL placement.
## What is the testing strategy for this PR?
New sqllogictest cases. Each one fails without the fix and passes with
it:
- `aggregate.slt`: both queries above at `target_partitions = 1`. An
`EXPLAIN` checks the combined `mode=Single ... lim=[2]` aggregate, and a
row-count check keeps the results deterministic.
- `aggregates_topk.slt`: a Parquet-backed table. `EXPLAIN` checks that
`lim=[10]` stays on both aggregates for `SELECT DISTINCT k ... ORDER BY
k LIMIT 10`, and that `lim=[2]` stays for a MAX TopK query whose scan
gets a semi-join dynamic filter. Result checks are included.
## Are there any user-facing changes?
No API changes. The listed queries plan again, and TopK aggregation
applies again on Parquet scans that receive dynamic filters.
---
.../src/combine_partial_final_agg.rs | 11 ++-
datafusion/physical-plan/src/aggregates/mod.rs | 35 +++++++--
datafusion/sqllogictest/test_files/aggregate.slt | 56 +++++++++++++++
.../sqllogictest/test_files/aggregates_topk.slt | 82 ++++++++++++++++++++++
4 files changed, 173 insertions(+), 11 deletions(-)
diff --git a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
index 74a48f2a1b..7097520061 100644
--- a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
+++ b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
@@ -21,7 +21,6 @@
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, AggregateKind, AggregateMode, PhysicalGroupBy,
@@ -103,7 +102,7 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
};
// 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`,
@@ -126,11 +125,9 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
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"
- );
- }
+ // Step 2 may limit only the final aggregate, e.g. when its
+ // grouping differs from the partial's. Keep both stages.
+ _ => return Ok(Transformed::no(plan)),
};
combined_agg.kind = combined_aggr_kind;
Some(Arc::new(combined_agg))
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index d098272688..e528cf467a 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -2520,11 +2520,38 @@ impl ExecutionPlan for AggregateExec {
Arc::clone(&self.input_schema),
Arc::clone(&self.schema),
)?;
- // Reapply a DISTINCT limit only if the new input remains eligible.
- if let AggregateKind::DistinctLimit { limit, .. } = &self.kind
- && let Some(optimized) =
+ // Reapply a limit only if the new input remains eligible.
+ let optimized = match &self.kind {
+ AggregateKind::General { .. } => None,
+ AggregateKind::DistinctLimit { limit, .. } => {
me.clone().try_optimize_distinct_soft_limit(*limit)
- {
+ }
+ AggregateKind::TopKMinMax {
+ aggr_expr,
+ limit,
+ descending,
+ nulls_first,
+ ..
+ } => me.clone().try_optimize_topk(
+ *limit,
+ aggr_expr.name(),
+ SortOptions::new(*descending, *nulls_first),
+ ),
+ // TopK DISTINCT orders by its single grouping key and
+ // ignores `nulls_first`.
+ AggregateKind::TopKDistinct {
+ group_by,
+ limit,
+ descending,
+ } => group_by.expr().first().and_then(|(_, alias)| {
+ me.clone().try_optimize_topk(
+ *limit,
+ alias,
+ SortOptions::new(*descending, false),
+ )
+ }),
+ };
+ if let Some(optimized) = optimized {
me = optimized.data;
}
me.dynamic_filter.clone_from(&self.dynamic_filter);
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
index d72372dee4..ac0331975b 100644
--- a/datafusion/sqllogictest/test_files/aggregate.slt
+++ b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -7619,6 +7619,62 @@ staging 1
statement ok
DROP TABLE distinct_limit_window_table;
+# `LimitedDistinctAggregation` limits only the final aggregate when its grouping
+# differs from the partial's (columns not at their output positions, or
+# expression keys). `CombinePartialFinalAggregate` must leave such a pair
+# uncombined instead of failing to plan.
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+statement ok
+CREATE TABLE distinct_limit_combine AS
+SELECT * FROM (VALUES (1, 10), (2, 20), (3, 30), (4, 40), (5, 50), (6, 60)) AS v(a, b);
+
+query TT
+EXPLAIN SELECT b, a FROM distinct_limit_combine GROUP BY b, a LIMIT 2;
+----
+logical_plan
+01)Limit: skip=0, fetch=2
+02)--Aggregate: groupBy=[[distinct_limit_combine.b, distinct_limit_combine.a]], aggr=[[]]
+03)----TableScan: distinct_limit_combine projection=[a, b]
+physical_plan
+01)GlobalLimitExec: skip=0, fetch=2
+02)--AggregateExec: mode=Final, gby=[b@0 as b, a@1 as a], aggr=[], lim=[2]
+03)----AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# LIMIT without ORDER BY may return any two groups
+query II
+SELECT count(*), count(*) FILTER (WHERE b = a * 10)
+FROM (SELECT b, a FROM distinct_limit_combine GROUP BY b, a LIMIT 2);
+----
+2 2
+
+query TT
+EXPLAIN SELECT a % 3 FROM distinct_limit_combine GROUP BY a % 3 LIMIT 2;
+----
+logical_plan
+01)Limit: skip=0, fetch=2
+02)--Aggregate: groupBy=[[distinct_limit_combine.a % Int64(3)]], aggr=[[]]
+03)----TableScan: distinct_limit_combine projection=[a]
+physical_plan
+01)GlobalLimitExec: skip=0, fetch=2
+02)--AggregateExec: mode=Final, gby=[distinct_limit_combine.a % Int64(3)@0 as distinct_limit_combine.a % Int64(3)], aggr=[], lim=[2]
+03)----AggregateExec: mode=Partial, gby=[a@0 % 3 as distinct_limit_combine.a % Int64(3)], aggr=[]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT count(*), count(DISTINCT m)
+FROM (SELECT a % 3 AS m FROM distinct_limit_combine GROUP BY a % 3 LIMIT 2);
+----
+2 2
+
+statement ok
+DROP TABLE distinct_limit_combine;
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
#
# regr_*() tests
#
diff --git a/datafusion/sqllogictest/test_files/aggregates_topk.slt b/datafusion/sqllogictest/test_files/aggregates_topk.slt
index e84db491e7..50ee8cc5ca 100644
--- a/datafusion/sqllogictest/test_files/aggregates_topk.slt
+++ b/datafusion/sqllogictest/test_files/aggregates_topk.slt
@@ -789,3 +789,85 @@ drop table topk_two_groups;
statement ok
drop table t0;
+
+# Rebuilding an aggregate must keep its TopK limit. Post-optimization filter
+# pushdown rebuilds a Parquet scan that accepts a dynamic filter, which
+# recomputes the aggregates above it.
+statement ok
+COPY (SELECT (v % 100)::int AS k, v AS x FROM generate_series(1, 1000) AS s(v))
+TO 'test_files/scratch/aggregates_topk/topk.parquet' STORED AS PARQUET;
+
+statement ok
+CREATE EXTERNAL TABLE topk_parquet STORED AS PARQUET
+LOCATION 'test_files/scratch/aggregates_topk/topk.parquet';
+
+# The TopK dynamic filter on `k` is pushed through the aggregates into the scan
+query TT
+EXPLAIN SELECT DISTINCT k FROM topk_parquet ORDER BY k LIMIT 10;
+----
+logical_plan
+01)Sort: topk_parquet.k ASC NULLS LAST, fetch=10
+02)--Aggregate: groupBy=[[topk_parquet.k]], aggr=[[]]
+03)----TableScan: topk_parquet projection=[k]
+physical_plan
+01)SortPreservingMergeExec: [k@0 ASC NULLS LAST], fetch=10
+02)--SortExec: TopK(fetch=10), expr=[k@0 ASC NULLS LAST], preserve_partitioning=[true]
+03)----AggregateExec: mode=FinalPartitioned, gby=[k@0 as k], aggr=[], lim=[10]
+04)------RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=1
+05)--------AggregateExec: mode=Partial, gby=[k@0 as k], aggr=[], lim=[10]
+06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/aggregates_topk/topk.parquet]]}, projection=[k], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible
+
+query I
+SELECT DISTINCT k FROM topk_parquet ORDER BY k LIMIT 10;
+----
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+
+# The semi join's dynamic filter is pushed into the scan below the TopK MAX aggregates
+statement ok
+CREATE TABLE topk_parquet_keys AS VALUES (1), (2), (3);
+
+query TT
+EXPLAIN SELECT k, max(x) FROM topk_parquet
+WHERE k IN (SELECT column1 FROM topk_parquet_keys)
+GROUP BY k ORDER BY max(x) DESC NULLS LAST LIMIT 2;
+----
+logical_plan
+01)Sort: max(topk_parquet.x) DESC NULLS LAST, fetch=2
+02)--Aggregate: groupBy=[[topk_parquet.k]], aggr=[[max(topk_parquet.x)]]
+03)----LeftSemi Join: CAST(topk_parquet.k AS Int64) = __correlated_sq_1.column1
+04)------TableScan: topk_parquet projection=[k, x]
+05)------SubqueryAlias: __correlated_sq_1
+06)--------TableScan: topk_parquet_keys projection=[column1]
+physical_plan
+01)SortPreservingMergeExec: [max(topk_parquet.x)@1 DESC NULLS LAST], fetch=2
+02)--SortExec: TopK(fetch=2), expr=[max(topk_parquet.x)@1 DESC NULLS LAST], preserve_partitioning=[true]
+03)----AggregateExec: mode=FinalPartitioned, gby=[k@0 as k], aggr=[max(topk_parquet.x)], lim=[2]
+04)------RepartitionExec: partitioning=Hash([k@0], 4), input_partitions=4
+05)--------AggregateExec: mode=Partial, gby=[k@0 as k], aggr=[max(topk_parquet.x)], lim=[2]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true
+07)------------HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(column1@0, CAST(topk_parquet.k AS Int64)@2)], projection=[k@0, x@1]
+08)--------------DataSourceExec: partitions=1, partition_sizes=[1]
+09)--------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/aggregates_topk/topk.parquet]]}, projection=[k, x, CAST(k@0 AS Int64) as CAST(topk_parquet.k AS Int64)], output_ordering=[x@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible
+
+query II
+SELECT k, max(x) FROM topk_parquet
+WHERE k IN (SELECT column1 FROM topk_parquet_keys)
+GROUP BY k ORDER BY max(x) DESC NULLS LAST LIMIT 2;
+----
+3 903
+2 902
+
+statement ok
+drop table topk_parquet_keys;
+
+statement ok
+drop table topk_parquet;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org