jayzhan211 opened a new pull request, #27416:
URL: https://github.com/apache/datafusion/pull/27416
## Which issue does this PR close?
- Closes #27219.
## Rationale for this change
When a query declares a `MATERIALIZED` CTE whose body groups by the sort
column, `ORDER BY ... LIMIT n` above that query returns wrong rows. The query
below returns no rows instead of 3:
```sql
SET datafusion.sql_parser.dialect = 'PostgreSQL'; -- needed to parse
MATERIALIZED
CREATE TABLE t AS SELECT value AS x FROM generate_series(1, 200);
SELECT * FROM (
WITH c AS MATERIALIZED (SELECT x FROM t GROUP BY x)
SELECT x FROM c WHERE x > 100
) s ORDER BY x LIMIT 3;
```
| Query | main | This PR |
|---|---|---|
| Above | no rows | 101, 102, 103 |
| Same CTE, `SELECT c.x, m.n FROM c CROSS JOIN (SELECT count(*) AS n FROM c)
m` | `n` = 3 | `n` = 200 |
The physical optimizer rule `TopKAggregation` starts at `SortExec:
TopK(fetch=3)` and walks down through nodes that keep the row count.
`MaterializedCteExec` says it keeps the row count, and its first child is the
CTE body. So the walk goes into the body and sets `lim=[3]` on `GROUP BY x`.
The body runs once, and every reference to the CTE reads its output, so every
reference sees only the groups that the limit kept. `MaterializedCteExec`
outputs only its second child, the continuation; the body is never its output.
The same bug is on `branch-56` (56.0.0-rc1), and this patch applies there
cleanly.
This PR stops the walk at `MaterializedCteExec`.
## What changes are included in this PR?
`datafusion/physical-optimizer/src/topk_aggregation.rs`: one new branch
stops the walk at `MaterializedCteExec`. `cte_materialized.slt`: the two
queries above.
The decision I'd most like reviewed: the rule names `MaterializedCteExec`
instead of stopping at every node with more than one child. The general rule
would also remove a `lim` that main sets today, on an aggregate that is the
first input of a `UnionExec`, which is a valid limit:
```text
SortExec: TopK(fetch=3), expr=[x@0 ASC NULLS LAST]
UnionExec
AggregateExec: mode=FinalPartitioned, gby=[x@0 as x], aggr=[], lim=[3]
...
```
<details>
<summary>Why the walk doesn't go into the continuation instead</summary>
Limiting an aggregate in the continuation would be valid, but it is a new
optimization, and it only matters when `ORDER BY ... LIMIT` is outside the
`WITH` query. When it is inside, the sort is already in the continuation and
the rule applies there as before:
```text
MaterializedCteExec: name=c
DataSourceExec: ...
SortPreservingMergeExec: [y@0 ASC NULLS LAST], fetch=3
SortExec: TopK(fetch=3), expr=[y@0 ASC NULLS LAST]
AggregateExec: mode=FinalPartitioned, gby=[y@0 as y], aggr=[], lim=[3]
```
</details>
Who is affected: queries that use `AS MATERIALIZED` (new in 56, PostgreSQL
dialect only) with `ORDER BY` + `LIMIT` above the `WITH` query and a `GROUP BY`
on the sort column in the CTE body. No other plan changes. No API change.
## What is the testing strategy for this PR?
- The two queries above are added to `cte_materialized.slt`. Both fail on
main and pass with this PR.
- The full sqllogictest suite passes (306 files) with no other plan changes,
and so do the `datafusion-physical-optimizer` unit tests.
- I compared the plans of several `UNION ALL`, scalar subquery, and
`MATERIALIZED` CTE queries before and after this change. Only the plan of the
query above changed: its body aggregates no longer have `lim=[3]`.
## Are there any user-facing changes?
Yes: these queries now return correct results. There are no API changes.
--
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]
