fix: preserve input ordering required by limits

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -2245,6 +2245,29 @@ async fn test_do_not_pushdown_through_limit() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_limit_preserves_input_order_through_projection() -> Result<()> {
+    let schema = create_test_schema()?;
+    let source = memory_exec(&schema);
+    let input = sort_exec([sort_expr("non_nullable_col", &schema)].into(), source);
+    let projection = projection_exec(
+        vec![(col("nullable_col", &schema)?, "nullable_col".to_string())],
+        input,
+    )?;
+    let physical_plan = Arc::new(GlobalLimitExec::new(projection, 1, None)) as _;
+
+    let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true);
+    assert_snapshot!(test.run(), @r"
+    Input / Optimized Plan:
+    GlobalLimitExec: skip=1, fetch=None
+      ProjectionExec: expr=[nullable_col@0 as nullable_col]
+        SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false]
+          DataSourceExec: partitions=1, partition_sizes=[0]
+    ");
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_remove_unnecessary_spm1() -> Result<()> {
     let schema = create_test_schema()?;

datafusion/core/tests/physical_optimizer/ensure_requirements.rs

@@ -299,6 +299,36 @@ fn test_projection_over_multi_partition_sort_limit() {
     ");
 }
 
+/// A limit consumes the global ordering established below a projection even
+/// when the projection removes the sort key from its output schema. The
+/// ordering must survive a second `EnsureRequirements` pass, which sees the
+/// `SortPreservingMergeExec` introduced by the first pass.
+#[test]
+fn test_limit_preserves_hidden_sort_across_distribution_passes() {
+    let source = Arc::new(MockMultiPartitionExec::new(16));
+    let sort_expr = sort_expr_on("a", 0, true, true);
+    let sort = Arc::new(SortExec::new(sort_expr, source));
+
+    // Hide the sort key so the projection itself has no output ordering.
+    let projection = Arc::new(
+        ProjectionExec::try_new(
+            vec![(Arc::new(Column::new("b", 1)) as _, "b".to_string())],
+            sort,
+        )
+        .unwrap(),
+    );
+    let limit: Arc<dyn ExecutionPlan> =
+        Arc::new(GlobalLimitExec::new(projection, 1, None));
+
+    assert_ensure_requirements_plan!(limit, @r"
+    GlobalLimitExec: skip=1, fetch=None
+      ProjectionExec: expr=[b@1 as b]
+        SortPreservingMergeExec: [a@0 DESC]
+          SortExec: expr=[a@0 DESC], preserve_partitioning=[true]
+            MockMultiPartitionExec
+    ");
+}
+
 // ========================================================================
 // Single partition tests (no unnecessary operators)
 // ========================================================================

datafusion/ffi/src/execution_plan.rs

@@ -50,6 +50,10 @@ pub struct FFI_ExecutionPlan {
     /// Return a vector of children plans
     pub children: unsafe extern "C" fn(plan: &Self) -> SVec<FFI_ExecutionPlan>,
 
+    /// Return whether each child input's ordering must be preserved.
+    pub requires_input_order_preservation:
+        unsafe extern "C" fn(plan: &Self) -> SVec<bool>,
+
     pub with_new_children:
         unsafe extern "C" fn(plan: &Self, children: SVec<Self>) -> FFI_Result<Self>,
 
@@ -138,6 +142,15 @@ unsafe extern "C" fn children_fn_wrapper(
         .collect()
 }
 
+unsafe extern "C" fn requires_input_order_preservation_fn_wrapper(
+    plan: &FFI_ExecutionPlan,
+) -> SVec<bool> {
+    plan.inner()
+        .requires_input_order_preservation()
+        .into_iter()
+        .collect()
+}
+
 unsafe extern "C" fn with_new_children_fn_wrapper(
     plan: &FFI_ExecutionPlan,
     children: SVec<FFI_ExecutionPlan>,
@@ -306,6 +319,8 @@ impl FFI_ExecutionPlan {
         Self {
             properties: properties_fn_wrapper,
             children: children_fn_wrapper,
+            requires_input_order_preservation:
+                requires_input_order_preservation_fn_wrapper,
             with_new_children: with_new_children_fn_wrapper,
             name: name_fn_wrapper,
             execute: execute_fn_wrapper,
@@ -416,6 +431,12 @@ impl ExecutionPlan for ForeignExecutionPlan {
         self.children.iter().collect()
     }
 
+    fn requires_input_order_preservation(&self) -> Vec<bool> {
+        unsafe { (self.plan.requires_input_order_preservation)(&self.plan) }
+            .into_iter()
+            .collect()
+    }
+
     fn with_new_children(
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
@@ -485,6 +506,7 @@ pub mod tests {
         children: Vec<Arc<dyn ExecutionPlan>>,
         metrics: Option<MetricsSet>,
         statistics: Option<Statistics>,
+        requires_input_order_preservation: bool,
     }
 
     impl EmptyExec {
@@ -499,9 +521,18 @@ pub mod tests {
                 children: Vec::default(),
                 metrics: None,
                 statistics: None,
+                requires_input_order_preservation: false,
             }
         }
 
+        pub fn with_input_order_preservation(
+            mut self,
+            requires_input_order_preservation: bool,
+        ) -> Self {
+            self.requires_input_order_preservation = requires_input_order_preservation;
+            self
+        }
+
         pub fn with_metrics(mut self, metrics: MetricsSet) -> Self {
             self.metrics = Some(metrics);
             self
@@ -545,9 +576,14 @@ pub mod tests {
                 children,
                 metrics: self.metrics.clone(),
                 statistics: self.statistics.clone(),
+                requires_input_order_preservation: self.requires_input_order_preservation,
             }))
         }
 
+        fn requires_input_order_preservation(&self) -> Vec<bool> {
+            vec![self.requires_input_order_preservation; self.children.len()]
+        }
+
         fn execute(
             &self,
             _partition: usize,
@@ -600,6 +636,25 @@ pub mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_ffi_execution_plan_input_order_preservation_round_trip() -> Result<()> {
+        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+            arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Float32, false),
+        ]));
+        let child = Arc::new(EmptyExec::new(Arc::clone(&schema))) as _;
+        let original_plan =
+            Arc::new(EmptyExec::new(schema).with_input_order_preservation(true))
+                .with_new_children(vec![child])?;
+
+        let mut local_plan = FFI_ExecutionPlan::new(original_plan, None);
+        local_plan.library_marker_id = crate::mock_foreign_marker_id;
+        let foreign_plan: Arc<dyn ExecutionPlan> = (&local_plan).try_into()?;
+
+        assert_eq!(foreign_plan.requires_input_order_preservation(), vec![true]);
+
+        Ok(())
+    }
+
     #[test]
     fn test_ffi_execution_plan_children() -> Result<()> {
         let schema = Arc::new(arrow::datatypes::Schema::new(vec![

datafusion/ffi/src/tests/mod.rs

@@ -107,6 +107,8 @@ pub struct ForeignLibraryModule {
 
     pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan,
 
+    pub create_order_sensitive_exec: extern "C" fn() -> FFI_ExecutionPlan,
+
     pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan,
 
     pub create_table_with_statistics:
@@ -161,6 +163,15 @@ pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan {
     FFI_ExecutionPlan::new(plan, None)
 }
 
+pub(crate) extern "C" fn create_order_sensitive_exec() -> FFI_ExecutionPlan {
+    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+    let child = Arc::new(EmptyExec::new(Arc::clone(&schema))) as _;
+    let plan = Arc::new(EmptyExec::new(schema).with_input_order_preservation(true))
+        .with_new_children(vec![child])
+        .expect("order-sensitive test plan should accept one child");
+    FFI_ExecutionPlan::new(plan, None)
+}
+
 /// Returns canonical statistics used by both the producer and consumer sides of
 /// the integration tests so round-trips can be asserted without hard-coding
 /// the values in two places.
@@ -259,6 +270,7 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule {
         create_rank_udwf: create_ffi_rank_func,
         create_extension_options: config::create_extension_options,
         create_empty_exec,
+        create_order_sensitive_exec,
         create_exec_with_statistics,
         create_table_with_statistics,
         create_physical_optimizer_rule:

datafusion/ffi/tests/ffi_execution_plan.rs

@@ -28,6 +28,19 @@ mod tests {
     use datafusion_physical_plan::ExecutionPlan;
     use std::sync::Arc;
 
+    #[test]
+    fn test_ffi_execution_plan_input_order_preservation_cross_library()
+    -> Result<(), DataFusionError> {
+        let module = get_module()?;
+        let plan = (module.create_order_sensitive_exec)();
+        let plan: Arc<dyn ExecutionPlan> = (&plan).try_into()?;
+
+        assert!(plan.is::<ForeignExecutionPlan>());
+        assert_eq!(plan.requires_input_order_preservation(), vec![true]);
+
+        Ok(())
+    }
+
     #[test]
     #[expect(deprecated)]
     fn test_ffi_execution_plan_partition_statistics_cross_library()

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -839,12 +839,21 @@ struct RemovedDistOps {
 
 fn remove_dist_changing_operators(
     mut distribution_context: DistributionContext,
+    preserve_output_order: bool,
 ) -> Result<RemovedDistOps> {
     let mut removed_fetch = None;
     while is_repartition(&distribution_context.plan)
         || is_coalesce_partitions(&distribution_context.plan)
         || is_sort_preserving_merge(&distribution_context.plan)
     {
+        // Keep an order-producing distribution operator when an ancestor
+        // semantically depends on its output sequence. Once such an operator
+        // is removed, an intervening projection may make its ordering
+        // impossible to reconstruct.
+        if preserve_output_order && distribution_context.plan.output_ordering().is_some()
+        {
+            break;
+        }
         // Preserve fetch from SPM or CoalescePartitions before removing (#14150).
         if let Some(fetch) = distribution_context.plan.fetch() {
             removed_fetch = Some(
@@ -936,6 +945,7 @@ struct DistributionChildState {
     context: DistributionContext,
     required_input_ordering: Option<OrderingRequirements>,
     maintains_input_order: bool,
+    preserve_output_order: bool,
     requirement: Distribution,
 }
 
@@ -1126,13 +1136,46 @@ fn enforce_distribution_relationships(
 /// This function is intended to be used in a bottom up traversal, as it
 /// can first repartition (or newly partition) at the datasources -- these
 /// source partitions may be later repartitioned with additional data exchange operators.
+pub fn ensure_distribution(
+    dist_context: DistributionContext,
+    config: &ConfigOptions,
+) -> Result<Transformed<DistributionContext>> {
+    ensure_distribution_inner(dist_context, config, false)
+}
+
+/// Enforces distribution bottom-up while carrying semantic order-preservation
+/// requirements top-down.
+pub(super) fn ensure_distribution_with_order_preservation(
+    mut dist_context: DistributionContext,
+    config: &ConfigOptions,
+    preserve_output_order: bool,
+) -> Result<DistributionContext> {
+    let own_requirements = dist_context.plan.requires_input_order_preservation();
+    let maintains_input_order = dist_context.plan.maintains_input_order();
+    dist_context.children = izip!(
+        dist_context.children,
+        own_requirements,
+        maintains_input_order
+    )
+    .map(|(child, own_requirement, maintains_input_order)| {
+        let preserve_child_order =
+            own_requirement || (preserve_output_order && maintains_input_order);
+        ensure_distribution_with_order_preservation(child, config, preserve_child_order)
+    })
+    .collect::<Result<Vec<_>>>()?;
+
+    ensure_distribution_inner(dist_context, config, preserve_output_order)
+        .map(|transformed| transformed.data)
+}
+
 #[expect(
     deprecated,
     reason = "HashPartitioned is accepted during the KeyPartitioned migration"
 )]
-pub fn ensure_distribution(
+fn ensure_distribution_inner(
     dist_context: DistributionContext,
     config: &ConfigOptions,
+    preserve_output_order: bool,
 ) -> Result<Transformed<DistributionContext>> {
     let dist_context = update_children(dist_context)?;
 
@@ -1171,7 +1214,7 @@ pub fn ensure_distribution(
                 children,
             },
         removed_fetch,
-    } = remove_dist_changing_operators(dist_context)?;
+    } = remove_dist_changing_operators(dist_context, preserve_output_order)?;
 
     if let Some(exec) = plan.downcast_ref::<WindowAggExec>() {
         if let Some(updated_window) = get_best_fitting_window(
@@ -1238,6 +1281,7 @@ pub fn ensure_distribution(
         children.into_iter(),
         plan.required_input_ordering(),
         plan.maintains_input_order(),
+        plan.requires_input_order_preservation(),
         repartition_status_flags.into_iter()
     )
     .map(
@@ -1246,13 +1290,16 @@ pub fn ensure_distribution(
             mut child,
             required_input_ordering,
             maintains,
+            own_order_preservation,
             RepartitionRequirementStatus {
                 requirement,
                 roundrobin_beneficial,
                 roundrobin_beneficial_stats,
                 hash_necessary,
             },
         )| {
+            let preserve_child_order =
+                own_order_preservation || (preserve_output_order && maintains);
             // Allow subset satisfaction when:
             // 1. Current partition count >= threshold
             // 2. Not a partitioned join since must use exact hash matching for joins
@@ -1391,6 +1438,7 @@ pub fn ensure_distribution(
                 context: child,
                 required_input_ordering,
                 maintains_input_order: maintains,
+                preserve_output_order: preserve_child_order,
                 requirement,
             })
         },
@@ -1414,6 +1462,7 @@ pub fn ensure_distribution(
                  mut context,
                  required_input_ordering,
                  maintains_input_order,
+                 preserve_output_order,
                  requirement,
              }| {
                 let streaming_benefit = if context.data {
@@ -1436,6 +1485,7 @@ pub fn ensure_distribution(
                     if (!ordering_satisfied || !order_preserving_variants_desirable)
                         && !streaming_benefit
                         && context.data
+                        && !preserve_output_order
                     {
                         context = replace_order_preserving_variants(context)?;
                         // If ordering requirements were satisfied before repartitioning,
@@ -1464,14 +1514,18 @@ pub fn ensure_distribution(
                             // ordering is pointless. However, if it does maintain
                             // input order, we keep order-preserving variants so
                             // ordering can flow through to ancestors that need it.
-                            if !maintains_input_order && !streaming_benefit {
+                            if !maintains_input_order
+                                && !streaming_benefit
+                                && !preserve_output_order
+                            {
                                 context = replace_order_preserving_variants(context)?;
                             }
                         }
                         Distribution::UnspecifiedDistribution => {
                             // Since ordering is lost, trying to preserve ordering is pointless
-                            if !maintains_input_order
-                                || plan.is::<OutputRequirementExec>()
+                            if (!maintains_input_order
+                                || plan.is::<OutputRequirementExec>())
+                                && !preserve_output_order
                             {
                                 context = replace_order_preserving_variants(context)?;
                             }