Add standalone NetworkCoalesceExec test and test with n_tasks > 1.

Author: JSOD11

src/protobuf/distributed_codec.rs

@@ -536,11 +536,32 @@ mod tests {
     }
 
     #[test]
-    fn test_roundtrip_isolator_flight_coalesce() -> datafusion::common::Result<()> {
+    fn test_roundtrip_single_flight_coalesce() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
         let registry = MemoryFunctionRegistry::new();
 
         let schema = schema_i32("e");
+        let plan: Arc<dyn ExecutionPlan> = Arc::new(new_network_coalesce_tasks_exec(
+            Partitioning::RoundRobinBatch(3),
+            schema,
+            dummy_stage(),
+        ));
+
+        let mut buf = Vec::new();
+        codec.try_encode(plan.clone(), &mut buf)?;
+
+        let decoded = codec.try_decode(&buf, &[empty_exec()], &registry)?;
+        assert_eq!(repr(&plan), repr(&decoded));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_roundtrip_isolator_flight_coalesce() -> datafusion::common::Result<()> {
+        let codec = DistributedCodec;
+        let registry = MemoryFunctionRegistry::new();
+
+        let schema = schema_i32("f");
         let flight = Arc::new(new_network_coalesce_tasks_exec(
             Partitioning::UnknownPartitioning(1),
             schema,
@@ -564,7 +585,7 @@ mod tests {
         let codec = DistributedCodec;
         let registry = MemoryFunctionRegistry::new();
 
-        let schema = schema_i32("c");
+        let schema = schema_i32("g");
         let left = Arc::new(new_network_coalesce_tasks_exec(
             Partitioning::RoundRobinBatch(2),
             schema.clone(),
@@ -578,7 +599,7 @@ mod tests {
 
         let union = Arc::new(UnionExec::new(vec![left.clone(), right.clone()]));
         let plan: Arc<dyn ExecutionPlan> =
-            Arc::new(PartitionIsolatorExec::new_ready(union.clone(), 1)?);
+            Arc::new(PartitionIsolatorExec::new_ready(union.clone(), 3)?);
 
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;