jayzhan211 opened a new pull request, #25983:
URL: https://github.com/apache/datafusion/pull/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.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
