Skip to content

Commit 5719930

Browse files
adriangbclaude
adriangb
and
claude
authored
fix: drop the redundant Limit above a Sort in the same PushDownLimit visit (#25672)
## Which issue does this PR close? - Closes #25671. ## Rationale for this change For `Limit(skip=0, fetch=n)` over a `Sort` without a fetch, `PushDownLimit` sets `Sort.fetch = n` but keeps the `Limit`. It removes the `Limit` only on its next visit, in the next optimizer pass. Thus every `ORDER BY ... LIMIT n` query runs 3 logical optimizer passes where 2 are enough. The final plan does not change. Only planning time does. A query engine that wraps every query in a row-cap `LIMIT` over an `ORDER BY` hits this on nearly every query. Reproduction, `datafusion-cli` on `main` (95bb0a0): ```sql CREATE TABLE t (a INT, ts TIMESTAMP) AS VALUES (1, TIMESTAMP '2026-01-01T00:00:00'); EXPLAIN VERBOSE SELECT a FROM t ORDER BY ts DESC LIMIT 10; ``` Rules that change the plan, per pass (the `logical_plan after <rule>` rows that are not `SAME TEXT AS ABOVE`): ```text main: pass 0: push_down_limit, optimize_projections pass 1: push_down_limit <- removes the Limit it kept in pass 0 pass 2: (no change) this PR: pass 0: push_down_limit, optimize_projections pass 1: (no change) ``` The final `logical_plan` is identical on both: ```text Projection: t.a Sort: t.ts DESC NULLS FIRST, fetch=10 TableScan: t projection=[a, ts] ``` ## What changes are included in this PR? In `rewrite_limit`, `LogicalPlan::Sort` branch, when the rule gives the `Sort` a new fetch and `skip == 0`, it now returns the `Sort` and drops the `Limit` in the same visit. With `skip > 0` the `Limit` stays, as before. When the `Sort` already has an equal or smaller fetch and `skip == 0`, `main` already dropped the `Limit`. Both `skip == 0` cases now share one branch. One detail needs attention. The optimizer applies `PushDownLimit` top-down and then recurses into the children of the node that the rule returns. Before this change, the `Sort` was a child of the returned `Limit`, so the same pass also visited it and applied the TopK-through-join pushdown (`push_topk_through_join`). Now the `Sort` is the returned node, so the traversal does not visit it again in this pass. The `skip == 0` branch therefore calls `push_topk_through_join` on the `Sort` directly. Without that call, `Limit -> Sort -> Left Join` would move the extra pass from the `Limit` removal to the TopK pushdown instead of removing it. The same was already true on `main` for a `Sort` that already had a fetch, so that case also gets the TopK pushdown in the same pass now. The branch always returns `Transformed::yes`, also when `push_topk_through_join` returns `Transformed::no`, because the `Limit` is removed in every case. ## What is the testing strategy for this PR? - `push_down_limit::test::limit_push_down_sort`: snapshot updated. After one pass the plan is `Sort: ..., fetch=10` without the `Limit` above it. - New `push_down_limit::test::limit_push_down_sort_settles_in_first_pass`: runs `PushDownLimit` with `max_passes = 3` and asserts that the optimizer stops after 2 passes (pass 1 leaves the plan unchanged). On `main` it runs 3. - `topk_through_join::test::topk_pushed_through_limit_then_sort_with_two_passes` is renamed to `topk_pushed_through_limit_then_sort_in_one_pass` and now uses `max_passes = 1`, with the same expected plan. It fails if the new branch does not call `push_topk_through_join`. - New `topk_through_join::test::topk_pushed_through_limit_over_sort_with_{same,smaller}_fetch_in_one_pass`: `Limit` over a `Sort` that already has `fetch=3`, over a Left Join, with `max_passes = 1`. The TopK moves below the join in the one pass. Both fail on the earlier code. - New `push_down_limit::test::limit_push_down_sort_reports_transformed`: calls `PushDownLimit::rewrite` on `Limit -> Sort -> TableScan` (no join) and asserts `transformed == true`. It fails if the branch returns the `Transformed` from `push_topk_through_join` as it is. - `user_defined_plan.rs`: the example `TopKOptimizerRule` also matches a `Sort` with a fetch (see below). The `topk_plan` and `topk_invariants*` tests failed without that change. I checked that the optimizer tests fail without the fix. Results: - `cargo test --profile ci -p datafusion-optimizer`: 901 + 26 + 5 passed, 0 failed. - `cargo test --profile ci -p datafusion-sqllogictest --test sqllogictests`: 522 of 522 files pass. No expected plan changes, because the final plans are the same. - `cargo fmt --all -- --check` and `cargo clippy --profile ci -p datafusion-optimizer --all-targets -- -D warnings`: clean. ## Are there any user-facing changes? Final plans are the same. Planning an `ORDER BY ... LIMIT` query without an `OFFSET` takes one less optimizer pass. A custom optimizer rule that runs after `PushDownLimit` and matches `Limit` over `Sort` no longer sees that pattern in the first pass. It sees a `Sort` with a fetch instead. On `main` such a rule also saw only the `Sort` with a fetch from the second pass on. The `TopKOptimizerRule` example in `datafusion/core/tests/user_defined/user_defined_plan.rs` had this dependency, so it now also matches a `Sort` with a fetch. --------- Co-authored-by: Claude Opus 5.5 <[email protected]>
1 parent 88306ff commit 5719930

3 files changed

Lines changed: 136 additions & 22 deletions

File tree

‎datafusion/core/tests/user_defined/user_defined_plan.rs‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -510,9 +510,30 @@ impl OptimizerRule for TopKOptimizerRule {
510510
plan: LogicalPlan,
511511
_config: &dyn OptimizerConfig,
512512
) -> Result<Transformed<LogicalPlan>, DataFusionError> {
513-
// Note: this code simply looks for the pattern of a Limit followed by a
514-
// Sort and replaces it by a TopK node. It does not handle many
515-
// edge cases (e.g multiple sort columns, sort ASC / DESC), etc.
513+
// Note: this code simply looks for a Sort with a fetch, or a Limit
514+
// followed by a Sort, and replaces it by a TopK node. It does not
515+
// handle many edge cases (e.g multiple sort columns, sort ASC / DESC),
516+
// etc.
517+
//
518+
// `PushDownLimit` gives a Sort below a Limit the Limit's fetch and
519+
// removes the Limit, so a Sort with a fetch is the usual shape.
520+
if let LogicalPlan::Sort(Sort {
521+
expr,
522+
input,
523+
fetch: Some(fetch),
524+
}) = &plan
525+
&& expr.len() == 1
526+
{
527+
return Ok(Transformed::yes(LogicalPlan::Extension(Extension {
528+
node: Arc::new(TopKPlanNode {
529+
k: *fetch,
530+
input: input.as_ref().clone(),
531+
expr: expr[0].clone(),
532+
invariant_mock: self.invariant_mock.clone(),
533+
}),
534+
})));
535+
}
536+
516537
let LogicalPlan::Limit(ref limit) = plan else {
517538
return Ok(Transformed::no(plan));
518539
};

‎datafusion/optimizer/src/push_down_limit.rs‎

Lines changed: 55 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -163,16 +163,24 @@ fn rewrite_limit(mut limit: Limit) -> Result<Transformed<LogicalPlan>> {
163163
let sort_fetch = skip + fetch;
164164
Some(sort.fetch.map(|f| f.min(sort_fetch)).unwrap_or(sort_fetch))
165165
};
166-
if new_fetch == sort.fetch {
167-
if skip > 0 {
168-
original_limit(skip, fetch, LogicalPlan::Sort(sort))
166+
let fetch_changed = new_fetch != sort.fetch;
167+
sort.fetch = new_fetch;
168+
if skip > 0 {
169+
if fetch_changed {
170+
limit.input = Arc::new(LogicalPlan::Sort(sort));
171+
Ok(Transformed::yes(LogicalPlan::Limit(limit)))
169172
} else {
170-
Ok(Transformed::yes(LogicalPlan::Sort(sort)))
173+
original_limit(skip, fetch, LogicalPlan::Sort(sort))
171174
}
172175
} else {
173-
sort.fetch = new_fetch;
174-
limit.input = Arc::new(LogicalPlan::Sort(sort));
175-
Ok(Transformed::yes(LogicalPlan::Limit(limit)))
176+
// With `skip = 0` the Sort's fetch already caps the output, so
177+
// drop the Limit now instead of on the next pass. The Sort then
178+
// replaces the Limit as the visited node and its own visit is
179+
// skipped, so apply the TopK pushdown here. The plan changed
180+
// (the Limit is gone) even if the TopK does not move.
181+
let mut result = push_topk_through_join(sort)?;
182+
result.transformed = true;
183+
Ok(result)
176184
}
177185
}
178186
LogicalPlan::Projection(mut proj) => {
@@ -660,17 +668,53 @@ mod test {
660668
.limit(0, Some(10))?
661669
.build()?;
662670

663-
// Should push down limit to sort
671+
// Should push down limit to sort, and drop the now redundant limit
664672
assert_optimized_plan_equal!(
665673
plan,
666674
@r"
667-
Limit: skip=0, fetch=10
668-
Sort: test.a ASC NULLS LAST, fetch=10
669-
TableScan: test
675+
Sort: test.a ASC NULLS LAST, fetch=10
676+
TableScan: test
670677
"
671678
)
672679
}
673680

681+
/// `Limit(skip=0) -> Sort` settles in the first pass: the second pass
682+
/// leaves the plan unchanged, so the optimizer stops after 2 passes.
683+
#[test]
684+
fn limit_push_down_sort_settles_in_first_pass() -> Result<()> {
685+
let table_scan = test_table_scan()?;
686+
687+
let plan = LogicalPlanBuilder::from(table_scan)
688+
.sort_by(vec![col("a")])?
689+
.limit(0, Some(10))?
690+
.build()?;
691+
692+
let optimizer_ctx = OptimizerContext::new().with_max_passes(3);
693+
let optimizer =
694+
crate::Optimizer::with_rules(vec![Arc::new(PushDownLimit::new())]);
695+
let mut passes = 0;
696+
optimizer.optimize(plan, &optimizer_ctx, |_, _| passes += 1)?;
697+
assert_eq!(passes, 2);
698+
Ok(())
699+
}
700+
701+
/// Dropping the `Limit` above a `Sort` changes the plan even when the
702+
/// TopK has no join to move through, so the rule must report it.
703+
#[test]
704+
fn limit_push_down_sort_reports_transformed() -> Result<()> {
705+
for sort_fetch in [None, Some(10), Some(5)] {
706+
let plan = LogicalPlanBuilder::from(test_table_scan()?)
707+
.sort_with_limit(vec![col("a").sort(true, false)], sort_fetch)?
708+
.limit(0, Some(10))?
709+
.build()?;
710+
711+
let result = PushDownLimit::new().rewrite(plan, &OptimizerContext::new())?;
712+
assert!(result.transformed, "sort fetch {sort_fetch:?}");
713+
assert!(matches!(result.data, LogicalPlan::Sort(_)));
714+
}
715+
Ok(())
716+
}
717+
674718
#[test]
675719
fn limit_push_down_sort_skip() -> Result<()> {
676720
let table_scan = test_table_scan()?;

‎datafusion/optimizer/src/push_down_limit/topk_through_join.rs‎

Lines changed: 57 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1380,14 +1380,12 @@ mod test {
13801380
}
13811381

13821382
/// `Limit(skip=0) → Sort(no fetch) → Join`: `rewrite_limit` merges the
1383-
/// Limit's fetch into the Sort on one visit, but only strips the now-
1384-
/// redundant `Limit(skip=0)` wrapper on a *later* visit to that same
1385-
/// node (it only recognizes the wrapper as redundant once the Sort
1386-
/// already carries the matching fetch). Confirms the TopK-through-join
1387-
/// pushdown and the Limit/Sort merge both still converge to the fully
1388-
/// collapsed form with two passes.
1383+
/// Limit's fetch into the Sort and drops the redundant `Limit(skip=0)`
1384+
/// in the same visit. Confirms that the TopK-through-join pushdown
1385+
/// still applies to that Sort, so one pass reaches the fully collapsed
1386+
/// form.
13891387
#[test]
1390-
fn topk_pushed_through_limit_then_sort_with_two_passes() -> Result<()> {
1388+
fn topk_pushed_through_limit_then_sort_in_one_pass() -> Result<()> {
13911389
let t1 = test_table_scan_with_name("t1")?;
13921390
let t2 = test_table_scan_with_name("t2")?;
13931391

@@ -1402,7 +1400,7 @@ mod test {
14021400
.limit(0, Some(3))?
14031401
.build()?;
14041402

1405-
let optimizer_ctx = OptimizerContext::new().with_max_passes(2);
1403+
let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
14061404
let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> =
14071405
vec![Arc::new(PushDownLimit::new())];
14081406
assert_optimized_plan_eq_snapshot!(
@@ -1419,6 +1417,57 @@ mod test {
14191417
)
14201418
}
14211419

1420+
/// A `Limit` over a `Sort` that already has an equal or smaller fetch is
1421+
/// dropped, and the `Sort` replaces it as the visited node. The TopK must
1422+
/// still move through the join in the same pass.
1423+
fn optimize_limit_over_sort_with_fetch(limit_fetch: usize) -> Result<LogicalPlan> {
1424+
let t1 = test_table_scan_with_name("t1")?;
1425+
let t2 = test_table_scan_with_name("t2")?;
1426+
1427+
let plan = LogicalPlanBuilder::from(t1)
1428+
.join(
1429+
LogicalPlanBuilder::from(t2).build()?,
1430+
JoinType::Left,
1431+
(vec!["a"], vec!["a"]),
1432+
None,
1433+
)?
1434+
.sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
1435+
.limit(0, Some(limit_fetch))?
1436+
.build()?;
1437+
1438+
let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
1439+
crate::Optimizer::with_rules(vec![Arc::new(PushDownLimit::new())]).optimize(
1440+
plan,
1441+
&optimizer_ctx,
1442+
|_, _| {},
1443+
)
1444+
}
1445+
1446+
#[test]
1447+
fn topk_pushed_through_limit_over_sort_with_same_fetch_in_one_pass() -> Result<()> {
1448+
insta::assert_snapshot!(optimize_limit_over_sort_with_fetch(3)?, @r"
1449+
Sort: t1.b ASC NULLS LAST, fetch=3
1450+
Left Join: t1.a = t2.a
1451+
Sort: t1.b ASC NULLS LAST, fetch=3
1452+
TableScan: t1
1453+
TableScan: t2
1454+
");
1455+
Ok(())
1456+
}
1457+
1458+
#[test]
1459+
fn topk_pushed_through_limit_over_sort_with_smaller_fetch_in_one_pass() -> Result<()>
1460+
{
1461+
insta::assert_snapshot!(optimize_limit_over_sort_with_fetch(5)?, @r"
1462+
Sort: t1.b ASC NULLS LAST, fetch=3
1463+
Left Join: t1.a = t2.a
1464+
Sort: t1.b ASC NULLS LAST, fetch=3
1465+
TableScan: t1
1466+
TableScan: t2
1467+
");
1468+
Ok(())
1469+
}
1470+
14221471
/// Running the optimizer a second time on its own output must be a
14231472
/// no-op: a rule that keeps rewriting an already-optimized plan (e.g.
14241473
/// oscillating between two equivalent forms) would never let the

0 commit comments

Comments
 (0)