[GitHub] [arrow-ballista] avantgardnerio commented on a diff in pull request #310: Add REST API to generate DOT graph for individual query stage

[GitBox <gi...@apache.org>]

avantgardnerio commented on code in PR #310:
URL: https://github.com/apache/arrow-ballista/pull/310#discussion_r985847487


##########
ballista/rust/scheduler/src/api/handlers.rs:
##########
@@ -163,20 +163,41 @@ pub(crate) async fn get_job_summary<T: AsLogicalPlan, U: AsExecutionPlan>(
         })),
     }
 }
+
 /// Generate a dot graph for the specified job id and return as plain text
 pub(crate) async fn get_job_dot_graph<T: AsLogicalPlan, U: AsExecutionPlan>(
     data_server: SchedulerServer<T, U>,
     job_id: String,
 ) -> Result<String, Rejection> {
-    let graph = data_server
+    let maybe_graph = data_server

Review Comment:
   I liked just "graph" :cry: 



##########
ballista/rust/scheduler/src/state/execution_graph/execution_stage.rs:
##########
@@ -73,6 +73,30 @@ impl Debug for ExecutionStage {
     }
 }
 
+impl ExecutionStage {
+    /// Get the name of the variant
+    pub(crate) fn name(&self) -> &str {
+        match self {
+            ExecutionStage::UnResolved(_) => "Unresolved",

Review Comment:
   Should this simply be the `impl Display` of `ExecutionStage`?



##########
ballista/rust/scheduler/src/state/execution_graph_dot.rs:
##########
@@ -74,53 +91,8 @@ impl ExecutionGraphDot {
             let stage = stages.get(id).unwrap(); // safe unwrap
             let stage_name = format!("stage_{}", id);
             writeln!(&mut dot, "\tsubgraph cluster{} {{", cluster)?;
-            match stage {
-                ExecutionStage::UnResolved(stage) => {
-                    writeln!(&mut dot, "\t\tlabel = \"Stage {} [UnResolved]\";", id)?;
-                    stage_meta.push(write_stage_plan(
-                        &mut dot,
-                        &stage_name,
-                        &stage.plan,
-                        0,
-                    )?);
-                }
-                ExecutionStage::Resolved(stage) => {
-                    writeln!(&mut dot, "\t\tlabel = \"Stage {} [Resolved]\";", id)?;
-                    stage_meta.push(write_stage_plan(
-                        &mut dot,
-                        &stage_name,
-                        &stage.plan,
-                        0,
-                    )?);
-                }
-                ExecutionStage::Running(stage) => {
-                    writeln!(&mut dot, "\t\tlabel = \"Stage {} [Running]\";", id)?;
-                    stage_meta.push(write_stage_plan(
-                        &mut dot,
-                        &stage_name,
-                        &stage.plan,
-                        0,
-                    )?);
-                }
-                ExecutionStage::Successful(stage) => {
-                    writeln!(&mut dot, "\t\tlabel = \"Stage {} [Completed]\";", id)?;
-                    stage_meta.push(write_stage_plan(
-                        &mut dot,
-                        &stage_name,
-                        &stage.plan,
-                        0,
-                    )?);
-                }
-                ExecutionStage::Failed(stage) => {
-                    writeln!(&mut dot, "\t\tlabel = \"Stage {} [FAILED]\";", id)?;
-                    stage_meta.push(write_stage_plan(
-                        &mut dot,
-                        &stage_name,
-                        &stage.plan,
-                        0,
-                    )?);
-                }
-            }
+            writeln!(&mut dot, "\t\tlabel = \"Stage {} [{}]\";", id, stage.name())?;

Review Comment:
   Nice refactor :smiley: 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org