distributed_physical_optimizer_rule: move metrics_collection to NetworkBoundary trait

This is a small refactor which moves metrics_collection to the NetworkBoundary trait
so all network boundaries must obey the metrics collecting protocol.

Author: jayshrivastava

src/distributed_physical_optimizer_rule.rs

@@ -1,5 +1,8 @@
 use super::{NetworkShuffleExec, PartitionIsolatorExec, StageExec};
 use crate::execution_plans::NetworkCoalesceExec;
+use crate::metrics::proto::MetricsSetProto;
+use crate::protobuf::StageKey;
+use dashmap::DashMap;
 use datafusion::common::plan_err;
 use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::datasource::source::DataSourceExec;
@@ -318,6 +321,20 @@ pub trait NetworkBoundary: ExecutionPlan {
         }
         Ok(Arc::clone(children.first().unwrap()))
     }
+
+    /// metrics_collection is used to collect metrics from child tasks. It is empty when a
+    /// [NetworkBoundary] is instantiated (deserialized, created via new() etc...).
+    /// Metrics are populated by executing() the [NetworkBoundary]. It's expected that the
+    /// collection is complete after the [NetworkBoundary] has been executed. It is undefined
+    /// what this returns during execution.
+    ///
+    /// An instance may receive metrics for 0 to N child tasks, where N is the number of tasks
+    /// in the stage it is reading from. This is because, by convention, the ArrowFlightEndpoint
+    /// sends metrics for a task to the last [NetworkBoundary] to read from it, which may or may
+    /// not be this instance.
+    fn metrics_collection(&self) -> Option<Arc<DashMap<StageKey, Vec<MetricsSetProto>>>> {
+        None
+    }
 }
 
 /// Error thrown during distributed planning that prompts the planner to change something and

src/execution_plans/metrics.rs

@@ -1,3 +1,4 @@
+use crate::distributed_physical_optimizer_rule::NetworkBoundary;
 use crate::execution_plans::{NetworkCoalesceExec, NetworkShuffleExec, StageExec};
 use crate::metrics::proto::{MetricsSetProto, metrics_set_proto_to_df};
 use arrow::ipc::writer::DictionaryTracker;
@@ -50,26 +51,24 @@ impl TreeNodeRewriter for TaskMetricsCollector {
     type Node = Arc<dyn ExecutionPlan>;
 
     fn f_down(&mut self, plan: Self::Node) -> Result<Transformed<Self::Node>> {
-        // If the plan is an NetworkShuffleExec, assume it has collected metrics already
+        // If the plan is a NetwordBoundary, assume it has collected metrics already
         // from child tasks.
         let metrics_collection =
             if let Some(node) = plan.as_any().downcast_ref::<NetworkShuffleExec>() {
-                let NetworkShuffleExec::Ready(ready) = node else {
-                    return internal_err!(
-                        "unexpected NetworkShuffleExec::Pending during metrics collection"
-                    );
-                };
-                Some(Arc::clone(&ready.metrics_collection))
+                node.metrics_collection()
+                    .map(Some)
+                    .ok_or(DataFusionError::Internal(
+                        "could not collect metrics from NetworkShuffleExec".to_string(),
+                    ))
             } else if let Some(node) = plan.as_any().downcast_ref::<NetworkCoalesceExec>() {
-                let NetworkCoalesceExec::Ready(ready) = node else {
-                    return internal_err!(
-                        "unexpected NetworkCoalesceExec::Pending during metrics collection"
-                    );
-                };
-                Some(Arc::clone(&ready.metrics_collection))
+                node.metrics_collection()
+                    .map(Some)
+                    .ok_or(DataFusionError::Internal(
+                        "could not collect metrics from NetworkCoalesceExec".to_string(),
+                    ))
             } else {
-                None
-            };
+                Ok(None)
+            }?;
 
         if let Some(metrics_collection) = metrics_collection {
             for mut entry in metrics_collection.iter_mut() {

src/execution_plans/network_coalesce.rs

@@ -90,15 +90,7 @@ pub struct NetworkCoalesceReady {
     pub(crate) properties: PlanProperties,
     pub(crate) stage_num: usize,
     pub(crate) input_tasks: usize,
-    /// metrics_collection is used to collect metrics from child tasks. It is empty when an
-    /// is instantiated (deserialized, created via [NetworkCoalesceExec::new_ready] etc...).
-    /// Metrics are populated in this map via [NetworkCoalesceExec::execute].
-    ///
-    /// An instance may receive metrics for 0 to N child tasks, where N is the number of tasks in
-    /// the stage it is reading from. This is because, by convention, the ArrowFlightEndpoint
-    /// sends metrics for a task to the last NetworkCoalesceExec to read from it, which may or may
-    /// not be this instance.
-    pub(crate) metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
+    pub(crate) child_task_metrics: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
 }
 
 impl NetworkCoalesceExec {
@@ -164,7 +156,7 @@ impl NetworkBoundary for NetworkCoalesceExec {
             }),
             stage_num,
             input_tasks: pending.input_tasks,
-            metrics_collection: Default::default(),
+            child_task_metrics: Default::default(),
         };
 
         Ok(Arc::new(Self::Ready(ready)))
@@ -185,10 +177,17 @@ impl NetworkBoundary for NetworkCoalesceExec {
                 }),
                 stage_num: ready.stage_num,
                 input_tasks,
-                metrics_collection: Arc::clone(&ready.metrics_collection),
+                child_task_metrics: Arc::clone(&ready.child_task_metrics),
             }),
         })
     }
+
+    fn metrics_collection(&self) -> Option<Arc<DashMap<StageKey, Vec<MetricsSetProto>>>> {
+        match self {
+            NetworkCoalesceExec::Pending(_) => None,
+            NetworkCoalesceExec::Ready(v) => Some(v.child_task_metrics.clone()),
+        }
+    }
 }
 
 impl DisplayAs for NetworkCoalesceExec {
@@ -298,7 +297,7 @@ impl ExecutionPlan for NetworkCoalesceExec {
             return internal_err!("NetworkCoalesceExec: task is unassigned, cannot proceed");
         };
 
-        let metrics_collection_capture = self_ready.metrics_collection.clone();
+        let metrics_collection_capture = self_ready.child_task_metrics.clone();
         let stream = async move {
             let channel = channel_resolver.get_channel_for_url(&url).await?;
             let stream = FlightServiceClient::new(channel)

src/execution_plans/network_shuffle.rs

@@ -136,15 +136,7 @@ pub struct NetworkShuffleReadyExec {
     /// the properties we advertise for this execution plan
     pub(crate) properties: PlanProperties,
     pub(crate) stage_num: usize,
-    /// metrics_collection is used to collect metrics from child tasks. It is empty when an
-    /// is instantiated (deserialized, created via [NetworkShuffleExec::new_ready] etc...).
-    /// Metrics are populated in this map via [NetworkShuffleExec::execute].
-    ///
-    /// An instance may receive metrics for 0 to N child tasks, where N is the number of tasks in
-    /// the stage it is reading from. This is because, by convention, the ArrowFlightEndpoint
-    /// sends metrics for a task to the last NetworkShuffleExec to read from it, which may or may
-    /// not be this instance.
-    pub(crate) metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
+    pub(crate) child_task_metrics: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
 }
 
 impl NetworkShuffleExec {
@@ -209,7 +201,7 @@ impl NetworkBoundary for NetworkShuffleExec {
             NetworkShuffleExec::Ready(prev) => NetworkShuffleExec::Ready(NetworkShuffleReadyExec {
                 properties: prev.properties.clone(),
                 stage_num: prev.stage_num,
-                metrics_collection: Arc::clone(&prev.metrics_collection),
+                child_task_metrics: Arc::clone(&prev.child_task_metrics),
             }),
         })
     }
@@ -226,11 +218,18 @@ impl NetworkBoundary for NetworkShuffleExec {
         let ready = NetworkShuffleReadyExec {
             properties: pending.repartition_exec.properties().clone(),
             stage_num,
-            metrics_collection: Default::default(),
+            child_task_metrics: Default::default(),
         };
 
         Ok(Arc::new(Self::Ready(ready)))
     }
+
+    fn metrics_collection(&self) -> Option<Arc<DashMap<StageKey, Vec<MetricsSetProto>>>> {
+        match self {
+            NetworkShuffleExec::Pending(_) => None,
+            NetworkShuffleExec::Ready(v) => Some(v.child_task_metrics.clone()),
+        }
+    }
 }
 
 impl DisplayAs for NetworkShuffleExec {
@@ -331,7 +330,7 @@ impl ExecutionPlan for NetworkShuffleExec {
                 },
             );
 
-            let metrics_collection_capture = self_ready.metrics_collection.clone();
+            let metrics_collection_capture = self_ready.child_task_metrics.clone();
             async move {
                 let url = task.url.ok_or(internal_datafusion_err!(
                     "NetworkShuffleExec: task is unassigned, cannot proceed"

src/protobuf/distributed_codec.rs

@@ -234,7 +234,7 @@ fn new_network_hash_shuffle_exec(
             Boundedness::Bounded,
         ),
         stage_num,
-        metrics_collection: Default::default(),
+        child_task_metrics: Default::default(),
     })
 }
 
@@ -268,7 +268,7 @@ fn new_network_coalesce_tasks_exec(
         ),
         stage_num,
         input_tasks,
-        metrics_collection: Default::default(),
+        child_task_metrics: Default::default(),
     })
 }