jayzhan211 commented on code in PR #26773:
URL: https://github.com/apache/datafusion/pull/26773#discussion_r4429777494
##########
datafusion/physical-optimizer/src/limit_pushdown_past_window.rs:
##########
@@ -178,23 +198,27 @@ fn apply_limit(
if !node.is::<SortExec>() && !node.is::<SortPreservingMergeExec>() {
return None;
}
- let latest = ctx.limit.take();
- let Some(fetch) = latest else {
- ctx.limit = None;
- ctx.lookahead = 0;
+ // Taking the state both reads the chain we were walking and forgets it.
+ let TraverseState {
+ limit, lookahead, ..
+ } = std::mem::take(ctx);
+ let Some(fetch) = limit else {
return Some(Transformed::no(Arc::clone(node)));
};
// A fetch that overflows can't drop any rows, so skip the pushdown rather
// than turning the sort into a TopK that holds every row and can't spill.
- let Some(fetch) = fetch.checked_add(ctx.lookahead) else {
- ctx.lookahead = 0;
+ let Some(fetch) = fetch.checked_add(lookahead) else {
return Some(Transformed::no(Arc::clone(node)));
};
let fetch = match node.fetch() {
None => fetch,
Some(existing) => cmp::min(existing, fetch),
};
- Some(Transformed::complete(node.with_fetch(Some(fetch)).unwrap()))
+ // Carry on past this sort rather than stopping. `ctx` is back to its
default
+ // above, which is exactly what it held at the start of the traversal, so
+ // walking the subtree below is equivalent to running the rule on it afresh
+ // -- and that subtree may hold another independent limit over another
window.
+ Some(Transformed::yes(node.with_fetch(Some(fetch)).unwrap()))
Review Comment:
@cool9850311 thanks, both earlier points are fixed and the tests catch them
(I reverted each guard and the new cases went red). One more thing this line
exposes: `grow_limit` combines stacked windows with `max_lookahead`, but when
an outer window reads an inner window's output the lookaheads add up. Main gets
this wrong at the top level too, but there the old `Stop` kept nested chains
safe. With this PR, the nested query below comes back wrong: main leaves the
inner sort full, this branch makes it `TopK(fetch=5)` but it needs 6.
```sql
SELECT empno, salary, s2, sum(empno) OVER (ORDER BY salary ROWS UNBOUNDED
PRECEDING) AS s3
FROM (
SELECT empno, salary, sum(s1) OVER (ORDER BY empno ROWS BETWEEN CURRENT
ROW AND 1 FOLLOWING) AS s2
FROM (SELECT empno, salary, sum(salary) OVER (ORDER BY empno ROWS BETWEEN
CURRENT ROW AND 2 FOLLOWING) AS s1 FROM employees)
ORDER BY empno LIMIT 3
)
ORDER BY salary LIMIT 3;
-- enable_window_limits=false: 363 1865307672 15119054465 662
-- this branch: 363 1865307672 11103612124 662
```
Adding per window fixes it. I ran this locally: `window_limits.slt` and the
unit tests pass, and five EXPLAINs in `window.slt` raise their fetch by 1–5
(10→11, 10→15, 15→16), which is conservative there since those windows don't
read each other. Could you fold it in with the query above as an slt case?
```diff
fn grow_limit(window: &BoundedWindowAggExec, ctx: &mut TraverseState) ->
bool {
let mut max_rel = 0;
+ let mut max_end = 0;
@@
- ctx.max_lookahead(end_bound);
+ max_end = max_end.max(end_bound);
}
@@
- ctx.max_lookahead(ctx.lookahead.saturating_add(max_rel));
+ // A window reads ahead of every row the windows above it need, so
stacked
+ // windows add their lookaheads; expressions in one window share theirs.
+ ctx.lookahead = ctx
+ .lookahead
+ .saturating_add(max_end)
+ .saturating_add(max_rel);
true
}
```
`TraverseState::max_lookahead` then has no callers and can go.
--
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]
