wip save on graphviz formatting

Author: robtandy

src/stage/display.rs

@@ -29,7 +29,7 @@ use datafusion::{
     },
 };
 
-use crate::common::util::display_plan_with_partition_in_out;
+use crate::{common::util::display_plan_with_partition_in_out, task::format_pg};
 
 use super::ExecutionStage;
 
@@ -96,14 +96,59 @@ pub fn display_stage_graphviz(stage: &ExecutionStage) -> Result<String> {
 
     writeln!(f, "digraph G {{")?;
     writeln!(f, "  node[shape=rect];")?;
+    writeln!(f, "  rankdir=BT;")?;
+    writeln!(f, "  ranksep=2;")?;
+    writeln!(f, "  edge[colorscheme=rdylbu11, penwidth=2.0];")?;
 
-    let mut stack = vec![stage];
+    // we'll keep a stack of stage ref, parrent stage ref
+    let mut stack: Vec<(&ExecutionStage, Option<&ExecutionStage>)> = vec![(stage, None)];
 
-    while !stack.is_empty() {
-        writeln!(f, "  subgraph cluster_{} {{];", stage.num)?;
+    while let Some((stage, parent)) = stack.pop() {
+        writeln!(f, "  subgraph cluster_{} {{", stage.num)?;
         writeln!(f, "    label=\"{}\";", stage.name())?;
-        writeln!(f, "    labeljust=l")?;
+        writeln!(f, "    labeljust=l;")?;
+
+        stage.tasks.iter().try_for_each(|task| {
+            writeln!(
+                f,
+                "    \"{}_{}\"[label = \"{}\"]",
+                stage.num,
+                format_pg(&task.partition_group),
+                format_pg(&task.partition_group)
+            )?;
+
+            if let Some(our_parent) = parent {
+                our_parent.tasks.iter().try_for_each(|ptask| {
+                    ptask.partition_group.iter().try_for_each(|partition| {
+                        writeln!(
+                            f,
+                            "    \"{}_{}\" -> \"{}_{}\"[tailport=n, headport=s, color={}]",
+                            stage.num,
+                            format_pg(&task.partition_group),
+                            our_parent.num,
+                            format_pg(&ptask.partition_group),
+                            partition + 1
+                        )
+                    })?;
+                    writeln!(
+                        f,
+                        "    \"{}_{}\" -> \"{}_{}\"[tailport=n, headport=s, color={}]",
+                        stage.num,
+                        format_pg(&task.partition_group),
+                        our_parent.num,
+                        format_pg(&ptask.partition_group),
+                        &ptask.partition_group[0] + 1,
+                    )
+                })?;
+            }
+
+            Ok::<(), std::fmt::Error>(())
+        })?;
         writeln!(f, "  }}")?;
+
+        for child in stage.child_stages_iter() {
+            stack.push((child, Some(stage)));
+        }
     }
 
     writeln!(f, "}}")?;

src/task.rs

@@ -106,7 +106,7 @@ impl Display for ExecutionTask {
     }
 }
 
-fn format_pg(partition_group: &[u64]) -> String {
+pub(crate) fn format_pg(partition_group: &[u64]) -> String {
     if partition_group.len() > 2 {
         format!(
             "{}..{}",

tests/stage_planning.rs

@@ -10,15 +10,16 @@ mod tests {
     use datafusion::physical_plan::{displayable, execute_stream};
     use datafusion::prelude::{SessionConfig, SessionContext};
     use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
+    use datafusion_distributed::stage::{display_stage_graphviz, ExecutionStage};
     use futures::TryStreamExt;
     use std::error::Error;
     use std::sync::Arc;
 
     #[tokio::test]
     async fn stage_planning() -> Result<(), Box<dyn Error>> {
-        let rule = DistributedPhysicalOptimizerRule::default(); //.with_maximum_partitions_per_task(3);
+        let config = SessionConfig::new().with_target_partitions(3);
 
-        let config = SessionConfig::new().with_target_partitions(10);
+        let rule = DistributedPhysicalOptimizerRule::default().with_maximum_partitions_per_task(2);
 
         let state = SessionStateBuilder::new()
             .with_default_features()
@@ -57,6 +58,10 @@ mod tests {
         let physical_str = displayable(physical.as_ref()).indent(true);
         println!("\n\nPhysical Plan:\n{}", physical_str);
 
+        let physical_str =
+            display_stage_graphviz(physical.as_any().downcast_ref::<ExecutionStage>().unwrap())?;
+        println!("\n\nPhysical Plan:\n{}", physical_str);
+
         assert_snapshot!(physical_str,
             @r"
         ",