jayzhan211 opened a new issue, #27219:
URL: https://github.com/apache/datafusion/issues/27219
### Describe the bug
`ORDER BY x LIMIT n` above a `MATERIALIZED` CTE whose body groups by `x`
returns wrong rows, often none. The physical optimizer rule `TopKAggregation`
puts the limit on the aggregate inside the CTE body, so the body keeps only n
groups, and every reference to the CTE sees only those. With `NOT MATERIALIZED`
the rows are right. The feature (#25580) is on `branch-56`, so 56.0.0 would
ship with this.
### To Reproduce
In `datafusion-cli` (the PostgreSQL dialect is needed to parse
`MATERIALIZED`):
```sql
SET datafusion.sql_parser.dialect = 'PostgreSQL';
CREATE TABLE t AS SELECT value AS x FROM generate_series(1, 200);
-- Query 1
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 2
SELECT * FROM (
WITH c AS MATERIALIZED (SELECT x FROM t GROUP BY x)
SELECT c.x, m.n FROM c CROSS JOIN (SELECT count(*) AS n FROM c) m
) s ORDER BY x LIMIT 3;
```
### Expected behavior
| Query | DataFusion returns | Expected |
|---|---|---|
| 1 | no rows | 101, 102, 103 |
| 2 | (1, 3), (2, 3), (3, 3) | (1, 200), (2, 200), (3, 200) |
With `NOT MATERIALIZED`, query 1 returns 101, 102, 103.
### Additional context
`EXPLAIN` of query 1 shows the limit inside the body:
```text
SortExec: TopK(fetch=3), expr=[x@0 ASC NULLS LAST],
preserve_partitioning=[true]
MaterializedCteExec: name=c
AggregateExec: mode=FinalPartitioned, gby=[x@0 as x], aggr=[], lim=[3]
RepartitionExec: partitioning=Hash([x@0], 12), input_partitions=1
AggregateExec: mode=Partial, gby=[x@0 as x], aggr=[], lim=[3]
DataSourceExec: partitions=1, partition_sizes=[1]
FilterExec: x@0 > 100
MaterializedCteScanExec: name=c
```
- Cause: `TopKAggregation::transform_sort` walks down with `transform_down`
(`datafusion/physical-optimizer/src/topk_aggregation.rs:91`) through every node
whose `cardinality_effect` is `Equal` (`:83`). `MaterializedCteExec` reports
`Equal` (`datafusion/physical-plan/src/materialized_cte.rs:623`) and lists the
body as its first child (`:534-536`). So the walk enters the body, finds `GROUP
BY x` matching the sort column `x`, and limits it. The body is not the node's
output; the continuation is.
- Fix direction: the walk should only enter children whose rows reach the
node's output (for `MaterializedCteExec`, the continuation).
cc @adriangb (author of #25580)
### Willingness to contribute
I'm not planning to fix this myself, so anyone is welcome to pick it up.
Feel free to ping me with questions about the cause or when you'd like a review.
--
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]
