Fix PushdownSort dropping LIMIT when eliminating SortExec (#21744)

## Which issue does this PR close?

When `PushdownSort` removes a `SortExec` because a source returns
`Exact` (guaranteeing ordering), any `fetch` (LIMIT) on the `SortExec`
is silently dropped if the underlying plan does not support
`with_fetch()`.

For example, `ProjectionExec` supports `try_pushdown_sort` (delegating
to its child) but does not implement `with_fetch()`. A plan like
`SortExec(fetch=10) → ProjectionExec → source` that gets sort-eliminated
loses the limit.

## What changes are included in this PR?

In the `Exact` branch of `PushdownSort`, when the eliminated `SortExec`
carried a `fetch`:

1. Try `with_fetch()` on the pushed-down source first
2. If `with_fetch()` returns `None`, fall back to wrapping with
`GlobalLimitExec`

## Are these changes tested?

Yes. Three new unit tests:

- `test_sort_pushdown_exact_no_fetch_no_limit` — Exact elimination
without fetch: no limit wrapper added
- `test_sort_pushdown_exact_preserves_fetch_with_global_limit` — Exact
elimination with fetch, source does NOT support `with_fetch()`:
`GlobalLimitExec` wrapper added
- `test_sort_pushdown_exact_preserves_fetch_with_source_support` — Exact
elimination with fetch, source supports `with_fetch()`: limit pushed
into source directly

## Are there any user-facing changes?

No.

Author: sgrebnov

datafusion/core/tests/physical_optimizer/pushdown_sort.rs

@@ -32,10 +32,10 @@ use datafusion_physical_optimizer::pushdown_sort::PushdownSort;
 use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
-    OptimizationTest, coalesce_partitions_exec, parquet_exec, parquet_exec_with_sort,
-    projection_exec, projection_exec_with_alias, repartition_exec, schema,
-    simple_projection_exec, sort_exec, sort_exec_with_fetch, sort_expr, sort_expr_named,
-    test_scan_with_ordering,
+    OptimizationTest, TestScan, coalesce_partitions_exec, parquet_exec,
+    parquet_exec_with_sort, projection_exec, projection_exec_with_alias,
+    repartition_exec, schema, simple_projection_exec, sort_exec, sort_exec_with_fetch,
+    sort_expr, sort_expr_named, test_scan_with_ordering,
 };
 
 #[test]
@@ -996,3 +996,91 @@ fn test_sort_pushdown_with_test_scan_arbitrary_ordering() {
     "
     );
 }
+
+// ============================================================================
+// EXACT PUSHDOWN TESTS (source guarantees ordering, SortExec removed)
+// ============================================================================
+
+#[test]
+fn test_sort_pushdown_exact_no_fetch_no_limit() {
+    // When a source returns Exact (without fetch), the SortExec should be
+    // removed entirely with no GlobalLimitExec wrapper.
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let b = sort_expr("b", &schema);
+    let source =
+        Arc::new(TestScan::new(schema.clone(), vec![]).with_exact_pushdown(true));
+
+    let ordering = LexOrdering::new(vec![a, b.reverse()]).unwrap();
+    let plan = sort_exec(ordering, source);
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: expr=[a@0 ASC, b@1 DESC NULLS LAST], preserve_partitioning=[false]
+        -   TestScan
+      output:
+        Ok:
+          - TestScan: requested_ordering=[a@0 ASC, b@1 DESC NULLS LAST]
+    "
+    );
+}
+
+#[test]
+fn test_sort_pushdown_exact_preserves_fetch_with_global_limit() {
+    // When a source returns Exact but does NOT support with_fetch(),
+    // the optimizer must wrap the result with GlobalLimitExec to preserve
+    // the LIMIT from the eliminated SortExec.
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let source =
+        Arc::new(TestScan::new(schema.clone(), vec![]).with_exact_pushdown(true));
+
+    let ordering = LexOrdering::new(vec![a]).unwrap();
+    let plan = sort_exec_with_fetch(ordering, Some(10), source);
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[false]
+        -   TestScan
+      output:
+        Ok:
+          - GlobalLimitExec: skip=0, fetch=10
+          -   TestScan: requested_ordering=[a@0 ASC]
+    "
+    );
+}
+
+#[test]
+fn test_sort_pushdown_exact_preserves_fetch_with_source_support() {
+    // When a source returns Exact AND supports with_fetch(),
+    // the limit should be pushed into the source directly (no GlobalLimitExec).
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let source = Arc::new(
+        TestScan::new(schema.clone(), vec![])
+            .with_exact_pushdown(true)
+            .with_supports_fetch(true),
+    );
+
+    let ordering = LexOrdering::new(vec![a]).unwrap();
+    let plan = sort_exec_with_fetch(ordering, Some(10), source);
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[false]
+        -   TestScan
+      output:
+        Ok:
+          - TestScan: requested_ordering=[a@0 ASC], fetch=10
+    "
+    );
+}

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -830,15 +830,30 @@ pub fn sort_expr_named(name: &str, index: usize) -> PhysicalSortExpr {
     }
 }
 
-/// A test data source that can display any requested ordering
-/// This is useful for testing sort pushdown behavior
+/// A test data source that can display any requested ordering.
+/// This is useful for testing sort pushdown behavior.
+///
+/// ## Configuration
+///
+/// - `exact_pushdown`: if `true`, `try_pushdown_sort` returns `Exact`
+///   (source guarantees ordering, SortExec can be removed); if `false`
+///   (default), returns `Inexact` (SortExec kept).
+/// - `supports_fetch`: if `true`, `with_fetch()` returns `Some` so the
+///   optimizer can push a LIMIT into the source; if `false` (default),
+///   `with_fetch()` returns `None`, forcing a `GlobalLimitExec` wrapper.
 #[derive(Debug, Clone)]
 pub struct TestScan {
     schema: SchemaRef,
     output_ordering: Vec<LexOrdering>,
     plan_properties: Arc<PlanProperties>,
     // Store the requested ordering for display
     requested_ordering: Option<LexOrdering>,
+    /// If true, `try_pushdown_sort` returns `Exact` instead of `Inexact`.
+    exact_pushdown: bool,
+    /// If true, `with_fetch()` returns `Some(...)` (source absorbs the limit).
+    supports_fetch: bool,
+    /// The fetch (LIMIT) value pushed into this scan via `with_fetch()`.
+    fetch: Option<usize>,
 }
 
 impl TestScan {
@@ -872,41 +887,60 @@ impl TestScan {
             output_ordering,
             plan_properties: Arc::new(plan_properties),
             requested_ordering: None,
+            exact_pushdown: false,
+            supports_fetch: false,
+            fetch: None,
         }
     }
 
     /// Create a TestScan with a single output ordering
     pub fn with_ordering(schema: SchemaRef, ordering: LexOrdering) -> Self {
         Self::new(schema, vec![ordering])
     }
+
+    /// Set whether `try_pushdown_sort` returns `Exact` (true) or `Inexact` (false).
+    pub fn with_exact_pushdown(mut self, exact: bool) -> Self {
+        self.exact_pushdown = exact;
+        self
+    }
+
+    /// Set whether `with_fetch()` returns `Some` (true) or `None` (false).
+    pub fn with_supports_fetch(mut self, supports: bool) -> Self {
+        self.supports_fetch = supports;
+        self
+    }
 }
 
 impl DisplayAs for TestScan {
     fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
         match t {
             DisplayFormatType::Default | DisplayFormatType::Verbose => {
                 write!(f, "TestScan")?;
+                let mut sep = ": ";
                 if !self.output_ordering.is_empty() {
-                    write!(f, ": output_ordering=[")?;
-                    // Format the ordering in a readable way
+                    write!(f, "{sep}output_ordering=[")?;
                     for (i, sort_expr) in self.output_ordering[0].iter().enumerate() {
                         if i > 0 {
                             write!(f, ", ")?;
                         }
                         write!(f, "{sort_expr}")?;
                     }
                     write!(f, "]")?;
+                    sep = ", ";
                 }
-                // This is the key part - show what ordering was requested
                 if let Some(ref req) = self.requested_ordering {
-                    write!(f, ", requested_ordering=[")?;
+                    write!(f, "{sep}requested_ordering=[")?;
                     for (i, sort_expr) in req.iter().enumerate() {
                         if i > 0 {
                             write!(f, ", ")?;
                         }
                         write!(f, "{sort_expr}")?;
                     }
                     write!(f, "]")?;
+                    sep = ", ";
+                }
+                if let Some(fetch) = self.fetch {
+                    write!(f, "{sep}fetch={fetch}")?;
                 }
                 Ok(())
             }
@@ -953,6 +987,20 @@ impl ExecutionPlan for TestScan {
         Ok(Arc::new(Statistics::new_unknown(&self.schema)))
     }
 
+    fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
+        if self.supports_fetch {
+            let mut new_scan = self.clone();
+            new_scan.fetch = fetch;
+            Some(Arc::new(new_scan))
+        } else {
+            None
+        }
+    }
+
+    fn fetch(&self) -> Option<usize> {
+        self.fetch
+    }
+
     // This is the key method - implement sort pushdown
     fn try_pushdown_sort(
         &self,
@@ -965,10 +1013,27 @@ impl ExecutionPlan for TestScan {
         let mut new_scan = self.clone();
         new_scan.requested_ordering = requested_ordering;
 
-        // Always return Inexact to keep the Sort node (like Phase 1 behavior)
-        Ok(SortOrderPushdownResult::Inexact {
-            inner: Arc::new(new_scan),
-        })
+        if self.exact_pushdown {
+            // Update plan properties to reflect the guaranteed ordering
+            let orderings: Vec<Vec<PhysicalSortExpr>> = vec![order.to_vec()];
+            let eq_properties = EquivalenceProperties::new_with_orderings(
+                Arc::clone(&self.schema),
+                orderings,
+            );
+            new_scan.plan_properties = Arc::new(PlanProperties::new(
+                eq_properties,
+                Partitioning::UnknownPartitioning(1),
+                EmissionType::Incremental,
+                Boundedness::Bounded,
+            ));
+            Ok(SortOrderPushdownResult::Exact {
+                inner: Arc::new(new_scan),
+            })
+        } else {
+            Ok(SortOrderPushdownResult::Inexact {
+                inner: Arc::new(new_scan),
+            })
+        }
     }
 
     fn apply_expressions(

datafusion/physical-optimizer/src/pushdown_sort.rs

@@ -61,6 +61,7 @@ use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::SortOrderPushdownResult;
 use datafusion_physical_plan::buffer::BufferExec;
+use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use std::sync::Arc;
@@ -104,8 +105,12 @@ impl PhysicalOptimizerRule for PushdownSort {
                 let required_ordering = sort_child.expr();
                 match sort_input.try_pushdown_sort(required_ordering)? {
                     SortOrderPushdownResult::Exact { inner } => {
+                        // Preserve fetch (LIMIT) from the eliminated SortExec.
+                        // Use LocalLimitExec (not Global) since input is multi-partition.
                         let inner = if let Some(fetch) = sort_child.fetch() {
-                            inner.with_fetch(Some(fetch)).unwrap_or(inner)
+                            inner.with_fetch(Some(fetch)).unwrap_or_else(|| {
+                                Arc::new(LocalLimitExec::new(inner, fetch))
+                            })
                         } else {
                             inner
                         };
@@ -149,14 +154,19 @@ impl PhysicalOptimizerRule for PushdownSort {
             match sort_input.try_pushdown_sort(required_ordering)? {
                 SortOrderPushdownResult::Exact { inner } => {
                     // Data source guarantees perfect ordering - remove the Sort operator.
-                    // Preserve the fetch (LIMIT) from the original SortExec so the
-                    // data source can stop reading early.
-                    let inner = if let Some(fetch) = sort_exec.fetch() {
-                        inner.with_fetch(Some(fetch)).unwrap_or(inner)
+                    //
+                    // If the SortExec carried a fetch (LIMIT), we must preserve it.
+                    // First try pushing the limit into the source via `with_fetch()`.
+                    // If the source doesn't support `with_fetch`, fall back to
+                    // wrapping with GlobalLimitExec.
+                    if let Some(fetch) = sort_exec.fetch() {
+                        let inner = inner.with_fetch(Some(fetch)).unwrap_or_else(|| {
+                            Arc::new(GlobalLimitExec::new(inner, 0, Some(fetch)))
+                        });
+                        Ok(Transformed::yes(inner))
                     } else {
-                        inner
-                    };
-                    Ok(Transformed::yes(inner))
+                        Ok(Transformed::yes(inner))
+                    }
                 }
                 SortOrderPushdownResult::Inexact { inner } => {
                     // Data source is optimized for the ordering but not perfectly sorted