Fix CoalescePartitionsExec proto serialization (apache#15824) (#299)

* add fetch to CoalescePartitionsExecNode

* gen proto code

* Add test

* fix

* fix build

* Fix test build

* remove comments

Co-authored-by: 张林伟 <lewiszlw520@gmail.com>

Author: 2 people

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

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

datafusion/proto/src/generated/pbjson.rs

@@ -291,7 +291,10 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode {
             PhysicalPlanType::Merge(merge) => {
                 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(|n| n as usize)),
+                ))
             }
             PhysicalPlanType::Repartition(repart) => {
                 let input: Arc<dyn ExecutionPlan> = into_physical_plan(
@@ -1671,6 +1674,7 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode {
                 physical_plan_type: Some(PhysicalPlanType::Merge(Box::new(
                     protobuf::CoalescePartitionsExecNode {
                         input: Some(Box::new(input)),
+                        fetch: exec.fetch().map(|n| n as u32),
                     },
                 ))),
             });

datafusion/proto/src/generated/prost.rs

@@ -58,6 +58,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, NthValue,
@@ -641,7 +642,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]));
@@ -657,6 +658,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 scan_config = FileScanConfig {