Generalize functions for NetworkCoalesceExec creation (#162)

Author: jonathanc-n

src/distributed_physical_optimizer_rule.rs

@@ -160,7 +160,7 @@ impl DistributedPhysicalOptimizerRule {
                 {
                     return Ok(Transformed::no(plan));
                 }
-                let node = NetworkCoalesceExec::from_coalesce_partitions_exec(node, tasks)?;
+                let node = NetworkCoalesceExec::from_input_exec(node, tasks)?;
 
                 let plan = plan.with_new_children(vec![Arc::new(node)])?;
 
@@ -174,7 +174,7 @@ impl DistributedPhysicalOptimizerRule {
                 plan.as_any().downcast_ref::<SortPreservingMergeExec>(),
                 self.network_coalesce_tasks,
             ) {
-                let node = NetworkCoalesceExec::from_sort_preserving_merge_exec(node, tasks)?;
+                let node = NetworkCoalesceExec::from_input_exec(node, tasks)?;
 
                 let plan = plan.with_new_children(vec![Arc::new(node)])?;
 

src/execution_plans/network_coalesce.rs

@@ -16,8 +16,6 @@ use dashmap::DashMap;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
-use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
-use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
 use futures::{TryFutureExt, TryStreamExt};
@@ -101,21 +99,9 @@ pub struct NetworkCoalesceReady {
 }
 
 impl NetworkCoalesceExec {
-    pub fn from_coalesce_partitions_exec(
-        input: &CoalescePartitionsExec,
-        input_tasks: usize,
-    ) -> Result<Self, DataFusionError> {
-        Self::from_input(input, input_tasks)
-    }
-
-    pub fn from_sort_preserving_merge_exec(
-        input: &SortPreservingMergeExec,
-        input_tasks: usize,
-    ) -> Result<Self, DataFusionError> {
-        Self::from_input(input, input_tasks)
-    }
-
-    pub fn from_input(
+    /// Creates a new [NetworkCoalesceExec] node from a [CoalescePartitionsExec] and
+    /// [SortPreservingMergeExec].
+    pub fn from_input_exec(
         input: &dyn ExecutionPlan,
         input_tasks: usize,
     ) -> Result<Self, DataFusionError> {