[SPARK-26103][SQL] Limit the length of debug strings for query plans

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala

@@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.errors._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.JoinType
 import org.apache.spark.sql.catalyst.plans.physical.{BroadcastMode, Partitioning}
-import org.apache.spark.sql.catalyst.util.truncatedString
+import org.apache.spark.sql.catalyst.util.{truncatedString, withSizeLimitedWriter}
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.StorageLevel
 
@@ -473,7 +473,9 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
   def treeString(verbose: Boolean, addSuffix: Boolean = false): String = {
     val writer = new StringBuilderWriter()
     try {
-      treeString(writer, verbose, addSuffix)
+      withSizeLimitedWriter(writer) { w =>
+        treeString(w, verbose, addSuffix)
+      }
       writer.toString
     } finally {
       writer.close()

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/SizeLimitedWriter.scala

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.util
+
+import java.io.Writer
+
+class WriterSizeException(val attemptedSize: Long, val charLimit: Long) extends Exception(
+  s"Attempted to write $attemptedSize characters to a writer that is limited to $charLimit")
+
+/**
+ * This class is used to control the size of generated writers.  Guarantees that the total number
+ * of characters written will be less than the specified size.
+ *
+ * Checks size before writing and throws a WriterSizeException if the total size would count the
+ * limit.
+ */
+class SizeLimitedWriter(underlying: Writer, charLimit: Long) extends Writer {
+
+  var charsWritten: Long = 0
+
+  override def write(cbuf: Array[Char], off: Int, len: Int): Unit = {
+    val newLength = charsWritten + Math.min(cbuf.length - off, len)
+    if (newLength > charLimit) {
+      throw new WriterSizeException(newLength, charLimit)
+    }
+    charsWritten = newLength
+    underlying.write(cbuf, off, len)
+  }
+
+  override def flush(): Unit = underlying.flush()
+
+  override def close(): Unit = underlying.close()
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala

@@ -202,6 +202,26 @@ package object util extends Logging {
   /** Shorthand for calling truncatedString() without start or end strings. */
   def truncatedString[T](seq: Seq[T], sep: String): String = truncatedString(seq, "", sep, "")
 
+  /** Whether we have warned about plan string truncation yet. */
+  private val planSizeWarningPrinted = new AtomicBoolean(false)
+
+  def withSizeLimitedWriter[T](writer: Writer)(f: (Writer) => T): Option[T] = {
+    try {
+      val limited = new SizeLimitedWriter(writer, SQLConf.get.maxPlanStringLength)
+      Some(f(limited))
+    }
+    catch {
+      case e: WriterSizeException =>
+        writer.write("...")
+        if (planSizeWarningPrinted.compareAndSet(false, true)) {
+          logWarning(
+            "Truncated the string representation of a plan since it was too long. This " +
+              s"behavior can be adjusted by setting '${SQLConf.MAX_PLAN_STRING_LENGTH.key}'.")
+        }
+        None
+    }
+  }
+
   /* FIX ME
   implicit class debugLogging(a: Any) {
     def debugLogging() {

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1610,6 +1610,12 @@ object SQLConf {
       """ "... N more fields" placeholder.""")
     .intConf
     .createWithDefault(25)
+
+  val MAX_PLAN_STRING_LENGTH = buildConf("spark.sql.debug.maxPlanLength")
+    .doc("Maximum number of characters to output for a plan in debug output.  If the plan is " +
+      "longer, it will end with a ... and further output will be truncated.")
+    .longConf
+    .createWithDefault(8192)
 }
 
 /**
@@ -2030,6 +2036,8 @@ class SQLConf extends Serializable with Logging {
 
   def maxToStringFields: Int = getConf(SQLConf.MAX_TO_STRING_FIELDS)
 
+  def maxPlanStringLength: Long = getConf(SQLConf.MAX_PLAN_STRING_LENGTH)
+
   /** ********************** SQLConf functionality methods ************ */
 
   /** Set Spark SQL configuration properties. */

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala

@@ -21,12 +21,10 @@ import java.math.BigInteger
 import java.util.UUID
 
 import scala.collection.mutable.ArrayBuffer
-
 import org.json4s.JsonAST._
 import org.json4s.JsonDSL._
 import org.json4s.jackson.JsonMethods
 import org.json4s.jackson.JsonMethods._
-
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow, TableIdentifier}
 import org.apache.spark.sql.catalyst.catalog._
@@ -36,6 +34,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
 import org.apache.spark.sql.catalyst.plans.{LeftOuter, NaturalJoin}
 import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, Union}
 import org.apache.spark.sql.catalyst.plans.physical.{IdentityBroadcastMode, RoundRobinPartitioning, SinglePartition}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.StorageLevel
 
@@ -595,4 +594,14 @@ class TreeNodeSuite extends SparkFunSuite {
     val expected = Coalesce(Stream(Literal(1), Literal(3)))
     assert(result === expected)
   }
+
+  test("toString() tree depth") {
+    val ds = (1 until 100).foldLeft(Literal("TestLiteral"): Expression) { case (treeNode, x) =>
+      Add(Literal(x), treeNode)
+    }
+
+    val planString = ds.treeString
+    assert(planString.endsWith("..."))
+    assert(planString.length <= SQLConf.get.maxPlanStringLength)
+  }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/SizeLimitedWriterSuite.scala

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.util
+
+import org.apache.commons.io.output.StringBuilderWriter
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.util.StringUtils._
+
+class SizeLimitedWriterSuite extends SparkFunSuite {
+
+  test("normal writer under limit") {
+    val writer = new StringBuilderWriter()
+    val limited = new SizeLimitedWriter(writer, 100)
+    limited.write("test")
+    limited.write("test")
+
+    assert(writer.toString === "testtest")
+  }
+
+  test("filter pattern") {
+    val writer = new StringBuilderWriter()
+    val limited = new SizeLimitedWriter(writer, 5)
+    assertThrows[WriterSizeException] {
+      limited.write("test")
+      limited.write("test")
+    }
+    assert(writer.toString === "test")
+  }
+}