fix: sort pushed below GlobalLimitExec ignores skip, returning too few rows

datafusion/core/tests/physical_optimizer/ensure_requirements.rs

@@ -1531,3 +1531,32 @@ fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result
 
     Ok(())
 }
+
+// ========================================================================
+// 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`.
+#[test]
+fn test_sort_pushed_below_limit_with_skip_keeps_skip_rows() -> 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)));
+    let sort: Arc<dyn ExecutionPlan> =
+        Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), 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 DESC], preserve_partitioning=[false]
+        CoalescePartitionsExec
+          MockMultiPartitionExec
+    ");
+    Ok(())
+}

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

@@ -79,6 +79,22 @@ impl Default for ParentRequirements {
 
 pub type SortPushDown = PlanContext<ParentRequirements>;
 
+/// Number of input rows `plan` needs from its children in order to produce
+/// its own `fetch` rows. This is `plan.fetch()` for every operator except
+/// [`GlobalLimitExec`], which discards `skip` rows first and therefore needs
+/// `skip + fetch` input rows. Using the bare `fetch` there would turn
+/// `LIMIT 10 OFFSET 5` into `TopK(10)` below the limit, i.e. 5 result rows.
+///
+/// Note this is distinct from the fetch a parent imposes on `plan`'s *output*
+/// (`ParentRequirements::fetch`), for which `plan.fetch()` is the right bound.
+fn input_fetch(plan: &Arc<dyn ExecutionPlan>) -> Option<usize> {
+    let fetch = plan.fetch()?;
+    let skip = plan
+        .downcast_ref::<GlobalLimitExec>()
+        .map_or(0, |limit| limit.skip());
+    Some(fetch + skip)
+}
+
 /// Assigns the ordering requirement of the root node to the its children.
 pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) {
     let reqs = sort_push_down.plan.required_input_ordering();
@@ -364,7 +380,7 @@ fn pushdown_sorts_helper(
         // For operators that can take a sort pushdown, continue with updated
         // requirements. If this node already outputs single partition (e.g. SPM),
         // don't push SinglePartition to children.
-        let current_fetch = sort_push_down.plan.fetch();
+        let current_fetch = input_fetch(&sort_push_down.plan);
         let dists = sort_push_down
             .plan
             .input_distribution_requirements()

datafusion/sqllogictest/test_files/limit.slt

@@ -339,6 +339,29 @@ 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.
+query TT
+EXPLAIN SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a;
+----
+logical_plan
+01)Sort: t1.a ASC NULLS LAST
+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]
+
+query I
+SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a;
+----
+4
+5
+6
+7
+
 # The aggregate does not need to be computed because the input statistics are exact and
 # an OFFSET, but no LIMIT, is specified.
 query TT