Fix CoalescePartitionsExec proto serialization (#15824)

* add fetch to CoalescePartitionsExecNode

* gen proto code

* Add test

* fix

* fix build

* Fix test build

* remove comments

Author: lewiszlw

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -3447,9 +3447,8 @@ fn test_parallelize_sort_preserves_fetch() -> Result<()> {
     let schema = create_test_schema3()?;
     let parquet_exec = parquet_exec(&schema);
     let coalesced = Arc::new(CoalescePartitionsExec::new(parquet_exec.clone()));
-    let top_coalesced = CoalescePartitionsExec::new(coalesced.clone())
-        .with_fetch(Some(10))
-        .unwrap();
+    let top_coalesced =
+        Arc::new(CoalescePartitionsExec::new(coalesced.clone()).with_fetch(Some(10)));
 
     let requirements = PlanWithCorrespondingCoalescePartitions::new(
         top_coalesced.clone(),

datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -1272,9 +1272,8 @@ fn test_plan_with_order_preserving_variants_preserves_fetch() -> Result<()> {
         "a", &schema,
     )];
     let parquet_exec = parquet_exec_sorted(&schema, parquet_sort_exprs);
-    let coalesced = CoalescePartitionsExec::new(parquet_exec.clone())
-        .with_fetch(Some(10))
-        .unwrap();
+    let coalesced =
+        Arc::new(CoalescePartitionsExec::new(parquet_exec.clone()).with_fetch(Some(10)));
 
     // Test sort's fetch is greater than coalesce fetch, return error because it's not reasonable
     let requirements = OrderPreservationContext::new(

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -1035,10 +1035,9 @@ pub fn replace_order_preserving_variants(
 
     if is_sort_preserving_merge(&context.plan) {
         let child_plan = Arc::clone(&context.children[0].plan);
-        // It's safe to unwrap because `CoalescePartitionsExec` supports `fetch`.
-        context.plan = CoalescePartitionsExec::new(child_plan)
-            .with_fetch(context.plan.fetch())
-            .unwrap();
+        context.plan = Arc::new(
+            CoalescePartitionsExec::new(child_plan).with_fetch(context.plan.fetch()),
+        );
         return Ok(context);
     } else if let Some(repartition) =
         context.plan.as_any().downcast_ref::<RepartitionExec>()

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

@@ -409,10 +409,10 @@ pub fn parallelize_sorts(
 
         Ok(Transformed::yes(
             PlanWithCorrespondingCoalescePartitions::new(
-                // Safe to unwrap, because `CoalescePartitionsExec` has a fetch
-                CoalescePartitionsExec::new(Arc::clone(&requirements.plan))
-                    .with_fetch(fetch)
-                    .unwrap(),
+                Arc::new(
+                    CoalescePartitionsExec::new(Arc::clone(&requirements.plan))
+                        .with_fetch(fetch),
+                ),
                 false,
                 vec![requirements],
             ),

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -59,6 +59,12 @@ impl CoalescePartitionsExec {
         }
     }
 
+    /// Update fetch with the argument
+    pub fn with_fetch(mut self, fetch: Option<usize>) -> Self {
+        self.fetch = fetch;
+        self
+    }
+
     /// Input execution plan
     pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
         &self.input

datafusion/proto/proto/datafusion.proto

@@ -1217,6 +1217,7 @@ message CoalesceBatchesExecNode {
 
 message CoalescePartitionsExecNode {
   PhysicalPlanNode input = 1;
+  optional uint32 fetch = 2;
 }
 
 message PhysicalHashRepartition {

datafusion/proto/src/generated/pbjson.rs

@@ -792,7 +792,10 @@ impl protobuf::PhysicalPlanNode {
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let input: Arc<dyn ExecutionPlan> =
             into_physical_plan(&merge.input, registry, runtime, extension_codec)?;
-        Ok(Arc::new(CoalescePartitionsExec::new(input)))
+        Ok(Arc::new(
+            CoalescePartitionsExec::new(input)
+                .with_fetch(merge.fetch.map(|f| f as usize)),
+        ))
     }
 
     fn try_into_repartition_physical_plan(
@@ -2354,6 +2357,7 @@ impl protobuf::PhysicalPlanNode {
             physical_plan_type: Some(PhysicalPlanType::Merge(Box::new(
                 protobuf::CoalescePartitionsExecNode {
                     input: Some(Box::new(input)),
+                    fetch: exec.fetch().map(|f| f as u32),
                 },
             ))),
         })

datafusion/proto/src/generated/prost.rs

@@ -66,6 +66,7 @@ use datafusion::physical_plan::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy,
 };
 use datafusion::physical_plan::analyze::AnalyzeExec;
+use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::expressions::{
     binary, cast, col, in_list, like, lit, BinaryExpr, Column, NotExpr, PhysicalSortExpr,
@@ -709,7 +710,7 @@ fn roundtrip_sort_preserve_partitioning() -> Result<()> {
 }
 
 #[test]
-fn roundtrip_coalesce_with_fetch() -> Result<()> {
+fn roundtrip_coalesce_batches_with_fetch() -> 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]));
@@ -725,6 +726,22 @@ fn roundtrip_coalesce_with_fetch() -> Result<()> {
     ))
 }
 
+#[test]
+fn roundtrip_coalesce_partitions_with_fetch() -> 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]));
+
+    roundtrip_test(Arc::new(CoalescePartitionsExec::new(Arc::new(
+        EmptyExec::new(schema.clone()),
+    ))))?;
+
+    roundtrip_test(Arc::new(
+        CoalescePartitionsExec::new(Arc::new(EmptyExec::new(schema)))
+            .with_fetch(Some(10)),
+    ))
+}
+
 #[test]
 fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> {
     let file_schema =