fix(proto): serialize Global/LocalLimitExec required_ordering

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -34,7 +34,7 @@ use arrow::datatypes::{DataType, SchemaRef};
 use datafusion_common::config::{ConfigOptions, CsvOptions};
 use datafusion_common::tree_node::{TreeNode, TransformedResult};
 use datafusion_common::{create_array, DataFusionError, NullEquality, Result, TableReference};
-use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
+use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr_common::operator::Operator;
 use datafusion_expr::{JoinType, SortExpr};
@@ -57,6 +57,8 @@ use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan;
 use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext};
 use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown, assign_initial_requirements, pushdown_sorts};
 use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
+use datafusion_physical_optimizer::limit_pushdown::LimitPushdown;
+use datafusion_physical_optimizer::projection_pushdown::ProjectionPushdown;
 use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion::prelude::*;
@@ -2297,6 +2299,55 @@ async fn test_remove_unnecessary_spm2() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_spm_fetch_preserves_ordering_through_child_rewrite() -> Result<()> {
+    let schema = create_test_schema()?;
+    let ordering: LexOrdering = [sort_expr("non_nullable_col", &schema)].into();
+    let source = parquet_exec_with_sort(Arc::clone(&schema), vec![ordering.clone()]);
+    let projection = projection_exec(
+        vec![
+            (col("nullable_col", &schema)?, "nullable_col".to_string()),
+            (
+                col("non_nullable_col", &schema)?,
+                "non_nullable_col".to_string(),
+            ),
+        ],
+        source,
+    )?;
+    let plan = sort_preserving_merge_exec_with_fetch(ordering.clone(), projection, 100);
+
+    let optimized = PlanWithCorrespondingSort::new_default(plan)
+        .transform_up(ensure_sorting)?
+        .data;
+    let optimized = check_integrity(optimized)?.plan;
+    let limit = optimized
+        .downcast_ref::<LocalLimitExec>()
+        .expect("SPM fetch should become a local limit");
+    assert_eq!(limit.fetch(), 100);
+    assert_eq!(limit.required_ordering().as_ref(), Some(&ordering));
+
+    let config = ConfigOptions::new();
+    let optimized = ProjectionPushdown::new().optimize(optimized, &config)?;
+    let limit = optimized
+        .downcast_ref::<LocalLimitExec>()
+        .expect("projection rewrite should retain the local limit");
+    assert_eq!(limit.required_ordering().as_ref(), Some(&ordering));
+    assert!(limit.input().is::<DataSourceExec>());
+
+    let optimized = LimitPushdown::new().optimize(optimized, &config)?;
+    let source = optimized
+        .downcast_ref::<DataSourceExec>()
+        .expect("limit should be pushed into the parquet scan");
+    let config = source
+        .data_source()
+        .downcast_ref::<FileScanConfig>()
+        .expect("parquet scan should use FileScanConfig");
+    assert_eq!(config.limit, Some(100));
+    assert!(config.preserve_order);
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_change_wrong_sorting() -> Result<()> {
     let schema = create_test_schema()?;

datafusion/physical-expr-common/src/sort_expr.rs

@@ -265,7 +265,8 @@ pub fn sort_exprs_try_to_proto<E: std::borrow::Borrow<PhysicalSortExpr>>(
 /// [`LexRequirement`], because callers differ in what an empty list means:
 /// `LexOrdering::new` / `LexRequirement::new` return `None` for it, which is
 /// "no ordering declared" for a scan and an error for an operator that requires
-/// one.
+/// one. Callers with the former convention can use
+/// [`optional_ordering_try_from_proto`] instead.
 ///
 /// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode
 #[cfg(feature = "proto")]
@@ -279,6 +280,25 @@ pub fn sort_exprs_try_from_proto(
         .collect()
 }
 
+/// Serialize an optional [`LexOrdering`], encoding `None` as an empty list.
+#[cfg(feature = "proto")]
+pub fn optional_ordering_try_to_proto(
+    ordering: Option<&LexOrdering>,
+    ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+) -> Result<Vec<datafusion_proto_models::protobuf::PhysicalSortExprNode>> {
+    sort_exprs_try_to_proto(ordering.into_iter().flatten(), ctx)
+}
+
+/// Counterpart of [`optional_ordering_try_to_proto`]: an empty list decodes
+/// as `None`.
+#[cfg(feature = "proto")]
+pub fn optional_ordering_try_from_proto(
+    nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode],
+    ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+) -> Result<Option<LexOrdering>> {
+    Ok(LexOrdering::new(sort_exprs_try_from_proto(nodes, ctx)?))
+}
+
 impl PartialEq for PhysicalSortExpr {
     fn eq(&self, other: &Self) -> bool {
         self.options == other.options && self.expr.eq(&other.expr)

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

@@ -461,13 +461,17 @@ pub fn ensure_sorting(
     } else if is_sort_preserving_merge(&requirements.plan)
         && child_node.plan.output_partitioning().partition_count() <= 1
     {
-        // This `SortPreservingMergeExec` is unnecessary, input already has a
-        // single partition and no fetch is required.
-        let mut child_node = requirements.children.swap_remove(0);
+        // This `SortPreservingMergeExec` is unnecessary because its input has a
+        // single partition.
+        let child_node = requirements.children.swap_remove(0);
         if let Some(fetch) = requirements.plan.fetch() {
-            // Add the limit exec if the original SPM had a fetch:
-            child_node.plan =
-                Arc::new(LocalLimitExec::new(Arc::clone(&child_node.plan), fetch));
+            let mut limit = LocalLimitExec::new(Arc::clone(&child_node.plan), fetch);
+            limit.set_required_ordering(requirements.plan.output_ordering().cloned());
+            return Ok(Transformed::yes(PlanContext::new(
+                Arc::new(limit),
+                false,
+                vec![child_node],
+            )));
         }
         return Ok(Transformed::yes(child_node));
     }

datafusion/physical-plan/src/joins/symmetric_hash_join.rs

@@ -699,20 +699,16 @@ impl ExecutionPlan for SymmetricHashJoinExec {
             })
             .transpose()?;
         let expr_ctx = ctx.expr_ctx();
-        let encode_sort_exprs =
-            |exprs: Option<&LexOrdering>| -> Result<Vec<protobuf::PhysicalSortExprNode>> {
-                exprs.map_or_else(
-                    || Ok(vec![]),
-                    |exprs| {
-                        datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto(
-                            exprs.iter(),
-                            &expr_ctx,
-                        )
-                    },
-                )
-            };
-        let left_sort_exprs = encode_sort_exprs(self.left_sort_exprs())?;
-        let right_sort_exprs = encode_sort_exprs(self.right_sort_exprs())?;
+        let left_sort_exprs =
+            datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto(
+                self.left_sort_exprs(),
+                &expr_ctx,
+            )?;
+        let right_sort_exprs =
+            datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto(
+                self.right_sort_exprs(),
+                &expr_ctx,
+            )?;
 
         Ok(Some(protobuf::PhysicalPlanNode {
             physical_plan_type: Some(
@@ -876,20 +872,16 @@ impl SymmetricHashJoinExec {
                 ))
             })
             .transpose()?;
-        let decode_sort_exprs = |sort_exprs: &[protobuf::PhysicalSortExprNode],
-                                 schema: &Schema|
-         -> Result<Option<LexOrdering>> {
-            let sort_exprs =
-                datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto(
-                    sort_exprs,
-                    &ctx.expr_ctx(schema),
-                )?;
-            Ok(LexOrdering::new(sort_exprs))
-        };
         let left_sort_exprs =
-            decode_sort_exprs(&sym_join.left_sort_exprs, left_schema.as_ref())?;
+            datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto(
+                &sym_join.left_sort_exprs,
+                &ctx.expr_ctx(left_schema.as_ref()),
+            )?;
         let right_sort_exprs =
-            decode_sort_exprs(&sym_join.right_sort_exprs, right_schema.as_ref())?;
+            datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto(
+                &sym_join.right_sort_exprs,
+                &ctx.expr_ctx(right_schema.as_ref()),
+            )?;
 
         Self::try_new(
             left,