From b035bbe7a2a8395bbe1d1163dc590510dab62b79 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Sat, 23 Jun 2018 07:35:41 +0000
Subject: [PATCH 1/2] Remove Blocks class.

---
 .../expressions/codegen/javaCode.scala        | 37 ++++++-------------
 1 file changed, 12 insertions(+), 25 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala
index 250ce48d059e..bdc1c572a79c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala
@@ -119,6 +119,7 @@ object JavaCode {
  * A trait representing a block of java code.
  */
 trait Block extends JavaCode {
+  import Block._
 
   // The expressions to be evaluated inside this block.
   def exprValues: Set[ExprValue]
@@ -148,14 +149,17 @@ trait Block extends JavaCode {
   }
 
   // Concatenates this block with other block.
-  def + (other: Block): Block
+  def + (other: Block): Block = other match {
+    case EmptyBlock => this
+    case _ => code"$this\n$other"
+  }
 }
 
 object Block {
 
   val CODE_BLOCK_BUFFER_LENGTH: Int = 512
 
-  implicit def blocksToBlock(blocks: Seq[Block]): Block = Blocks(blocks)
+  implicit def blocksToBlock(blocks: Seq[Block]): Block = blocks.reduceLeft(_ + _)
 
   implicit class BlockHelper(val sc: StringContext) extends AnyVal {
     def code(args: Any*): Block = {
@@ -199,9 +203,7 @@ object Block {
       }
       buf.append(strings.next)
     }
-    if (buf.nonEmpty) {
-      codeParts += buf.toString
-    }
+    codeParts += buf.toString
 
     (codeParts.toSeq, blockInputs.toSeq)
   }
@@ -209,7 +211,11 @@ object Block {
 
 /**
  * A block of java code. Including a sequence of code parts and some inputs to this block.
- * The actual java code is generated by embedding the inputs into the code parts.
+ * The actual java code is generated by embedding the inputs into the code parts. Here we keep
+ * inputs of `JavaCode` instead of simply folding them as a string of code, because we need to
+ * track expressions (`ExprValue`) in this code block. We need to be able to manipulate the
+ * expressions later without changing the behavior of this code block in some applications, e.g.,
+ * method splitting.
  */
 case class CodeBlock(codeParts: Seq[String], blockInputs: Seq[JavaCode]) extends Block {
   override lazy val exprValues: Set[ExprValue] = {
@@ -230,30 +236,11 @@ case class CodeBlock(codeParts: Seq[String], blockInputs: Seq[JavaCode]) extends
     }
     buf.toString
   }
-
-  override def + (other: Block): Block = other match {
-    case c: CodeBlock => Blocks(Seq(this, c))
-    case b: Blocks => Blocks(Seq(this) ++ b.blocks)
-    case EmptyBlock => this
-  }
-}
-
-case class Blocks(blocks: Seq[Block]) extends Block {
-  override lazy val exprValues: Set[ExprValue] = blocks.flatMap(_.exprValues).toSet
-  override lazy val code: String = blocks.map(_.toString).mkString("\n")
-
-  override def + (other: Block): Block = other match {
-    case c: CodeBlock => Blocks(blocks :+ c)
-    case b: Blocks => Blocks(blocks ++ b.blocks)
-    case EmptyBlock => this
-  }
 }
 
 object EmptyBlock extends Block with Serializable {
   override val code: String = ""
   override val exprValues: Set[ExprValue] = Set.empty
-
-  override def + (other: Block): Block = other
 }
 
 /**

From 9e1439720a55f5c80169ddb367a1095db188698b Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Sat, 23 Jun 2018 14:02:19 +0000
Subject: [PATCH 2/2] Early expand EmptyBlock.

---
 .../spark/sql/catalyst/expressions/codegen/javaCode.scala      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala
index bdc1c572a79c..44f63e21e93b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala
@@ -194,10 +194,11 @@ object Block {
     while (strings.hasNext) {
       val input = inputs.next
       input match {
-        case _: ExprValue | _: Block =>
+        case _: ExprValue | _: CodeBlock =>
           codeParts += buf.toString
           buf.clear
           blockInputs += input.asInstanceOf[JavaCode]
+        case EmptyBlock =>
         case _ =>
           buf.append(input)
       }