You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/25 11:56:35 UTC
(datafusion) branch main updated: fix: drop the redundant `Limit` above a `Sort` in the same `PushDownLimit` visit (#25672)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 5719930666 fix: drop the redundant `Limit` above a `Sort` in the same `PushDownLimit` visit (#25672)
5719930666 is described below
commit 5719930666fb590b318cd286cbb00cc473efe244
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Fri Sep 25 11:45:49 2026 +0000
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` (95bb0a0dfa):
```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 <no...@anthropic.com>
---
.../core/tests/user_defined/user_defined_plan.rs | 27 ++++++++-
datafusion/optimizer/src/push_down_limit.rs | 66 ++++++++++++++++++----
.../src/push_down_limit/topk_through_join.rs | 65 ++++++++++++++++++---
3 files changed, 136 insertions(+), 22 deletions(-)
diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs
index a1a8e3aa14..080b1112f3 100644
--- a/datafusion/core/tests/user_defined/user_defined_plan.rs
+++ b/datafusion/core/tests/user_defined/user_defined_plan.rs
@@ -510,9 +510,30 @@ impl OptimizerRule for TopKOptimizerRule {
plan: LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Transformed<LogicalPlan>, DataFusionError> {
- // Note: this code simply looks for the pattern of a Limit followed by a
- // Sort and replaces it by a TopK node. It does not handle many
- // edge cases (e.g multiple sort columns, sort ASC / DESC), etc.
+ // Note: this code simply looks for a Sort with a fetch, or a Limit
+ // followed by a Sort, and replaces it by a TopK node. It does not
+ // handle many edge cases (e.g multiple sort columns, sort ASC / DESC),
+ // etc.
+ //
+ // `PushDownLimit` gives a Sort below a Limit the Limit's fetch and
+ // removes the Limit, so a Sort with a fetch is the usual shape.
+ if let LogicalPlan::Sort(Sort {
+ expr,
+ input,
+ fetch: Some(fetch),
+ }) = &plan
+ && expr.len() == 1
+ {
+ return Ok(Transformed::yes(LogicalPlan::Extension(Extension {
+ node: Arc::new(TopKPlanNode {
+ k: *fetch,
+ input: input.as_ref().clone(),
+ expr: expr[0].clone(),
+ invariant_mock: self.invariant_mock.clone(),
+ }),
+ })));
+ }
+
let LogicalPlan::Limit(ref limit) = plan else {
return Ok(Transformed::no(plan));
};
diff --git a/datafusion/optimizer/src/push_down_limit.rs b/datafusion/optimizer/src/push_down_limit.rs
index 79c18fbdeb..8e07d9dc93 100644
--- a/datafusion/optimizer/src/push_down_limit.rs
+++ b/datafusion/optimizer/src/push_down_limit.rs
@@ -163,16 +163,24 @@ fn rewrite_limit(mut limit: Limit) -> Result<Transformed<LogicalPlan>> {
let sort_fetch = skip + fetch;
Some(sort.fetch.map(|f| f.min(sort_fetch)).unwrap_or(sort_fetch))
};
- if new_fetch == sort.fetch {
- if skip > 0 {
- original_limit(skip, fetch, LogicalPlan::Sort(sort))
+ let fetch_changed = new_fetch != sort.fetch;
+ sort.fetch = new_fetch;
+ if skip > 0 {
+ if fetch_changed {
+ limit.input = Arc::new(LogicalPlan::Sort(sort));
+ Ok(Transformed::yes(LogicalPlan::Limit(limit)))
} else {
- Ok(Transformed::yes(LogicalPlan::Sort(sort)))
+ original_limit(skip, fetch, LogicalPlan::Sort(sort))
}
} else {
- sort.fetch = new_fetch;
- limit.input = Arc::new(LogicalPlan::Sort(sort));
- Ok(Transformed::yes(LogicalPlan::Limit(limit)))
+ // With `skip = 0` the Sort's fetch already caps the output, so
+ // drop the Limit now instead of on the next pass. The Sort then
+ // replaces the Limit as the visited node and its own visit is
+ // skipped, so apply the TopK pushdown here. The plan changed
+ // (the Limit is gone) even if the TopK does not move.
+ let mut result = push_topk_through_join(sort)?;
+ result.transformed = true;
+ Ok(result)
}
}
LogicalPlan::Projection(mut proj) => {
@@ -660,17 +668,53 @@ mod test {
.limit(0, Some(10))?
.build()?;
- // Should push down limit to sort
+ // Should push down limit to sort, and drop the now redundant limit
assert_optimized_plan_equal!(
plan,
@r"
- Limit: skip=0, fetch=10
- Sort: test.a ASC NULLS LAST, fetch=10
- TableScan: test
+ Sort: test.a ASC NULLS LAST, fetch=10
+ TableScan: test
"
)
}
+ /// `Limit(skip=0) -> Sort` settles in the first pass: the second pass
+ /// leaves the plan unchanged, so the optimizer stops after 2 passes.
+ #[test]
+ fn limit_push_down_sort_settles_in_first_pass() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .sort_by(vec![col("a")])?
+ .limit(0, Some(10))?
+ .build()?;
+
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(3);
+ let optimizer =
+ crate::Optimizer::with_rules(vec![Arc::new(PushDownLimit::new())]);
+ let mut passes = 0;
+ optimizer.optimize(plan, &optimizer_ctx, |_, _| passes += 1)?;
+ assert_eq!(passes, 2);
+ Ok(())
+ }
+
+ /// Dropping the `Limit` above a `Sort` changes the plan even when the
+ /// TopK has no join to move through, so the rule must report it.
+ #[test]
+ fn limit_push_down_sort_reports_transformed() -> Result<()> {
+ for sort_fetch in [None, Some(10), Some(5)] {
+ let plan = LogicalPlanBuilder::from(test_table_scan()?)
+ .sort_with_limit(vec![col("a").sort(true, false)], sort_fetch)?
+ .limit(0, Some(10))?
+ .build()?;
+
+ let result = PushDownLimit::new().rewrite(plan, &OptimizerContext::new())?;
+ assert!(result.transformed, "sort fetch {sort_fetch:?}");
+ assert!(matches!(result.data, LogicalPlan::Sort(_)));
+ }
+ Ok(())
+ }
+
#[test]
fn limit_push_down_sort_skip() -> Result<()> {
let table_scan = test_table_scan()?;
diff --git a/datafusion/optimizer/src/push_down_limit/topk_through_join.rs b/datafusion/optimizer/src/push_down_limit/topk_through_join.rs
index b4141a5f3b..a0c699f45f 100644
--- a/datafusion/optimizer/src/push_down_limit/topk_through_join.rs
+++ b/datafusion/optimizer/src/push_down_limit/topk_through_join.rs
@@ -1380,14 +1380,12 @@ mod test {
}
/// `Limit(skip=0) → Sort(no fetch) → Join`: `rewrite_limit` merges the
- /// Limit's fetch into the Sort on one visit, but only strips the now-
- /// redundant `Limit(skip=0)` wrapper on a *later* visit to that same
- /// node (it only recognizes the wrapper as redundant once the Sort
- /// already carries the matching fetch). Confirms the TopK-through-join
- /// pushdown and the Limit/Sort merge both still converge to the fully
- /// collapsed form with two passes.
+ /// Limit's fetch into the Sort and drops the redundant `Limit(skip=0)`
+ /// in the same visit. Confirms that the TopK-through-join pushdown
+ /// still applies to that Sort, so one pass reaches the fully collapsed
+ /// form.
#[test]
- fn topk_pushed_through_limit_then_sort_with_two_passes() -> Result<()> {
+ fn topk_pushed_through_limit_then_sort_in_one_pass() -> Result<()> {
let t1 = test_table_scan_with_name("t1")?;
let t2 = test_table_scan_with_name("t2")?;
@@ -1402,7 +1400,7 @@ mod test {
.limit(0, Some(3))?
.build()?;
- let optimizer_ctx = OptimizerContext::new().with_max_passes(2);
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> =
vec![Arc::new(PushDownLimit::new())];
assert_optimized_plan_eq_snapshot!(
@@ -1419,6 +1417,57 @@ mod test {
)
}
+ /// A `Limit` over a `Sort` that already has an equal or smaller fetch is
+ /// dropped, and the `Sort` replaces it as the visited node. The TopK must
+ /// still move through the join in the same pass.
+ fn optimize_limit_over_sort_with_fetch(limit_fetch: usize) -> Result<LogicalPlan> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .limit(0, Some(limit_fetch))?
+ .build()?;
+
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+ crate::Optimizer::with_rules(vec![Arc::new(PushDownLimit::new())]).optimize(
+ plan,
+ &optimizer_ctx,
+ |_, _| {},
+ )
+ }
+
+ #[test]
+ fn topk_pushed_through_limit_over_sort_with_same_fetch_in_one_pass() -> Result<()> {
+ insta::assert_snapshot!(optimize_limit_over_sort_with_fetch(3)?, @r"
+ 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
+ ");
+ Ok(())
+ }
+
+ #[test]
+ fn topk_pushed_through_limit_over_sort_with_smaller_fetch_in_one_pass() -> Result<()>
+ {
+ insta::assert_snapshot!(optimize_limit_over_sort_with_fetch(5)?, @r"
+ 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
+ ");
+ Ok(())
+ }
+
/// Running the optimizer a second time on its own output must be a
/// no-op: a rule that keeps rewriting an already-optimized plan (e.g.
/// oscillating between two equivalent forms) would never let the
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org