25672: fix: drop the redundant Limit above a Sort in the same PushDownLimit visit

[martin-augment]

25672: To review by AI

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Advanced

Run ID: b0985f61-5e3f-42cd-8eb7-3da103c0e145

📥 Commits

Reviewing files that changed from the base of the PR and between 5678858 and 684da9c.

📒 Files selected for processing (3)

• datafusion/core/tests/user_defined/user_defined_plan.rs
• datafusion/optimizer/src/push_down_limit.rs
• datafusion/optimizer/src/push_down_limit/topk_through_join.rs

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

Summary

• PushDownLimit now removes Limit(skip=0) above a Sort during the same optimizer visit. It keeps the smaller of the existing and requested fetch values.
• The rule also attempts TopK-through-join pushdown during that visit. It reports a transformation when it removes the Limit, even if TopK does not move.
• Tests cover existing sort fetch values, one-pass TopK pushdown, and optimizer convergence.
• The TopKOptimizerRule test now rewrites a single-expression Sort with a literal fetch into a TopK node.

Tests

Test execution results were not provided.

Walkthrough

For a zero-skip Limit above a Sort, PushDownLimit now removes the Limit and invokes TopK-through-join pushdown in the same pass. Tests cover fetch values, optimizer pass counts, and TopK placement. TopKOptimizerRule also converts a fetched Sort with one sort expression into a TopK node.

Priority: ⬇️ Low

Merge Risk: ⚪ Minimal · up to 684da

The change preserves the query result limit while enabling the intended one-pass TopK optimization. No actionable merge risk remains beyond normal checks.

✨ Finishing Touches

📝 Generate docstrings

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

Code Review: PushDownLimit Sort/Fetch Optimization

Summary

The changes refactor rewrite_limit to drop a redundant Limit(skip=0) over a Sort in the same pass, and to apply the TopK-through-join pushdown immediately when the Sort replaces the Limit as the visited node. The test changes reflect the new single-pass convergence behavior.

Overall the logic appears sound, but there are several points worth scrutinizing.

---

Issues & Observations

1. result.transformed = true unconditionally overrides the inner result (medium)

let mut result = push_topk_through_join(sort)?;
result.transformed = true;
Ok(result)

This forcibly marks the plan as transformed regardless of what push_topk_through_join returned. The comment justifies this ("the Limit is gone"), which is correct — the Limit node is removed in this branch. However:

• It silently discards any transformed semantics from push_topk_through_join. If that helper ever returns Transformed::no for a legitimate reason (e.g., a future guard), this code would still claim a change. Consider result.transformed |= true semantics being explicit, or better, assert/document that the removal of Limit is the sole source of truth.
• More importantly, verify that push_topk_through_join does not itself rely on transformed being accurate for its own recursion/termination. Overriding it could mask a bug.

2. Behavior change: Limit is now dropped even when skip > 0 is false but fetch semantics differ (verify)

In the old code, the Limit was only dropped when new_fetch != sort.fetch (i.e., the fetch actually changed). Now, with skip == 0, the Limit is dropped unconditionally, even when new_fetch == sort.fetch.

This is the intended change, but double-check the case where sort.fetch is None and fetch is None... actually fetch here comes from the Limit, which always has a fetch in this branch. Still, confirm:

• When sort.fetch == Some(f) and limit.fetch == Some(f) with skip == 0: dropping the Limit is correct because Sort already caps at f. ✅
• When sort.fetch == Some(smaller) and limit.fetch == Some(larger): new_fetch = min(...) = smaller, Limit dropped. ✅ (covered by new test ..._smaller_fetch...)

The logic looks correct, but the asymmetry between the skip > 0 and skip == 0 branches is now significant and deserves a comment explaining why skip == 0 can drop unconditionally while skip > 0 cannot.

3. expr.len() == 1 guard in user_defined_plan.rs — silent fallthrough (low)

if let LogicalPlan::Sort(Sort { expr, input, fetch: Some(fetch) }) = &plan
    && expr.len() == 1
{
    return Ok(Transformed::yes(...));
}

If a Sort has fetch: Some(_) but expr.len() != 1, the code falls through to the Limit handling below, which will not match a Sort and return Transformed::no(plan). That's fine, but the multi-column case is silently ignored — consistent with the documented "does not handle many edge cases" note. No action required, but the comment could explicitly mention the single-column restriction.

4. Test limit_push_down_sort_reports_transformed asserts matches!(result.data, LogicalPlan::Sort(_)) (low)

This is a good assertion, but it only checks the top-level node type. Consider also asserting the fetch value on the resulting Sort to lock in the min behavior, e.g. for sort_fetch = Some(5) the result should be fetch = Some(5) (not Some(10)). The current test would pass even if the fetch were incorrectly set to the Limit's value.

5. Pass-count test brittleness (low)

let optimizer_ctx = OptimizerContext::new().with_max_passes(3);
...
assert_eq!(passes, 2);

This asserts an exact pass count. If the optimizer's convergence detection changes (e.g., a future rule reports transformed differently), this test breaks even though the final plan is correct. The companion snapshot tests already validate the output; consider whether the exact pass count is a meaningful contract or an over-specification. At minimum, the comment should note this is testing convergence, not a hard API guarantee.

6. topk_pushed_through_limit_over_sort_with_same_fetch_in_one_pass — snapshot shows duplicated Sort (informational)

The snapshot:

Sort: t1.b ASC NULLS LAST, fetch=3
  Left Join: t1.a = t2.a
    Sort: t1.b ASC NULLS LAST, fetch=3
      TableScan: t1
    TableScan: t2

The outer Sort and the inner Sort on t1 both carry fetch=3. This is expected (TopK pushed to the join input, outer Sort retained), but it's worth confirming this is the desired final form and not an artifact of the single-pass change. The previous two-pass test presumably produced the same shape; if so, fine.

7. Naming: push_topk_through_join called on a Sort with no join (low)

In the skip == 0 branch, push_topk_through_join(sort) is called even when the Sort's input is not a join. The function name implies join-specific behavior; calling it unconditionally is fine if it's a no-op for non-join inputs, but the comment "The Sort then replaces the Limit as the visited node ... so apply the TopK pushdown here" could be clearer that this is a general TopK pushdown entry point, not join-only.

---

Recommendations

1. Clarify the transformed override — add a comment or restructure so it's obvious the Limit removal is the reason, and confirm push_topk_through_join doesn't depend on its own transformed flag.
2. Add a comment explaining why skip == 0 permits unconditional Limit removal while skip > 0 does not.
3. Strengthen limit_push_down_sort_reports_transformed to assert the resulting Sort's fetch value.
4. Reconsider the exact pass-count assertion or document it as a convergence contract.
5. Verify the duplicated-Sort snapshot is the intended steady state (compare against the prior two-pass output).

No blocking correctness issues found; the core logic change (dropping the redundant Limit in the same pass and reporting transformed) appears correct and is well-covered by the new tests.

[claude]

Review

Traced through rewrite_plan_in_place/Rewriter (datafusion/optimizer/src/optimizer.rs) to confirm the core claim: under ApplyOrder::TopDown, f_down calls rule.rewrite() exactly once per node, then recurses into the children of the resulting plan — the replacement node itself is never re-visited in the same pass. So when rewrite_limit's skip == 0 branch replaces Limit(Sort) with Sort directly, that Sort genuinely won't get a second visit this pass, which is exactly why explicitly calling push_topk_through_join(sort) inline (rather than waiting for a later top-level LogicalPlan::Sort dispatch) is necessary and correct. Good catch, and the comments explaining this framework subtlety are worth keeping — this is exactly the kind of non-obvious "why" that's easy to get wrong.

Correctness

• Forcing result.transformed = true after push_topk_through_join in the skip == 0 branch is correct even when the TopK push itself is a no-op: dropping the Limit node is itself a plan change, and push_topk_through_join has no way to know a Limit was just removed above it, so it can under-report transformed. Without the override, a change could go unreported and the optimizer might treat the plan as already at fixpoint prematurely.
• The skip > 0 path is untouched in behavior (still wraps in Limit when the fetch changed, still returns the unchanged Limit otherwise) — the refactor there is just restructuring, not a behavior change.
• The user_defined_plan.rs test rule update is necessary, not incidental: since PushDownLimit now collapses Limit(skip=0) -> Sort into a bare Sort(fetch) in one pass, the example TopKOptimizerRule needs to recognize that shape directly (in addition to the old Limit -> Sort shape) or the example test would silently stop matching depending on how many passes ran before it.

Test coverage — solid. The new tests cover: single-pass convergence (limit_push_down_sort_settles_in_first_pass), the transformed flag across sort_fetch in {None, Some(10), Some(5)} (covers both the fetch-changed and fetch-unchanged sub-cases), and TopK-through-join collapsing to one pass for both "same fetch" and "smaller/tighter existing sort fetch" cases. One minor gap: the join-based tests only cover limit_fetch >= sort_fetch (3 and 5 against an existing sort fetch of 3); there's no join-based test where the Limit's fetch is smaller than the existing Sort fetch (i.e. fetch_changed = true through a join, not just the non-join case in push_down_limit.rs). Not a correctness concern given the logic is fetch-value-agnostic, but it'd close the matrix.

Style — consistent with the surrounding code; the let-chain pattern in the user_defined_plan.rs change matches the existing unchanged branch just below it, and the module-level doc comment in push_down_limit.rs remains accurate (the filter/outer-join multi-pass dependency it describes is unrelated to this fix and still applies).

Performance — net positive: this removes one full optimizer pass for a very common query shape (ORDER BY ... LIMIT), a small but real planning-time win, with no added overhead in the common non-join case (confirmed the join-walk in push_topk_through_join exits promptly via peel_transparent_layer returning None).

Security — none; this is a pure logical-plan rewrite with no I/O, no unsafe code, and no external input handling.

Nothing blocking. I wasn't able to run cargo fmt --all / cargo clippy --all-targets --all-features -- -D warnings in this environment — please confirm CI is green on those before merge per the repo's contribution guidelines.

[augmentcode]

🤖 Augment PR Summary

Summary: This PR removes a redundant zero-offset Limit above a Sort during PushDownLimit.

Changes:

• Merges the effective fetch into the `Sort` and returns that `Sort` immediately when `skip = 0`.
• Runs the TopK-through-join rewrite during that same visit, since the replacement `Sort` is not revisited as a node.
• Preserves the existing `Limit` wrapper for nonzero offsets, where it remains semantically necessary.
• Marks the rewrite as transformed even when no TopK join pushdown is possible, because the wrapper was removed.
• Updates the user-defined TopK example rule to recognize the fetch-bearing `Sort` shape produced by the native optimizer.
• Updates plan expectations and adds one-pass regression coverage for unfetched, equal-fetch, and tighter-fetch Sort cases.

Technical Notes: The change relies on `Sort.fetch` being semantically equivalent to `LIMIT fetch OFFSET 0`, while retaining the outer Sort after join pushdown to preserve correct output cardinality.

_{🤖 Was this summary useful? React with 👍 or 👎}

[augmentcode]

Review completed. No suggestions at this time.

Comment augment review to trigger a new review at any time.