Return input task count in NetworkBoundary (#204)

Author: gabotechs

src/distributed_planner/network_boundary.rs

@@ -35,6 +35,9 @@ pub trait NetworkBoundary: ExecutionPlan {
         input_tasks: usize,
     ) -> datafusion::common::Result<Arc<dyn NetworkBoundary>>;
 
+    /// Returns the input tasks assigned to this [NetworkBoundary].
+    fn input_task_count(&self) -> usize;
+
     /// Called when a [Stage] is correctly formed. The [NetworkBoundary] can use this
     /// information to perform any internal transformations necessary for distributed execution.
     ///

src/execution_plans/network_coalesce.rs

@@ -176,6 +176,13 @@ impl NetworkBoundary for NetworkCoalesceExec {
             }
         }))
     }
+
+    fn input_task_count(&self) -> usize {
+        match self {
+            Self::Pending(v) => v.input_tasks,
+            Self::Ready(v) => v.input_stage.tasks.len(),
+        }
+    }
 }
 
 impl DisplayAs for NetworkCoalesceExec {

src/execution_plans/network_shuffle.rs

@@ -203,6 +203,13 @@ impl NetworkBoundary for NetworkShuffleExec {
         }))
     }
 
+    fn input_task_count(&self) -> usize {
+        match self {
+            Self::Pending(v) => v.input_tasks,
+            Self::Ready(v) => v.input_stage.tasks.len(),
+        }
+    }
+
     fn with_input_stage(
         &self,
         input_stage: Stage,