initial work

Signed-off-by: Adam Gutglick <adamgsal@gmail.com>

Author: AdamGS

datafusion/core/tests/physical_optimizer/ensure_requirements.rs

@@ -1536,15 +1536,10 @@ fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result
 // Limits with a `skip`
 // ========================================================================
 
-/// A sort pushed below `GlobalLimitExec` with a non-zero `skip` must ask its
-/// input for `skip + fetch` rows; asking for only `fetch` rows used to leave
-/// `LIMIT 10 OFFSET 5` with 5 result rows.
-///
-/// This checks a single pass only: the sort is inserted by `pushdown_sorts`,
-/// which runs after `parallelize_sorts`, so a second pass would additionally
-/// parallelize it into `SortPreservingMergeExec` + partitioned `SortExec`.
+/// Keep sorting above an unordered limit so its input can stop early and only
+/// the rows retained after the offset need sorting.
 #[test]
-fn test_sort_pushed_below_limit_with_skip_keeps_skip_rows() -> Result<()> {
+fn test_keep_sort_above_unordered_limit_with_skip() -> Result<()> {
     let source = Arc::new(MockMultiPartitionExec::new(4));
     let coalesce = Arc::new(CoalescePartitionsExec::new(source));
     let limit = Arc::new(GlobalLimitExec::new(coalesce, 5, Some(10)));
@@ -1553,10 +1548,36 @@ fn test_sort_pushed_below_limit_with_skip_keeps_skip_rows() -> Result<()> {
 
     let optimized = optimize_and_sanity_check(sort)?;
     assert_snapshot!(plan_string(&optimized), @r"
-    GlobalLimitExec: skip=5, fetch=10
-      SortExec: TopK(fetch=15), expr=[a@0 DESC], preserve_partitioning=[false]
+    SortExec: expr=[a@0 DESC], preserve_partitioning=[false]
+      GlobalLimitExec: skip=5, fetch=10
         CoalescePartitionsExec
           MockMultiPartitionExec
     ");
     Ok(())
 }
+
+/// Refining an existing ordering can still push a TopK below the limit. It must
+/// retain `skip + fetch` rows so the offset does not reduce the result count.
+#[test]
+fn test_sort_pushed_below_ordered_limit_with_skip_keeps_skip_rows() -> Result<()> {
+    let source = Arc::new(MockMultiPartitionExec::new(4));
+    let mut ordering = sort_expr_on("a", 0, false, false);
+    let ordered_input = Arc::new(MockReqExec::new(
+        source,
+        Distribution::SinglePartition,
+        Some(ordering.clone()),
+    ));
+    let limit = Arc::new(GlobalLimitExec::new(ordered_input, 5, Some(10)));
+    ordering.extend(sort_expr_on("b", 1, false, false));
+    let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(ordering, limit));
+
+    let optimized = optimize_and_sanity_check(sort)?;
+    assert_snapshot!(plan_string(&optimized), @r"
+    GlobalLimitExec: skip=5, fetch=10
+      SortExec: TopK(fetch=15), expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false], sort_prefix=[a@0 ASC NULLS LAST]
+        MockReqExec
+          SortPreservingMergeExec: [a@0 ASC NULLS LAST]
+            MockMultiPartitionExec
+    ");
+    Ok(())
+}

datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs

@@ -566,8 +566,10 @@ fn pushdown_requirement_to_children(
             } else {
                 parent_required.clone()
             };
+
+        // Keep sorting above an unordered limit so the input can stop early.
         let Some(ordering) = plan.properties().output_ordering() else {
-            return Ok(Some(vec![Some(child_required)]));
+            return Ok(None);
         };
         if plan.properties().eq_properties.requirements_compatible(
             parent_required.first().clone(),

datafusion/sqllogictest/test_files/limit.slt

@@ -339,9 +339,8 @@ SELECT COUNT(*) FROM (SELECT a FROM t1 LIMIT 3 OFFSET 8);
 ----
 2
 
-# A sort above LIMIT ... OFFSET is pushed below the limit as a TopK. The TopK
-# has to keep `skip + fetch` rows (3 + 4 = 7) so that the limit can still skip
-# 3 rows and return 4; keeping only `fetch` rows returned a single row.
+# Keep the sort above an unordered LIMIT ... OFFSET so the scan can stop
+# after `skip + fetch` rows (3 + 4 = 7), and only the 4 selected rows are sorted.
 query TT
 EXPLAIN SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a;
 ----
@@ -350,9 +349,9 @@ logical_plan
 02)--Limit: skip=3, fetch=4
 03)----TableScan: t1 projection=[a], fetch=7
 physical_plan
-01)GlobalLimitExec: skip=3, fetch=4
-02)--SortExec: TopK(fetch=7), expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false]
-03)----DataSourceExec: partitions=1, partition_sizes=[1]
+01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false]
+02)--GlobalLimitExec: skip=3, fetch=4
+03)----DataSourceExec: partitions=1, partition_sizes=[1], fetch=7
 
 query I
 SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a;