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

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 extraChars: Long, val charLimit: Long) extends Exception(
+  s"Writer reached limit of $charLimit characters.  $extraChars extra characters ignored.")
+
+/**
+ * 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 {
+
+  private var charsWritten: Long = 0
+
+  override def write(cbuf: Array[Char], off: Int, len: Int): Unit = {
+    val charsToWrite = Math.min(charLimit - charsWritten, len).toInt
+    underlying.write(cbuf, off, charsToWrite)
+    charsWritten += charsToWrite
+    if (charsToWrite < len) {
+      throw new WriterSizeException(len - charsToWrite, charLimit)
+    }
+  }
+
+  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

@@ -204,6 +204,27 @@ package object util extends Logging {
     truncatedString(seq, "", sep, "", maxFields)
   }
 
+  /** 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 {
+      // Subtract 3 from the string length to leave room for the "..."
+      val limited = new SizeLimitedWriter(writer, SQLConf.get.maxPlanStringLength - 3)
+      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

@@ -1625,6 +1625,15 @@ object SQLConf {
     .intConf
     .createWithDefault(25)
 
+  val MAX_PLAN_STRING_LENGTH = buildConf("spark.sql.maxPlanLength")
+    .doc("Maximum number of characters to output for a plan string.  If the plan is " +
+      "longer, it will end with \"...\" and further output will be truncated.  The default " +
+      "setting always generates a full plan.  Set this to a lower value such as 8192 if plan " +
+      "strings are taking up too much memory or are causing OutOfMemory errors in the driver or " +
+      "UI processes.")
+    .longConf
+    .createWithDefault(Long.MaxValue)
+
   val SET_COMMAND_REJECTS_SPARK_CORE_CONFS =
     buildConf("spark.sql.legacy.setCommandRejectsSparkCoreConfs")
       .internal()
@@ -2061,6 +2070,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)
+
   def setCommandRejectsSparkCoreConfs: Boolean =
     getConf(SQLConf.SET_COMMAND_REJECTS_SPARK_CORE_CONFS)
 

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

@@ -33,9 +33,10 @@ import org.apache.spark.sql.catalyst.catalog._
 import org.apache.spark.sql.catalyst.dsl.expressions.DslString
 import org.apache.spark.sql.catalyst.expressions._
 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.{LeftOuter, NaturalJoin, SQLHelper}
 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
 
@@ -81,7 +82,7 @@ case class SelfReferenceUDF(
   def apply(key: String): Boolean = config.contains(key)
 }
 
-class TreeNodeSuite extends SparkFunSuite {
+class TreeNodeSuite extends SparkFunSuite with SQLHelper {
   test("top node changed") {
     val after = Literal(1) transform { case Literal(1, _) => Literal(2) }
     assert(after === Literal(2))
@@ -595,4 +596,25 @@ class TreeNodeSuite extends SparkFunSuite {
     val expected = Coalesce(Stream(Literal(1), Literal(3)))
     assert(result === expected)
   }
+
+  test("treeString default doesn't limit plan length") {
+    val ds = (1 until 20).foldLeft(Literal("TestLiteral"): Expression) { case (treeNode, x) =>
+      Add(Literal(x), treeNode)
+    }
+
+    val planString = ds.treeString
+    assert(!planString.endsWith("..."))
+  }
+
+  test("treeString limits plan length") {
+    withSQLConf(SQLConf.MAX_PLAN_STRING_LENGTH.key -> "200") {
+      val ds = (1 until 20).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("truncate at the limit") {
+    val writer = new StringBuilderWriter()
+    val limited = new SizeLimitedWriter(writer, 5)
+    assertThrows[WriterSizeException] {
+      limited.write("test")
+      limited.write("test")
+    }
+    assert(writer.toString === "testt")
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala

@@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, ReturnAnswer}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.StringUtils.StringConcat
 import org.apache.spark.sql.catalyst.util.truncatedString
+import org.apache.spark.sql.catalyst.util.withSizeLimitedWriter
 import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReuseExchange}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.util.Utils
@@ -112,6 +113,14 @@ class QueryExecution(
     concat.toString
   }
 
+//  private def writeOrError(writer: Writer)(f: Writer => Unit): Unit = {
+//    try
+//      withSizeLimitedWriter(writer)(f)
+//    catch {
+//      case e: AnalysisException => writer.write(e.toString)
+//    }
+//  }
+//
   private def writePlans(append: String => Unit, maxFields: Int): Unit = {
     val (verbose, addSuffix) = (true, false)
     append("== Parsed Logical Plan ==\n")