Serialize preserve_partitioning in SortExec (apache#5306)

datafusion/proto/proto/datafusion.proto

@@ -1266,6 +1266,7 @@ message SortExecNode {
   repeated PhysicalExprNode expr = 2;
   // Maximum number of highest/lowest rows to fetch; negative means no limit
   int64 fetch = 3;
+  bool preserve_partitioning = 4;
 }
 
 message SortPreservingMergeExecNode {

datafusion/proto/src/physical_plan/mod.rs

@@ -608,7 +608,12 @@ impl AsExecutionPlan for PhysicalPlanNode {
                 } else {
                     Some(sort.fetch as usize)
                 };
-                Ok(Arc::new(SortExec::try_new(exprs, input, fetch)?))
+                Ok(Arc::new(SortExec::new_with_partitioning(
+                    exprs,
+                    input,
+                    sort.preserve_partitioning,
+                    fetch,
+                )))
             }
             PhysicalPlanType::SortPreservingMerge(sort) => {
                 let input: Arc<dyn ExecutionPlan> =
@@ -1043,6 +1048,7 @@ impl AsExecutionPlan for PhysicalPlanNode {
                             Some(n) => n as i64,
                             _ => -1,
                         },
+                        preserve_partitioning: exec.preserve_partitioning(),
                     },
                 ))),
             })
@@ -1441,6 +1447,43 @@ mod roundtrip_tests {
         )?))
     }
 
+    #[test]
+    fn roundtrip_sort_preserve_partitioning() -> Result<()> {
+        let field_a = Field::new("a", DataType::Boolean, false);
+        let field_b = Field::new("b", DataType::Int64, false);
+        let schema = Arc::new(Schema::new(vec![field_a, field_b]));
+        let sort_exprs = vec![
+            PhysicalSortExpr {
+                expr: col("a", &schema)?,
+                options: SortOptions {
+                    descending: true,
+                    nulls_first: false,
+                },
+            },
+            PhysicalSortExpr {
+                expr: col("b", &schema)?,
+                options: SortOptions {
+                    descending: false,
+                    nulls_first: true,
+                },
+            },
+        ];
+
+        roundtrip_test(Arc::new(SortExec::new_with_partitioning(
+            sort_exprs.clone(),
+            Arc::new(EmptyExec::new(false, schema.clone())),
+            false,
+            None,
+        )))?;
+
+        roundtrip_test(Arc::new(SortExec::new_with_partitioning(
+            sort_exprs,
+            Arc::new(EmptyExec::new(false, schema)),
+            true,
+            None,
+        )))
+    }
+
     #[test]
     fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> {
         let scan_config = FileScanConfig {