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

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));
         };

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()?;

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