Public method for stage metrics

coralogix · Nov 17, 2022 · a802315 · a802315
1 parent 490bda5
commit a802315
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs
@@ -23,6 +23,7 @@ use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
+use datafusion::physical_plan::metrics::MetricsSet;
 use datafusion::physical_plan::{
     accept, ExecutionPlan, ExecutionPlanVisitor, Partitioning,
 };
@@ -210,6 +211,16 @@ impl ExecutionGraph {
         new_tid
     }
 
+    pub fn stage_metrics(
+        &self,
+    ) -> HashMap<usize, (&dyn ExecutionPlan, Option<&Vec<MetricsSet>>)> {
+        let mut result = HashMap::with_capacity(self.stages.len());
+        for (stage_id, stage) in self.stages.iter() {
+            result.insert(*stage_id, (stage.plan(), stage.stage_metrics()));
+        }
+        result
+    }
+
     pub(crate) fn stages(&self) -> &HashMap<usize, ExecutionStage> {
         &self.stages
     }

diff --git a/ballista/scheduler/src/state/execution_graph/execution_stage.rs b/ballista/scheduler/src/state/execution_graph/execution_stage.rs
@@ -97,6 +97,21 @@ impl ExecutionStage {
             ExecutionStage::Failed(stage) => stage.plan.as_ref(),
         }
     }
+
+    pub(crate) fn stage_metrics(&self) -> Option<&Vec<MetricsSet>> {
+        match self {
+            ExecutionStage::UnResolved(_) | ExecutionStage::Resolved(_) => None,
+            ExecutionStage::Running(RunningStage { stage_metrics, .. }) => {
+                stage_metrics.as_ref()
+            }
+            ExecutionStage::Successful(SuccessfulStage { stage_metrics, .. }) => {
+                Some(stage_metrics)
+            }
+            ExecutionStage::Failed(FailedStage { stage_metrics, .. }) => {
+                stage_metrics.as_ref()
+            }
+        }
+    }
 }
 
 /// For a stage whose input stages are not all completed, we say it's a unresolved stage