[SPARK-12094][SQL] Prettier tree string for TreeNode

liancheng · yhuai · commit f449a407f6f1 · 2015-12-02T09:36:20.000-08:00
When examining plans of complex queries with multiple joins, a pain point of mine is that, it's hard to immediately see the sibling node of a specific query plan node. This PR adds tree lines for the tree string of a `TreeNode`, so that the result can be visually more intuitive. Author: Cheng Lian <lian@databricks.com> Closes apache#10099 from liancheng/prettier-tree-string. (cherry picked from commit a1542ce) Signed-off-by: Yin Huai <yhuai@databricks.com>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
@@ -393,7 +393,7 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
   override def toString: String = treeString
 
   /** Returns a string representation of the nodes in this tree */
-  def treeString: String = generateTreeString(0, new StringBuilder).toString
+  def treeString: String = generateTreeString(0, Nil, new StringBuilder).toString
 
   /**
    * Returns a string representation of the nodes in this tree, where each operator is numbered.
@@ -419,12 +419,33 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
     }
   }
 
-  /** Appends the string represent of this node and its children to the given StringBuilder. */
-  protected def generateTreeString(depth: Int, builder: StringBuilder): StringBuilder = {
-    builder.append(" " * depth)
+  /**
+   * Appends the string represent of this node and its children to the given StringBuilder.
+   *
+   * The `i`-th element in `lastChildren` indicates whether the ancestor of the current node at
+   * depth `i + 1` is the last child of its own parent node.  The depth of the root node is 0, and
+   * `lastChildren` for the root node should be empty.
+   */
+  protected def generateTreeString(
+      depth: Int, lastChildren: Seq[Boolean], builder: StringBuilder): StringBuilder = {
+    if (depth > 0) {
+      lastChildren.init.foreach { isLast =>
+        val prefixFragment = if (isLast) "   " else ":  "
+        builder.append(prefixFragment)
+      }
+
+      val branch = if (lastChildren.last) "+- " else ":- "
+      builder.append(branch)
+    }
+
     builder.append(simpleString)
     builder.append("\n")
-    children.foreach(_.generateTreeString(depth + 1, builder))
+
+    if (children.nonEmpty) {
+      children.init.foreach(_.generateTreeString(depth + 1, lastChildren :+ false, builder))
+      children.last.generateTreeString(depth + 1, lastChildren :+ true, builder)
+    }
+
     builder
   }
 
