jayzhan211 commented on code in PR #26773:
URL: https://github.com/apache/datafusion/pull/26773#discussion_r4434439119
##########
datafusion/physical-optimizer/src/limit_pushdown_past_window.rs:
##########
@@ -141,23 +162,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) => existing.min(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, the stacked-lookahead fix holds up: I ran on/off
comparisons with `lead`, three stacked windows, `PARTITION BY`, and `OFFSET`,
and they all match. One case of #26659 is still left. An SPM's fetch also
limits everything below it, but once the walk is past a window that fetch never
starts a new chain: `ctx` resets after the SPM gets its fetch, and with no
limit in hand the SPM's own fetch is ignored. So the sort under a second,
partitioned window stays a full sort (main does the same). I should have caught
this when I suggested the SPM exception in an earlier round, sorry for the
extra round.
```sql
SELECT depname, empno, s1, sum(s1) OVER (ORDER BY depname, empno ROWS
BETWEEN CURRENT ROW AND 1 FOLLOWING) AS s2
FROM (SELECT depname, empno, sum(salary) OVER (PARTITION BY depname ORDER BY
empno ROWS BETWEEN CURRENT ROW AND 1 FOLLOWING) AS s1 FROM employees)
ORDER BY depname, empno LIMIT 5;
-- SortPreservingMergeExec fetch=6, but the SortExec under the inner window
stays a full sort
```
With the change below it becomes `TopK(fetch=7)`, and the results don't
change. I ran the window/limit slt files, the unit tests and clippy with it
(`window.slt:6417` still covers the case where the SPM itself gets the fetch).
Could you add the query above, plus the same query with `LIMIT 5` moved into
the subquery (no limit above the outer window → `TopK(fetch=6)`), as slt cases?
````suggestion
// Taking the state both reads the chain we were walking and forgets it.
let TraverseState {
limit, lookahead, ..
} = std::mem::take(ctx);
// 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 fetch = limit
.and_then(|limit| limit.checked_add(lookahead))
.map(|fetch| node.fetch().map_or(fetch, |existing|
existing.min(fetch)));
// A `SortPreservingMergeExec` outputs its input's first rows, so its
fetch
// also bounds everything below it: start the next chain from it.
if node.is::<SortPreservingMergeExec>() {
ctx.reset_limit(fetch.or_else(|| node.fetch()));
}
let Some(fetch) = fetch else {
return Some(Transformed::no(Arc::clone(node)));
};
// Carry on past this sort rather than stopping: the subtree below may
hold
// another independent limit over another window.
Some(Transformed::yes(node.with_fetch(Some(fetch)).unwrap()))
````
--
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]
