Merge remote-tracking branch 'origin/main' into neilc/perf-builder-replace

Author: neilconway

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(