jayzhan211 commented on code in PR #26773:
URL: https://github.com/apache/datafusion/pull/26773#discussion_r4400520842
##########
datafusion/physical-optimizer/src/limit_pushdown_past_window.rs:
##########
@@ -99,21 +107,21 @@ impl PhysicalOptimizerRule for LimitPushPastWindows {
}
// grab the latest limit we see
- if phase == Phase::FindOrGrow && get_limit(&node, &mut ctx) {
+ if ctx.phase == Phase::FindOrGrow && get_limit(&node, &mut ctx) {
Review Comment:
One shape still keeps a full sort: your strict RANGE case, but with a ROWS
outer window. The outer window can grow, so we stay in `Apply`, skip
`get_limit` at the inner `GlobalLimitExec`, and then its `LowerEqual`
cardinality resets and drops it. Only `SortPreservingMergeExec` is both a limit
source and an apply target, so the guard could cover just that. Fine to handle
in a follow-up if you prefer.
```sql
EXPLAIN SELECT empno, s1, sum(empno) OVER (ORDER BY empno ROWS UNBOUNDED
PRECEDING) AS s2
FROM (SELECT empno, salary, sum(salary) OVER (ORDER BY empno ROWS UNBOUNDED
PRECEDING) AS s1 FROM employees ORDER BY empno LIMIT 5)
ORDER BY empno LIMIT 2;
-- today: SortExec: expr=[empno@0 ASC NULLS LAST]; expected: SortExec:
TopK(fetch=5)
```
```rs
if (ctx.phase == Phase::FindOrGrow || !node.is::<SortPreservingMergeExec>())
&& get_limit(&node, &mut ctx)
{
ctx.phase = Phase::FindOrGrow;
return Ok(Transformed::no(node));
}
```
##########
datafusion/physical-optimizer/src/limit_pushdown_past_window.rs:
##########
@@ -188,17 +196,22 @@ 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)));
};
let fetch = match node.fetch() {
- None => fetch + ctx.lookahead,
- Some(existing) => cmp::min(existing, fetch + ctx.lookahead),
+ None => fetch + lookahead,
+ Some(existing) => cmp::min(existing, fetch + lookahead),
};
- 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 for the quick turnaround, both earlier points look good!
One thing this line changes: `Absolute` in `grow_limit` (around line 171) uses
`ctx.limit.unwrap_or(0)`, so with no limit in hand it invents one. On `main`
the `Stop` kept the walk out of everything below an applied sort. Now that we
continue, a nested window whose UDWF returns `Absolute(1)` gets its sort turned
into `TopK(fetch=1)` with no LIMIT above it. I checked with a unit test: `main`
keeps the full sort, this branch truncates. Since this PR exposes it and the
fix is small, could you fold it in here with a unit test (a `row_number`
wrapper overriding `limit_effect`)?
```rs
LimitEffect::Absolute(val) => {
// with no limit in hand there is nothing to raise
if let Some(cur) = ctx.limit {
ctx.limit = Some(cur.max(val));
}
}
```
--
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]
