feat: Allow tree explain format width to be customizable (#16827)

- Also fix a bug in the tree format.

Author: nuno-faria

datafusion/common/src/config.rs

@@ -840,6 +840,10 @@ config_namespace! {
         /// Display format of explain. Default is "indent".
         /// When set to "tree", it will print the plan in a tree-rendered format.
         pub format: String, default = "indent".to_string()
+
+        /// (format=tree only) Maximum total width of the rendered tree.
+        /// When set to 0, the tree will have no width limit.
+        pub tree_maximum_render_width: usize, default = 240
     }
 }
 

datafusion/core/src/physical_planner.rs

@@ -1856,6 +1856,7 @@ impl DefaultPhysicalPlanner {
                 stringified_plans.push(StringifiedPlan::new(
                     FinalPhysicalPlan,
                     displayable(optimized_plan.as_ref())
+                        .set_tree_maximum_render_width(config.tree_maximum_render_width)
                         .tree_render()
                         .to_string(),
                 ));

datafusion/physical-plan/src/display.rs

@@ -120,6 +120,8 @@ pub struct DisplayableExecutionPlan<'a> {
     show_statistics: bool,
     /// If schema should be displayed. See [`Self::set_show_schema`]
     show_schema: bool,
+    // (TreeRender) Maximum total width of the rendered tree
+    tree_maximum_render_width: usize,
 }
 
 impl<'a> DisplayableExecutionPlan<'a> {
@@ -131,6 +133,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
             show_metrics: ShowMetrics::None,
             show_statistics: false,
             show_schema: false,
+            tree_maximum_render_width: 240,
         }
     }
 
@@ -143,6 +146,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
             show_metrics: ShowMetrics::Aggregated,
             show_statistics: false,
             show_schema: false,
+            tree_maximum_render_width: 240,
         }
     }
 
@@ -155,6 +159,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
             show_metrics: ShowMetrics::Full,
             show_statistics: false,
             show_schema: false,
+            tree_maximum_render_width: 240,
         }
     }
 
@@ -173,6 +178,12 @@ impl<'a> DisplayableExecutionPlan<'a> {
         self
     }
 
+    /// Set the maximum render width for the tree format
+    pub fn set_tree_maximum_render_width(mut self, width: usize) -> Self {
+        self.tree_maximum_render_width = width;
+        self
+    }
+
     /// Return a `format`able structure that produces a single line
     /// per node.
     ///
@@ -270,14 +281,21 @@ impl<'a> DisplayableExecutionPlan<'a> {
     pub fn tree_render(&self) -> impl fmt::Display + 'a {
         struct Wrapper<'a> {
             plan: &'a dyn ExecutionPlan,
+            maximum_render_width: usize,
         }
         impl fmt::Display for Wrapper<'_> {
             fn fmt(&self, f: &mut Formatter) -> fmt::Result {
-                let mut visitor = TreeRenderVisitor { f };
+                let mut visitor = TreeRenderVisitor {
+                    f,
+                    maximum_render_width: self.maximum_render_width,
+                };
                 visitor.visit(self.plan)
             }
         }
-        Wrapper { plan: self.inner }
+        Wrapper {
+            plan: self.inner,
+            maximum_render_width: self.tree_maximum_render_width,
+        }
     }
 
     /// Return a single-line summary of the root of the plan
@@ -540,6 +558,8 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> {
 struct TreeRenderVisitor<'a, 'b> {
     /// Write to this formatter
     f: &'a mut Formatter<'b>,
+    /// Maximum total width of the rendered tree
+    maximum_render_width: usize,
 }
 
 impl TreeRenderVisitor<'_, '_> {
@@ -557,7 +577,6 @@ impl TreeRenderVisitor<'_, '_> {
     const HORIZONTAL: &'static str = "─"; // Horizontal line
 
     // TODO: Make these variables configurable.
-    const MAXIMUM_RENDER_WIDTH: usize = 240; // Maximum total width of the rendered tree
     const NODE_RENDER_WIDTH: usize = 29; // Width of each node's box
     const MAX_EXTRA_LINES: usize = 30; // Maximum number of extra info lines per node
 
@@ -592,6 +611,12 @@ impl TreeRenderVisitor<'_, '_> {
         y: usize,
     ) -> Result<(), fmt::Error> {
         for x in 0..root.width {
+            if self.maximum_render_width > 0
+                && x * Self::NODE_RENDER_WIDTH >= self.maximum_render_width
+            {
+                break;
+            }
+
             if root.has_node(x, y) {
                 write!(self.f, "{}", Self::LTCORNER)?;
                 write!(
@@ -662,7 +687,9 @@ impl TreeRenderVisitor<'_, '_> {
         // Render the actual node.
         for render_y in 0..=extra_height {
             for (x, _) in root.nodes.iter().enumerate().take(root.width) {
-                if x * Self::NODE_RENDER_WIDTH >= Self::MAXIMUM_RENDER_WIDTH {
+                if self.maximum_render_width > 0
+                    && x * Self::NODE_RENDER_WIDTH >= self.maximum_render_width
+                {
                     break;
                 }
 
@@ -780,7 +807,9 @@ impl TreeRenderVisitor<'_, '_> {
         y: usize,
     ) -> Result<(), fmt::Error> {
         for x in 0..=root.width {
-            if x * Self::NODE_RENDER_WIDTH >= Self::MAXIMUM_RENDER_WIDTH {
+            if self.maximum_render_width > 0
+                && x * Self::NODE_RENDER_WIDTH >= self.maximum_render_width
+            {
                 break;
             }
             let mut has_adjacent_nodes = false;