Added plan size limits to StringConcat

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala

@@ -308,10 +308,10 @@ object QueryPlan extends PredicateHelper {
    */
   def append[T <: QueryPlan[T]](
       plan: => QueryPlan[T],
-      append: String => Unit,
+      append: String => Boolean,
       verbose: Boolean,
       addSuffix: Boolean,
-      maxFields: Int = SQLConf.get.maxToStringFields): Unit = {
+      maxFields: Int = SQLConf.get.maxToStringFields): Boolean = {
     try {
       plan.treeString(append, verbose, addSuffix, maxFields)
     } catch {

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

@@ -487,10 +487,10 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
   }
 
   def treeString(
-      append: String => Unit,
+      append: String => Boolean,
       verbose: Boolean,
       addSuffix: Boolean,
-      maxFields: Int): Unit = {
+      maxFields: Int): Boolean = {
     generateTreeString(0, Nil, append, verbose, "", addSuffix, maxFields)
   }
 
@@ -554,11 +554,11 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
   def generateTreeString(
       depth: Int,
       lastChildren: Seq[Boolean],
-      append: String => Unit,
+      append: String => Boolean,
       verbose: Boolean,
       prefix: String = "",
       addSuffix: Boolean = false,
-      maxFields: Int): Unit = {
+      maxFields: Int): Boolean = {
 
     if (depth > 0) {
       lastChildren.init.foreach { isLast =>
@@ -591,6 +591,7 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
       children.last.generateTreeString(
         depth + 1, lastChildren :+ true, append, verbose, prefix, addSuffix, maxFields)
     }
+    append("")
   }
 
   /**

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

@@ -17,14 +17,17 @@
 
 package org.apache.spark.sql.catalyst.util
 
+import java.util.concurrent.atomic.AtomicBoolean
 import java.util.regex.{Pattern, PatternSyntaxException}
 
 import scala.collection.mutable.ArrayBuffer
 
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.unsafe.types.UTF8String
 
-object StringUtils {
+object StringUtils extends Logging {
 
   /**
    * Validate and convert SQL 'like' pattern to a Java regular expression.
@@ -92,31 +95,57 @@ object StringUtils {
 
   /**
    * Concatenation of sequence of strings to final string with cheap append method
-   * and one memory allocation for the final string.
+   * and one memory allocation for the final string.  Can also bound the final size of
+   * the string.
    */
-  class StringConcat {
+  class StringConcat(val maxLength: Int = Integer.MAX_VALUE) {
     private val strings = new ArrayBuffer[String]
     private var length: Int = 0
 
+    def atLimit: Boolean = length >= maxLength
+
     /**
      * Appends a string and accumulates its length to allocate a string buffer for all
-     * appended strings once in the toString method.
+     * appended strings once in the toString method.  Returns true if the string still
+     * has room for further appends before it hits its max limit.
      */
-    def append(s: String): Unit = {
-      if (s != null) {
+    def append(s: String): Boolean = {
+      if (!atLimit && s != null) {
         strings.append(s)
         length += s.length
       }
+      return !atLimit
     }
 
     /**
      * The method allocates memory for all appended strings, writes them to the memory and
      * returns concatenated string.
      */
     override def toString: String = {
-      val result = new java.lang.StringBuilder(length)
-      strings.foreach(result.append)
+      val finalLength = Math.min(length, maxLength)
+      val result = new java.lang.StringBuilder(finalLength)
+      strings.dropRight(1).foreach(result.append)
+      strings.lastOption.foreach { s =>
+        val lastLength = Math.min(s.length, maxLength - result.length())
+        result.append(s, 0, lastLength)
+      }
       result.toString
     }
   }
+
+  /** Whether we have warned about plan string truncation yet. */
+  private val planSizeWarningPrinted = new AtomicBoolean(false)
+
+  /** A string concatenator for plan strings.  Uses length from a configured value, and
+   *  prints a warning the first time a plan is truncated. */
+  class PlanStringConcat extends StringConcat(SQLConf.get.maxPlanStringLength) {
+    override def toString: String = {
+      if (atLimit && 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}'.")
+      }
+      super.toString
+    }
+  }
 }

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

@@ -204,27 +204,6 @@ 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

@@ -1627,12 +1627,11 @@ object SQLConf {
 
   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)
+      "longer, 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.")
+    .intConf
+    .createWithDefault(Int.MaxValue)
 
   val SET_COMMAND_REJECTS_SPARK_CORE_CONFS =
     buildConf("spark.sql.legacy.setCommandRejectsSparkCoreConfs")
@@ -2065,7 +2064,7 @@ 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 maxPlanStringLength: Int = 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/util/StringUtilsSuite.scala

@@ -56,4 +56,26 @@ class StringUtilsSuite extends SparkFunSuite {
     assert(concat("1", "2") == "12")
     assert(concat("abc", "\n", "123") == "abc\n123")
   }
+
+  test("string concatenation with limit") {
+    def concat(seq: String*): String = {
+      seq.foldLeft(new StringConcat(7))((acc, s) => {acc.append(s); acc}).toString
+    }
+    assert(concat("under") == "under")
+    assert(concat("under", "over", "extra") == "underov")
+    assert(concat("underover") == "underov")
+    assert(concat("under", "ov") == "underov")
+  }
+
+  test("string concatenation return value") {
+    assert(new StringConcat(7).append("under") == true)
+    assert(new StringConcat(7).append("underover") == false)
+    assert(new StringConcat(7).append("underov") == false)
+  }
+
+  test("string concatenation append after limit") {
+    val concat = new StringConcat(7)
+    concat.append("underover")
+    assert(concat.append("extra") == false)
+  }
 }

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

@@ -20,17 +20,15 @@ package org.apache.spark.sql.execution
 import java.io.{BufferedWriter, OutputStreamWriter}
 
 import org.apache.hadoop.fs.Path
-
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.{InternalRow, QueryPlanningTracker}
 import org.apache.spark.sql.catalyst.analysis.UnsupportedOperationChecker
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 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.StringUtils.{PlanStringConcat, 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
@@ -106,22 +104,14 @@ class QueryExecution(
     ReuseSubquery(sparkSession.sessionState.conf))
 
   def simpleString: String = withRedaction {
-    val concat = new StringConcat()
+    val concat = new PlanStringConcat()
     concat.append("== Physical Plan ==\n")
     QueryPlan.append(executedPlan, concat.append, verbose = false, addSuffix = false)
     concat.append("\n")
     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 = {
+  private def writePlans(append: String => Boolean, maxFields: Int): Boolean = {
     val (verbose, addSuffix) = (true, false)
     append("== Parsed Logical Plan ==\n")
     QueryPlan.append(logical, append, verbose, addSuffix, maxFields)
@@ -142,13 +132,13 @@ class QueryExecution(
   }
 
   override def toString: String = withRedaction {
-    val concat = new StringConcat()
+    val concat = new PlanStringConcat()
     writePlans(concat.append, SQLConf.get.maxToStringFields)
     concat.toString
   }
 
   def stringWithStats: String = withRedaction {
-    val concat = new StringConcat()
+    val concat = new PlanStringConcat()
     val maxFields = SQLConf.get.maxToStringFields
 
     // trigger to compute stats for logical plans
@@ -203,9 +193,12 @@ class QueryExecution(
       val filePath = new Path(path)
       val fs = filePath.getFileSystem(sparkSession.sessionState.newHadoopConf())
       val writer = new BufferedWriter(new OutputStreamWriter(fs.create(filePath)))
-
+      val append = (s: String) => {
+        writer.write(s)
+        true
+      }
       try {
-        writePlans(writer.write, maxFields)
+        writePlans(append, maxFields)
         writer.write("\n== Whole Stage Codegen ==\n")
         org.apache.spark.sql.execution.debug.writeCodegen(writer.write, executedPlan)
       } finally {