fix: preserve required ordering in limit pushdown

datafusion/core/tests/physical_optimizer/limit_pushdown.rs

@@ -19,11 +19,12 @@ 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, sort_exec, sort_preserving_merge_exec, stream_exec,
 };
 
 use arrow::compute::SortOptions;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::datasource::source::DataSourceExec;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
 use datafusion_expr::{JoinType, Operator};
@@ -36,6 +37,7 @@ use datafusion_physical_optimizer::limit_pushdown::LimitPushdown;
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::joins::NestedLoopJoinExec;
+use datafusion_physical_plan::limit::GlobalLimitExec;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::{ExecutionPlan, get_plan_string};
@@ -162,6 +164,129 @@ fn transforms_streaming_table_exec_into_fetching_version_and_keeps_the_global_li
     Ok(())
 }
 
+#[test]
+fn preserves_required_ordering_when_reinserting_global_limit() -> Result<()> {
+    let schema = create_schema();
+    let projection =
+        projection_exec(Arc::clone(&schema), empty_exec(Arc::clone(&schema)))?;
+    let ordering = LexOrdering::new([PhysicalSortExpr::new(
+        col("c1", schema.as_ref())?,
+        SortOptions::default(),
+    )])
+    .unwrap();
+
+    let mut limit = GlobalLimitExec::new(projection, 2, Some(5));
+    limit.set_required_ordering(Some(ordering.clone()));
+
+    let optimized =
+        LimitPushdown::new().optimize(Arc::new(limit), &ConfigOptions::new())?;
+    let limit = optimized
+        .as_ref()
+        .downcast_ref::<ProjectionExec>()
+        .unwrap()
+        .input()
+        .as_ref()
+        .downcast_ref::<GlobalLimitExec>()
+        .unwrap();
+
+    assert_eq!(limit.required_ordering().as_ref(), Some(&ordering));
+    Ok(())
+}
+
+#[test]
+fn preserves_required_ordering_when_reinserting_local_limit() -> Result<()> {
+    let schema = create_schema();
+    let projection =
+        projection_exec(Arc::clone(&schema), empty_exec(Arc::clone(&schema)))?;
+    let repartition = repartition_exec(projection)?;
+    let ordering = LexOrdering::new([PhysicalSortExpr::new(
+        col("c1", schema.as_ref())?,
+        SortOptions::default(),
+    )])
+    .unwrap();
+
+    let mut limit = datafusion_physical_plan::limit::LocalLimitExec::new(repartition, 5);
+    limit.set_required_ordering(Some(ordering.clone()));
+
+    let optimized =
+        LimitPushdown::new().optimize(Arc::new(limit), &ConfigOptions::new())?;
+    let limit = optimized
+        .as_ref()
+        .downcast_ref::<datafusion_physical_plan::limit::LocalLimitExec>()
+        .unwrap();
+
+    assert_eq!(limit.required_ordering().as_ref(), Some(&ordering));
+    Ok(())
+}
+
+/// Regression test for #24215: a lost `required_ordering` is not just an
+/// unused metadata field. `FileScanConfig::create_sibling_state` (the hook
+/// that decides whether sibling partition streams may steal each other's
+/// file-reading work at runtime) only disables that work-stealing when
+/// `preserve_order` is `true`. `LimitPushdown` derives `preserve_order` from
+/// `required_ordering` when it absorbs a fetch into a scan
+/// (`with_preserve_order(required_ordering.is_some())`), so a limit that
+/// reaches this rule with its ordering already lost will leave work-stealing
+/// enabled on the scan beneath it — the exact mechanism by which a multi-file
+/// scan can read files out of order under an `ORDER BY ... LIMIT`.
+///
+/// This test does not race the scheduler (that would be flaky); it proves
+/// the deterministic precondition: whether `required_ordering` reaches this
+/// rule controls whether the resulting scan disables work-stealing.
+#[test]
+fn required_ordering_disables_scan_work_stealing() -> Result<()> {
+    let mut work_stealing_enabled = ConfigOptions::new();
+    work_stealing_enabled
+        .execution
+        .enable_file_stream_work_stealing = true;
+
+    let schema = create_schema();
+    let ordering = LexOrdering::new([PhysicalSortExpr::new(
+        col("c1", schema.as_ref())?,
+        SortOptions::default(),
+    )])
+    .unwrap();
+
+    // Ordering intact: the fix threads `required_ordering` through to the
+    // scan, so work-stealing must stay disabled.
+    let mut limit_with_ordering =
+        GlobalLimitExec::new(parquet_exec(Arc::clone(&schema)), 0, Some(5));
+    limit_with_ordering.set_required_ordering(Some(ordering));
+    let optimized_with_ordering = LimitPushdown::new()
+        .optimize(Arc::new(limit_with_ordering), &ConfigOptions::new())?;
+    let scan_with_ordering = optimized_with_ordering
+        .as_ref()
+        .downcast_ref::<DataSourceExec>()
+        .expect("fetch should have been absorbed into the scan");
+    assert!(
+        scan_with_ordering
+            .data_source()
+            .create_sibling_state(&work_stealing_enabled)
+            .is_none(),
+        "an order-sensitive limit must disable work-stealing on the scan it reaches"
+    );
+
+    // Ordering absent (simulating the pre-fix loss): work-stealing is free to
+    // reorder file reads under the limit, which is the wrong-results bug.
+    let limit_without_ordering = GlobalLimitExec::new(parquet_exec(schema), 0, Some(5));
+    let optimized_without_ordering = LimitPushdown::new()
+        .optimize(Arc::new(limit_without_ordering), &ConfigOptions::new())?;
+    let scan_without_ordering = optimized_without_ordering
+        .as_ref()
+        .downcast_ref::<DataSourceExec>()
+        .expect("fetch should have been absorbed into the scan");
+    assert!(
+        scan_without_ordering
+            .data_source()
+            .create_sibling_state(&work_stealing_enabled)
+            .is_some(),
+        "a lost required_ordering leaves work-stealing enabled, allowing files \
+         to be read out of order under the limit"
+    );
+
+    Ok(())
+}
+
 fn join_on_columns(
     left_col: &str,
     right_col: &str,

datafusion/physical-optimizer/src/limit_pushdown.rs

@@ -70,6 +70,7 @@ use datafusion_common::error::Result;
 use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};
 use datafusion_common::utils::combine_limit;
+use datafusion_physical_expr::LexOrdering;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
@@ -92,11 +93,11 @@ pub struct LimitPushdown {}
 ///
 /// [`LimitPushdown`]: crate::limit_pushdown::LimitPushdown
 #[derive(Default, Clone, Debug)]
-pub struct GlobalRequirements {
+pub(crate) struct GlobalRequirements {
     fetch: Option<usize>,
     skip: usize,
     satisfied: bool,
-    preserve_order: bool,
+    required_ordering: Option<LexOrdering>,
 }
 
 impl LimitPushdown {
@@ -116,7 +117,7 @@ impl PhysicalOptimizerRule for LimitPushdown {
             fetch: None,
             skip: 0,
             satisfied: false,
-            preserve_order: false,
+            required_ordering: None,
         };
         pushdown_limits(plan, global_state)
     }
@@ -134,7 +135,7 @@ struct LimitInfo {
     input: Arc<dyn ExecutionPlan>,
     fetch: Option<usize>,
     skip: usize,
-    preserve_order: bool,
+    required_ordering: Option<LexOrdering>,
 }
 
 /// This function is the main helper function of the `LimitPushDown` rule.
@@ -144,7 +145,7 @@ struct LimitInfo {
 ///
 /// If a limit is encountered, a [`TreeNodeRecursion::Stop`] is returned. Otherwise,
 /// return a [`TreeNodeRecursion::Continue`].
-pub fn pushdown_limit_helper(
+pub(crate) fn pushdown_limit_helper(
     mut pushdown_plan: Arc<dyn ExecutionPlan>,
     mut global_state: GlobalRequirements,
 ) -> Result<(Transformed<Arc<dyn ExecutionPlan>>, GlobalRequirements)> {
@@ -160,7 +161,7 @@ pub fn pushdown_limit_helper(
         );
         global_state.skip = skip;
         global_state.fetch = fetch;
-        global_state.preserve_order = limit_info.preserve_order;
+        global_state.required_ordering = limit_info.required_ordering;
         global_state.satisfied = false;
 
         if let Some(fetch) = fetch
@@ -219,6 +220,7 @@ pub fn pushdown_limit_helper(
                     pushdown_plan,
                     global_state.skip,
                     None,
+                    global_state.required_ordering.clone(),
                 )),
                 global_state,
             ))
@@ -243,8 +245,12 @@ pub fn pushdown_limit_helper(
             // Execution plans can't (yet) handle skip, so if we have one,
             // we still need to add a global limit
             if global_state.skip > 0 {
-                new_plan =
-                    add_global_limit(new_plan, global_state.skip, global_state.fetch);
+                new_plan = add_global_limit(
+                    new_plan,
+                    global_state.skip,
+                    global_state.fetch,
+                    global_state.required_ordering.clone(),
+                );
             }
             global_state.fetch = skip_and_fetch;
             global_state.skip = 0;
@@ -260,6 +266,7 @@ pub fn pushdown_limit_helper(
                     pushdown_plan,
                     global_state.skip,
                     global_fetch,
+                    global_state.required_ordering.clone(),
                 )),
                 global_state,
             ))
@@ -278,7 +285,7 @@ pub fn pushdown_limit_helper(
         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)
+                    .with_preserve_order(global_state.required_ordering.is_some())
                     .unwrap_or(plan_with_fetch);
                 Ok((Transformed::yes(plan_with_preserve_order), global_state))
             } else {
@@ -288,20 +295,26 @@ pub fn pushdown_limit_helper(
             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)
+                    .with_preserve_order(global_state.required_ordering.is_some())
                     .unwrap_or(plan_with_fetch);
 
                 if global_skip > 0 {
                     add_global_limit(
                         plan_with_preserve_order,
                         global_skip,
                         Some(global_fetch),
+                        global_state.required_ordering.clone(),
                     )
                 } else {
                     plan_with_preserve_order
                 }
             } else {
-                add_limit(pushdown_plan, global_skip, global_fetch)
+                add_limit(
+                    pushdown_plan,
+                    global_skip,
+                    global_fetch,
+                    global_state.required_ordering.clone(),
+                )
             };
             Ok((Transformed::yes(pushdown_plan), global_state))
         }
@@ -417,15 +430,15 @@ fn extract_limit(plan: &Arc<dyn ExecutionPlan>) -> Option<LimitInfo> {
             input: Arc::clone(global_limit.input()),
             fetch: global_limit.fetch(),
             skip: global_limit.skip(),
-            preserve_order: global_limit.required_ordering().is_some(),
+            required_ordering: global_limit.required_ordering().clone(),
         })
     } else {
         plan.downcast_ref::<LocalLimitExec>()
             .map(|local_limit| LimitInfo {
                 input: Arc::clone(local_limit.input()),
                 fetch: Some(local_limit.fetch()),
                 skip: 0,
-                preserve_order: local_limit.required_ordering().is_some(),
+                required_ordering: local_limit.required_ordering().clone(),
             })
     }
 }
@@ -435,17 +448,22 @@ fn combines_input_partitions(plan: &Arc<dyn ExecutionPlan>) -> bool {
     plan.is::<CoalescePartitionsExec>() || plan.is::<SortPreservingMergeExec>()
 }
 
-/// Adds a limit to the plan, chooses between global and local limits based on
-/// skip value and the number of partitions.
+/// Adds a limit to the plan, choosing between global and local limits
+/// based on the skip value and the number of partitions.
 fn add_limit(
     pushdown_plan: Arc<dyn ExecutionPlan>,
     skip: usize,
     fetch: usize,
+    required_ordering: Option<LexOrdering>,
 ) -> Arc<dyn ExecutionPlan> {
     if skip > 0 || pushdown_plan.output_partitioning().partition_count() == 1 {
-        add_global_limit(pushdown_plan, skip, Some(fetch))
+        add_global_limit(pushdown_plan, skip, Some(fetch), required_ordering)
     } else {
-        Arc::new(LocalLimitExec::new(pushdown_plan, fetch + skip)) as _
+        let mut limit = LocalLimitExec::new(pushdown_plan, fetch + skip);
+
+        limit.set_required_ordering(required_ordering);
+
+        Arc::new(limit)
     }
 }
 
@@ -454,8 +472,10 @@ fn add_global_limit(
     pushdown_plan: Arc<dyn ExecutionPlan>,
     skip: usize,
     fetch: Option<usize>,
+    required_ordering: Option<LexOrdering>,
 ) -> Arc<dyn ExecutionPlan> {
-    Arc::new(GlobalLimitExec::new(pushdown_plan, skip, fetch)) as _
+    let mut limit = GlobalLimitExec::new(pushdown_plan, skip, fetch);
+    limit.set_required_ordering(required_ordering);
+    Arc::new(limit)
 }
-
 // See tests in datafusion/core/tests/physical_optimizer

datafusion/physical-optimizer/src/pushdown_sort.rs

@@ -67,7 +67,6 @@ use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use std::sync::Arc;
-
 /// A PhysicalOptimizerRule that attempts to push down sort requirements to data sources.
 ///
 /// See module-level documentation for details.
@@ -108,10 +107,16 @@ impl PhysicalOptimizerRule for PushdownSort {
                 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.
+                        // Use LocalLimitExec (not Global) since input is multi-partition:
+                        // this rule runs after distribution enforcement, and
+                        // GlobalLimitExec requires a single input partition.
                         let inner = if let Some(fetch) = sort_child.fetch() {
                             inner.with_fetch(Some(fetch)).unwrap_or_else(|| {
-                                Arc::new(LocalLimitExec::new(inner, fetch))
+                                let mut limit = LocalLimitExec::new(inner, fetch);
+                                limit.set_required_ordering(Some(
+                                    sort_child.expr().clone(),
+                                ));
+                                Arc::new(limit)
                             })
                         } else {
                             inner
@@ -171,7 +176,9 @@ impl PhysicalOptimizerRule for PushdownSort {
                     // 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)))
+                            let mut limit = GlobalLimitExec::new(inner, 0, Some(fetch));
+                            limit.set_required_ordering(Some(sort_exec.expr().clone()));
+                            Arc::new(limit)
                         });
                         Ok(Transformed::yes(inner))
                     } else {