fix: preserve scan order during sort-preserving merge limit pushdown

datafusion/core/tests/physical_optimizer/limit_pushdown.rs

@@ -19,13 +19,16 @@ use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
     coalesce_partitions_exec, global_limit_exec, hash_join_exec, local_limit_exec,
-    sort_exec, sort_preserving_merge_exec, stream_exec,
+    parquet_exec, parquet_exec_with_sort, sort_exec, sort_preserving_merge_exec,
+    sort_preserving_merge_exec_with_fetch, stream_exec,
 };
 
 use arrow::compute::SortOptions;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
+use datafusion_datasource::file_scan_config::FileScanConfig;
+use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr::{JoinType, Operator};
 use datafusion_physical_expr::Partitioning;
 use datafusion_physical_expr::expressions::{BinaryExpr, col, lit};
@@ -405,6 +408,64 @@ fn pushes_global_limit_into_multiple_fetch_plans() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn preserves_order_when_pushing_fetch_from_sort_preserving_merge() -> Result<()> {
+    let schema = create_schema();
+    let ordering: LexOrdering = [PhysicalSortExpr {
+        expr: col("c1", &schema)?,
+        options: SortOptions::default(),
+    }]
+    .into();
+    let scan = parquet_exec_with_sort(schema, vec![ordering.clone()]);
+    let local_limit = local_limit_exec(scan, 10);
+    let plan = sort_preserving_merge_exec_with_fetch(ordering, local_limit, 5);
+
+    let optimized = LimitPushdown::new().optimize(plan, &ConfigOptions::new())?;
+    let scan = optimized.children().swap_remove(0);
+    let scan = scan
+        .downcast_ref::<DataSourceExec>()
+        .expect("fetch should be pushed into the scan");
+    let config = scan
+        .data_source()
+        .downcast_ref::<FileScanConfig>()
+        .expect("parquet scan should use FileScanConfig");
+
+    assert_eq!(config.limit, Some(5));
+    assert!(config.preserve_order);
+
+    Ok(())
+}
+
+#[test]
+fn does_not_preserve_order_for_independent_limit_below_sort() -> Result<()> {
+    let schema = create_schema();
+    let ordering: LexOrdering = [PhysicalSortExpr {
+        expr: col("c1", &schema)?,
+        options: SortOptions::default(),
+    }]
+    .into();
+    let scan = parquet_exec(schema);
+    let local_limit = local_limit_exec(scan, 10);
+    let sort = sort_exec(ordering.clone(), local_limit);
+    let plan = sort_preserving_merge_exec_with_fetch(ordering, sort, 5);
+
+    let optimized = LimitPushdown::new().optimize(plan, &ConfigOptions::new())?;
+    let sort = optimized.children().swap_remove(0);
+    let scan = sort.children().swap_remove(0);
+    let scan = scan
+        .downcast_ref::<DataSourceExec>()
+        .expect("inner fetch should be pushed into the scan");
+    let config = scan
+        .data_source()
+        .downcast_ref::<FileScanConfig>()
+        .expect("parquet scan should use FileScanConfig");
+
+    assert_eq!(config.limit, Some(10));
+    assert!(!config.preserve_order);
+
+    Ok(())
+}
+
 #[test]
 fn keeps_pushed_local_limit_exec_when_there_are_multiple_input_partitions() -> Result<()>
 {

datafusion/physical-optimizer/src/limit_pushdown.rs

@@ -160,7 +160,7 @@ pub fn pushdown_limit_helper(
         );
         global_state.skip = skip;
         global_state.fetch = fetch;
-        global_state.preserve_order = limit_info.preserve_order;
+        global_state.preserve_order |= limit_info.preserve_order;
         global_state.satisfied = false;
 
         if let Some(fetch) = fetch
@@ -196,8 +196,10 @@ pub fn pushdown_limit_helper(
     }
 
     // If we have a non-limit operator with fetch capability, update global
-    // state as necessary:
+    // state as necessary. A fetched ordered merge selects the leading rows in
+    // its ordering, so preserve that ordering when pushing the fetch down.
     if pushdown_plan.fetch().is_some() {
+        global_state.preserve_order |= pushdown_plan.is::<SortPreservingMergeExec>();
         if global_state.skip == 0 {
             global_state.satisfied = true;
         }
@@ -269,36 +271,33 @@ pub fn pushdown_limit_helper(
         // to add a limit or a fetch. If the plan is already satisfied, we will try
         // to add the fetch info and return the plan.
 
-        // There's no push down, change fetch & skip to default values:
+        // This operator materializes the current fetch. Clear its requirements before
+        // descending so they are not combined with an independent limit below.
         let global_skip = global_state.skip;
+        let global_preserve_order = std::mem::take(&mut global_state.preserve_order);
         global_state.fetch = None;
         global_state.skip = 0;
 
-        let maybe_fetchable = pushdown_plan.with_fetch(skip_and_fetch);
+        let maybe_fetchable = pushdown_plan.with_fetch(skip_and_fetch).map(|plan| {
+            if global_preserve_order {
+                plan.with_preserve_order(true).unwrap_or(plan)
+            } else {
+                plan
+            }
+        });
         if global_state.satisfied {
             if let Some(plan_with_fetch) = maybe_fetchable {
-                let plan_with_preserve_order = plan_with_fetch
-                    .with_preserve_order(global_state.preserve_order)
-                    .unwrap_or(plan_with_fetch);
-                Ok((Transformed::yes(plan_with_preserve_order), global_state))
+                Ok((Transformed::yes(plan_with_fetch), global_state))
             } else {
                 Ok((Transformed::no(pushdown_plan), global_state))
             }
         } else {
             global_state.satisfied = true;
             pushdown_plan = if let Some(plan_with_fetch) = maybe_fetchable {
-                let plan_with_preserve_order = plan_with_fetch
-                    .with_preserve_order(global_state.preserve_order)
-                    .unwrap_or(plan_with_fetch);
-
                 if global_skip > 0 {
-                    add_global_limit(
-                        plan_with_preserve_order,
-                        global_skip,
-                        Some(global_fetch),
-                    )
+                    add_global_limit(plan_with_fetch, global_skip, Some(global_fetch))
                 } else {
-                    plan_with_preserve_order
+                    plan_with_fetch
                 }
             } else {
                 add_limit(pushdown_plan, global_skip, global_fetch)

datafusion/sqllogictest/test_files/limit_pruning.slt

@@ -113,6 +113,79 @@ select a from fully_matched_limit where a >= 3 order by a limit 4;
 statement ok
 drop table fully_matched_limit;
 
+# A fetched SortPreservingMergeExec must preserve scan order when pushing down
+# its limit. Otherwise, limit pruning can skip the first partially matched row
+# group in favor of the following fully matched row group.
+statement ok
+CREATE TABLE ordered_limit_source(id INT, passes INT) AS VALUES
+  (1, 0),
+  (2, 1),
+  (3, 1),
+  (4, 1),
+  (5, 1),
+  (6, 1),
+  (101, 1),
+  (102, 1),
+  (103, 1);
+
+query I
+COPY (SELECT * FROM ordered_limit_source WHERE id < 100 ORDER BY id)
+TO 'test_files/scratch/limit_pruning/ordered_limit/a.parquet'
+STORED AS PARQUET
+OPTIONS (
+  'format.max_row_group_size' '3'
+);
+----
+6
+
+query I
+COPY (SELECT * FROM ordered_limit_source WHERE id >= 100 ORDER BY id)
+TO 'test_files/scratch/limit_pruning/ordered_limit/b.parquet'
+STORED AS PARQUET
+OPTIONS (
+  'format.max_row_group_size' '3'
+);
+----
+3
+
+statement ok
+drop table ordered_limit_source;
+
+statement ok
+set datafusion.execution.target_partitions = 2;
+
+statement ok
+CREATE EXTERNAL TABLE ordered_limit(id INT, passes INT)
+STORED AS PARQUET
+LOCATION 'test_files/scratch/limit_pruning/ordered_limit/'
+WITH ORDER (id ASC);
+
+statement ok
+set datafusion.explain.physical_plan_only = true;
+
+query TT
+explain select * from ordered_limit where passes = 1 order by id limit 3;
+----
+physical_plan
+01)SortPreservingMergeExec: [id@0 ASC NULLS LAST], fetch=3
+02)--DataSourceExec: <slt:ignore>limit=3, output_ordering=[id@0 ASC NULLS LAST], file_type=parquet<slt:ignore>
+
+query II
+select * from ordered_limit where passes = 1 order by id limit 3;
+----
+2 1
+3 1
+4 1
+
+statement ok
+drop table ordered_limit;
+
+statement ok
+reset datafusion.explain.physical_plan_only;
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
 # limit_pruned_row_groups=0 total → 0 matched
 # because of order by, scan needs to preserve sort, so limit pruning is disabled
 query TT