From 90b46d9af134ee9d3d674b554e33e14da0af241c Mon Sep 17 00:00:00 2001 From: Cheng Lian Date: Thu, 26 Nov 2015 17:45:13 +0800 Subject: [PATCH 1/2] Prettier tree string for TreeNode --- .../spark/sql/catalyst/trees/TreeNode.scala | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) 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 index 35f087baccde..61a4d03e918b 100644 --- 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 @@ -368,7 +368,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. @@ -395,11 +395,24 @@ 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) - builder.append(simpleString) - builder.append("\n") - children.foreach(_.generateTreeString(depth + 1, builder)) + protected def generateTreeString( + depth: Int, lastChildren: Seq[Boolean], builder: StringBuilder): StringBuilder = { + val prefix = if (depth == 0) { + "" + } else { + (lastChildren.init.map { isLast => + if (isLast) " " else ": " + } :+ (if (lastChildren.last) s"+- " else ":- ")).mkString + } + + val head = prefix + simpleString + builder.append(head + "\n") + + if (children.nonEmpty) { + children.init.foreach(_.generateTreeString(depth + 1, lastChildren :+ false, builder)) + children.last.generateTreeString(depth + 1, lastChildren :+ true, builder) + } + builder } From 31a1c22d0454a5ed1386902ae083c15f5e0cb026 Mon Sep 17 00:00:00 2001 From: Cheng Lian Date: Wed, 2 Dec 2015 18:58:31 +0800 Subject: [PATCH 2/2] Minor styling changes and comments --- .../spark/sql/catalyst/trees/TreeNode.scala | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) 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 index 61a4d03e918b..0bd49a67019b 100644 --- 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 @@ -394,19 +394,27 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product { } } - /** Appends the string represent of this node and its children to the given StringBuilder. */ + /** + * 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 = { - val prefix = if (depth == 0) { - "" - } else { - (lastChildren.init.map { isLast => - if (isLast) " " else ": " - } :+ (if (lastChildren.last) s"+- " else ":- ")).mkString + if (depth > 0) { + lastChildren.init.foreach { isLast => + val prefixFragment = if (isLast) " " else ": " + builder.append(prefixFragment) + } + + val branch = if (lastChildren.last) "+- " else ":- " + builder.append(branch) } - val head = prefix + simpleString - builder.append(head + "\n") + builder.append(simpleString) + builder.append("\n") if (children.nonEmpty) { children.init.foreach(_.generateTreeString(depth + 1, lastChildren :+ false, builder))