jayzhan211 opened a new issue, #26659:
URL: https://github.com/apache/datafusion/issues/26659
### Is your feature request related to a problem or challenge?
`LimitPushPastWindows` turns the `SortExec` below a window into a TopK when
a `LIMIT` sits above the window. It only does this for the first such window it
meets. In a `UNION ALL` of two windowed queries with limits, the second branch
keeps a full sort:
```sql
CREATE TABLE t(a INT, x INT) AS VALUES (1, 1), (2, 2), (3, 3), (4, 4), (5,
5);
EXPLAIN SELECT * FROM (SELECT a, sum(x) OVER (ORDER BY a ROWS UNBOUNDED
PRECEDING) s FROM t ORDER BY a LIMIT 2)
UNION ALL
SELECT * FROM (SELECT a, sum(x) OVER (ORDER BY a ROWS UNBOUNDED PRECEDING) s
FROM t ORDER BY a LIMIT 3);
```
```
UnionExec
ProjectionExec: ...
GlobalLimitExec: skip=0, fetch=2
BoundedWindowAggExec: ...
SortExec: TopK(fetch=2), expr=[a@0 ASC NULLS LAST],
preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1]
ProjectionExec: ...
GlobalLimitExec: skip=0, fetch=3
BoundedWindowAggExec: ...
SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false]
<-- no TopK(fetch=3)
DataSourceExec: partitions=1, partition_sizes=[1]
```
The same happens for any later window in the plan, not only under a union.
### Describe the solution you'd like
Both branches get a TopK. Two things in
`datafusion/physical-optimizer/src/limit_pushdown_past_window.rs` cause this:
1. After setting the fetch, `apply_limit` returns
`Transformed::complete(...)` (`:201`). That sets `TreeNodeRecursion::Stop`,
which stops the whole `transform_down`, not just the subtree below the sort.
`TreeNodeRecursion::Jump` would skip only the sort's children.
2. `phase` (`:83`) becomes `Phase::Apply` at the first window and is never
set back. The `reset` closure (`:86-94`) clears the limit but not the phase, so
later limits are never picked up.
### Describe alternatives you've considered
None.
### Additional context
- The early stop was added in #17347 to avoid walking below a sort that
already got its fetch.
- A sqllogictest case like the one above in `window_limits.slt` would cover
it.
--
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]
