stage display formatting as indent and tree done

Author: robtandy

src/physical_optimizer.rs

@@ -105,12 +105,12 @@ struct StagePlanner {
     plan_head: Option<Arc<dyn ExecutionPlan>>,
     /// Current depth in the plan tree, as we walk the tree
     depth: usize,
-    /// Input stages collected so far. Each entry is a tuple of (tree depth, stage).
+    /// Input stages collected so far. Each entry is a tuple of (plan tree depth, stage).
     /// This allows us to keep track of the depth in the plan tree
     /// where we created the stage.   That way when we create a new
     /// stage, we can tell if it is a peer to the current input stages or
     /// should be a parent (if its depth is a smaller number)
-    input_stages: Vec<(usize, Arc<ExecutionStage>)>,
+    input_stages: Vec<(usize, ExecutionStage)>,
     /// current stage number
     stage_counter: usize,
     /// Optional codec to assist in serializing and deserializing any custom
@@ -135,32 +135,45 @@ impl StagePlanner {
     }
 
     fn finish(mut self) -> Result<Arc<ExecutionStage>> {
-        if self.input_stages.is_empty() {
-            Ok(Arc::new(ExecutionStage::new(
+        let stage = if self.input_stages.is_empty() {
+            ExecutionStage::new(
                 self.stage_counter,
                 self.plan_head
                     .take()
                     .ok_or_else(|| internal_datafusion_err!("No plan head set"))?,
                 vec![],
-            )))
+            )
         } else if self.depth < self.input_stages[0].0 {
             // There is more plan above the last stage we created, so we need to
             // create a new stage that includes the last plan head
-            Ok(Arc::new(ExecutionStage::new(
+            ExecutionStage::new(
                 self.stage_counter,
                 self.plan_head
                     .take()
                     .ok_or_else(|| internal_datafusion_err!("No plan head set"))?,
                 self.input_stages
-                    .iter()
-                    .map(|(_, stage)| stage.clone())
+                    .into_iter()
+                    .map(|(_, stage)| Arc::new(stage))
                     .collect(),
-            )))
+            )
         } else {
             // We have a plan head, and we are at the same depth as the last stage we created,
             // so we can just return the last stage
-            Ok(self.input_stages.last().unwrap().1.clone())
+            self.input_stages.last().unwrap().1.clone()
+        };
+
+        // assign the proper tree depth to each stage in the tree
+        fn assign_tree_depth(stage: &ExecutionStage, depth: usize) {
+            stage
+                .depth
+                .store(depth as u64, std::sync::atomic::Ordering::Relaxed);
+            for input in stage.child_stages_iter() {
+                assign_tree_depth(input, depth + 1);
+            }
         }
+        assign_tree_depth(&stage, 0);
+
+        Ok(Arc::new(stage))
     }
 }
 
@@ -191,26 +204,13 @@ impl TreeNodeRewriter for StagePlanner {
                 .map(|(_, stage)| stage.clone())
                 .collect::<Vec<_>>();
 
-            println!(
-                "\n\n\n ---------------------- \ncreating a stage, depth = {}, input_stages ={}, child_stages = {}, plan=\n{}",
-                self.depth,
-                self.input_stages
-                    .iter()
-                    .map(|(depth, stage)| format!("({},{})", depth, stage.num))
-                    .collect::<Vec<_>>()
-                    .join(", "),
-                child_stages
-                    .iter()
-                    .map(|stage| format!("{}", stage.num))
-                    .collect::<Vec<_>>()
-                    .join(", "),
-
-                displayable(plan.as_ref()).indent(false)
-            );
-
             self.input_stages.retain(|(depth, _)| *depth <= self.depth);
 
-            let mut stage = ExecutionStage::new(self.stage_counter, plan.clone(), child_stages);
+            let mut stage = ExecutionStage::new(
+                self.stage_counter,
+                plan.clone(),
+                child_stages.into_iter().map(Arc::new).collect(),
+            );
 
             if let Some(partitions_per_task) = self.partitions_per_task {
                 stage = stage.with_maximum_partitions_per_task(partitions_per_task);
@@ -219,7 +219,7 @@ impl TreeNodeRewriter for StagePlanner {
                 stage = stage.with_codec(codec.clone());
             }
 
-            self.input_stages.push((self.depth, Arc::new(stage)));
+            self.input_stages.push((self.depth, stage));
 
             // As we are walking up the plan tree, we've now put what we've encountered so far
             // into a stage.   We want to replace this plan now with an ArrowFlightReadExec

src/stage/display.rs

@@ -33,12 +33,51 @@ use crate::common::util::display_plan_with_partition_in_out;
 
 use super::ExecutionStage;
 
+// Unicode box-drawing characters for creating borders and connections.
+const LTCORNER: &str = "┌"; // Left top corner
+const RTCORNER: &str = "┐"; // Right top corner
+const LDCORNER: &str = "└"; // Left bottom corner
+const RDCORNER: &str = "┘"; // Right bottom corner
+
+const TMIDDLE: &str = "┬"; // Top T-junction (connects down)
+const LMIDDLE: &str = "├"; // Left T-junction (connects right)
+const DMIDDLE: &str = "┴"; // Bottom T-junction (connects up)
+
+const VERTICAL: &str = "│"; // Vertical line
+const HORIZONTAL: &str = "─"; // Horizontal line
+
 impl DisplayAs for ExecutionStage {
     fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match t {
-            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+            DisplayFormatType::Default => {
                 write!(f, "{}", self.name)
             }
+            DisplayFormatType::Verbose => {
+                writeln!(
+                    f,
+                    "{}{}{}{}",
+                    LTCORNER,
+                    HORIZONTAL.repeat(5),
+                    format!(" {} ", self.name),
+                    HORIZONTAL.repeat(50 - 7 - self.name.len())
+                )?;
+                let plan_str = display_plan_with_partition_in_out(self.plan.as_ref())
+                    .map_err(|_| std::fmt::Error {})?;
+                let plan_str = plan_str.replace(
+                    '\n',
+                    &format!("\n{}{}", "  ".repeat(self.depth()), VERTICAL),
+                );
+                writeln!(f, "{}{}{}", "  ".repeat(self.depth()), VERTICAL, plan_str)?;
+                write!(
+                    f,
+                    "{}{}{}",
+                    "  ".repeat(self.depth()),
+                    LDCORNER,
+                    HORIZONTAL.repeat(50)
+                )?;
+
+                Ok(())
+            }
             DisplayFormatType::TreeRender => write!(
                 f,
                 "{}",

src/stage/stage.rs

@@ -1,3 +1,5 @@
+use std::any::Any;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use datafusion::error::Result;
@@ -7,6 +9,7 @@ use datafusion::prelude::SessionContext;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 
 use itertools::Itertools;
+use tokio::sync::RwLock;
 
 use crate::task::ExecutionTask;
 
@@ -30,6 +33,8 @@ pub struct ExecutionStage {
     pub tasks: Vec<ExecutionTask>,
     /// An optional codec to assist in serializing and deserializing this stage
     pub codec: Option<Arc<dyn PhysicalExtensionCodec>>,
+    /// tree depth of our location in the stage tree, used for display only
+    pub(crate) depth: AtomicU64,
 }
 
 impl Clone for ExecutionStage {
@@ -43,6 +48,7 @@ impl Clone for ExecutionStage {
             inputs: self.inputs.to_vec(),
             tasks: self.tasks.clone(),
             codec: self.codec.clone(),
+            depth: AtomicU64::new(self.depth.load(Ordering::Relaxed)),
         }
     }
 }
@@ -78,6 +84,7 @@ impl ExecutionStage {
                 partition_group,
             }],
             codec: None,
+            depth: AtomicU64::new(0),
         }
     }
 
@@ -138,6 +145,10 @@ impl ExecutionStage {
         };
         format!("Stage {:<3}{}", self.num, child_str)
     }
+
+    pub(crate) fn depth(&self) -> usize {
+        self.depth.load(Ordering::Relaxed) as usize
+    }
 }
 
 impl ExecutionPlan for ExecutionStage {
@@ -164,6 +175,7 @@ impl ExecutionPlan for ExecutionStage {
             inputs: children,
             tasks: self.tasks.clone(),
             codec: self.codec.clone(),
+            depth: AtomicU64::new(self.depth.load(Ordering::Relaxed)),
         }))
     }
 

src/task.rs

@@ -96,7 +96,7 @@ impl Display for ExecutionTask {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "Task[partitions: {}, {}]",
+            "Task: partitions: {}\n{}]",
             format_pg(&self.partition_group),
             self.worker_addr
                 .as_ref()