cool9850311 commented on code in PR #26773:
URL: https://github.com/apache/datafusion/pull/26773#discussion_r4439491931


##########
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:
   Confirmed and fixed in 8559e7c, your patch as-is. main produces the identical
   plan for both your queries, so this predates the PR rather than regressing in
   it -- by a different route though: main stops the walk at the merge, this
   branch walks on with an emptied state.
   
   Two slt cases. Limit above the outer window: the merge gets 5 + 1 and that 6
   starts the next chain, so the inner sort gets 6 + 1. Limit inside the
   subquery: no limit in hand at the merge, so its own fetch=5 starts the chain
   and the sort gets 5 + 1. The second is the one that guards
   `fetch.or_else(|| node.fetch())` -- narrowing it to the computed limit leaves
   the first green and reds only the second. Results differ too, since the
   second query's last row has no following row for the outer window to read.



-- 
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]

Reply via email to