[SPARK-24121][SQL] Add API for handling expression code generation

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala

@@ -101,8 +101,7 @@ abstract class Expression extends TreeNode[Expression] {
     ctx.subExprEliminationExprs.get(this).map { subExprState =>
       // This expression is repeated which means that the code to evaluate it has already been added
       // as a function before. In that case, we just re-use it.
-      ExprCode(ctx.registerComment(this.toString), subExprState.isNull,
-        subExprState.value)
+      ExprCode(ctx.registerComment(this.toString), subExprState.isNull, subExprState.value)
     }.getOrElse {
       val isNull = ctx.freshName("isNull")
       val value = ctx.freshName("value")

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.expressions.codegen
 
 import java.lang.{Boolean => JBool}
 
+import scala.collection.mutable.ArrayBuffer
 import scala.language.{existentials, implicitConversions}
 
 import org.apache.spark.sql.types.{BooleanType, DataType}
@@ -130,6 +131,8 @@ trait Block extends JavaCode {
 
   def length: Int = toString.length
 
+  def nonEmpty: Boolean = toString.nonEmpty
+
   // The leading prefix that should be stripped from each line.
   // By default we strip blanks or control characters followed by '|' from the line.
   var _marginChar: Option[Char] = Some('|')
@@ -167,9 +170,40 @@ object Block {
           case other => throw new IllegalArgumentException(
             s"Can not interpolate ${other.getClass.getName} into code block.")
         }
-        CodeBlock(sc.parts, args)
+
+        val (codeParts, blockInputs) = foldLiteralArgs(sc.parts, args)
+        CodeBlock(codeParts, blockInputs)
+      }
+    }
+  }
+
+  // Folds eagerly the literal args into the code parts.
+  private def foldLiteralArgs(parts: Seq[String], args: Seq[Any]): (Seq[String], Seq[Any]) = {
+    val codeParts = ArrayBuffer.empty[String]
+    val blockInputs = ArrayBuffer.empty[Any]
+
+    val strings = parts.iterator
+    val inputs = args.iterator
+    val buf = new StringBuilder(Block.CODE_BLOCK_BUFFER_LENGTH)
+
+    buf append strings.next
+    while (strings.hasNext) {
+      val input = inputs.next
+      input match {
+        case _: ExprValue | _: Block =>
+          codeParts += buf.toString
+          buf.clear
+          blockInputs += input
+        case _ =>
+          buf append input
       }
+      buf append strings.next
+    }
+    if (buf.nonEmpty) {
+      codeParts += buf.toString
     }
+
+    (codeParts.toSeq, blockInputs.toSeq)
   }
 }
 
@@ -182,11 +216,10 @@ case class CodeBlock(codeParts: Seq[String], blockInputs: Seq[Any]) extends Bloc
     blockInputs.flatMap {
       case b: Block => b.exprValues
       case e: ExprValue => Set(e)
-      case _ => Set.empty[ExprValue]
     }.toSet
   }
 
-  override def code: String = {
+  override lazy val code: String = {
     val strings = codeParts.iterator
     val inputs = blockInputs.iterator
     val buf = new StringBuilder(Block.CODE_BLOCK_BUFFER_LENGTH)
@@ -207,7 +240,7 @@ case class CodeBlock(codeParts: Seq[String], blockInputs: Seq[Any]) extends Bloc
 
 case class Blocks(blocks: Seq[Block]) extends Block {
   override lazy val exprValues: Set[ExprValue] = blocks.flatMap(_.exprValues).toSet
-  override def code: String = blocks.map(_.toString).mkString("\n")
+  override lazy val code: String = blocks.map(_.toString).mkString("\n")
 
   override def + (other: Block): Block = other match {
     case c: CodeBlock => Blocks(blocks :+ c)
@@ -217,7 +250,7 @@ case class Blocks(blocks: Seq[Block]) extends Block {
 }
 
 object EmptyBlock extends Block with Serializable {
-  override def code: String = ""
+  override val code: String = ""
   override val exprValues: Set[ExprValue] = Set.empty
 
   override def + (other: Block): Block = other

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/inputFileBlock.scala

@@ -43,8 +43,9 @@ case class InputFileName() extends LeafExpression with Nondeterministic {
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val className = InputFileBlockHolder.getClass.getName.stripSuffix("$")
-    ev.copy(code = code"final ${CodeGenerator.javaType(dataType)} ${ev.value} = " +
-      code"$className.getInputFilePath();", isNull = FalseLiteral)
+    val typeDef = s"final ${CodeGenerator.javaType(dataType)}"
+    ev.copy(code = code"$typeDef ${ev.value} = $className.getInputFilePath();",
+      isNull = FalseLiteral)
   }
 }
 
@@ -66,8 +67,8 @@ case class InputFileBlockStart() extends LeafExpression with Nondeterministic {
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val className = InputFileBlockHolder.getClass.getName.stripSuffix("$")
-    ev.copy(code = code"final ${CodeGenerator.javaType(dataType)} ${ev.value} = " +
-      code"$className.getStartOffset();", isNull = FalseLiteral)
+    val typeDef = s"final ${CodeGenerator.javaType(dataType)}"
+    ev.copy(code = code"$typeDef ${ev.value} = $className.getStartOffset();", isNull = FalseLiteral)
   }
 }
 
@@ -89,7 +90,7 @@ case class InputFileBlockLength() extends LeafExpression with Nondeterministic {
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val className = InputFileBlockHolder.getClass.getName.stripSuffix("$")
-    ev.copy(code = code"final ${CodeGenerator.javaType(dataType)} ${ev.value} = " +
-      code"$className.getLength();", isNull = FalseLiteral)
+    val typeDef = s"final ${CodeGenerator.javaType(dataType)}"
+    ev.copy(code = code"$typeDef ${ev.value} = $className.getLength();", isNull = FalseLiteral)
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeBlockSuite.scala

@@ -23,12 +23,20 @@ import org.apache.spark.sql.types.{BooleanType, IntegerType}
 
 class CodeBlockSuite extends SparkFunSuite {
 
-  test("Block can interpolate string and ExprValue inputs") {
+  test("Block interpolates string and ExprValue inputs") {
     val isNull = JavaCode.isNullVariable("expr1_isNull")
-    val code = code"boolean ${isNull} = ${JavaCode.defaultLiteral(BooleanType)};"
+    val stringLiteral = "false"
+    val code = code"boolean $isNull = $stringLiteral;"
     assert(code.toString == "boolean expr1_isNull = false;")
   }
 
+  test("Literals are folded into string code parts instead of block inputs") {
+    val value = JavaCode.variable("expr1", IntegerType)
+    val intLiteral = 1
+    val code = code"int $value = $intLiteral;"
+    assert(code.asInstanceOf[CodeBlock].blockInputs === Seq(value))
+  }
+
   test("Block.stripMargin") {
     val isNull = JavaCode.isNullVariable("expr1_isNull")
     val value = JavaCode.variable("expr1", IntegerType)
@@ -92,26 +100,26 @@ class CodeBlockSuite extends SparkFunSuite {
   }
 
   test("Throws exception when interpolating unexcepted object in code block") {
-    val obj = TestClass(100)
+    val obj = Tuple2(1, 1)
     val e = intercept[IllegalArgumentException] {
       code"$obj"
     }
     assert(e.getMessage().contains(s"Can not interpolate ${obj.getClass.getName}"))
   }
 
   test("replace expr values in code block") {
-    val statement = JavaCode.expression("1 + 1", IntegerType)
+    val expr = JavaCode.expression("1 + 1", IntegerType)
     val isNull = JavaCode.isNullVariable("expr1_isNull")
     val exprInFunc = JavaCode.variable("expr1", IntegerType)
 
     val code =
       code"""
-           |callFunc(int $statement) {
+           |callFunc(int $expr) {
            |  boolean $isNull = false;
-           |  int $exprInFunc = $statement + 1;
+           |  int $exprInFunc = $expr + 1;
            |}""".stripMargin
 
-    val aliasedParam = JavaCode.variable("aliased", statement.javaType)
+    val aliasedParam = JavaCode.variable("aliased", expr.javaType)
     val aliasedInputs = code.asInstanceOf[CodeBlock].blockInputs.map {
       case _: SimpleExprValue => aliasedParam
       case other => other
@@ -126,5 +134,3 @@ class CodeBlockSuite extends SparkFunSuite {
     assert(aliasedCode.toString == expected.toString)
   }
 }
-
-private case class TestClass(a: Int)

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

@@ -59,7 +59,7 @@ private[sql] trait ColumnarBatchScan extends CodegenSupport {
     }
     val valueVar = ctx.freshName("value")
     val str = s"columnVector[$columnVar, $ordinal, ${dataType.simpleString}]"
-    val code = code"${ctx.registerComment(str)}\n" + (if (nullable) {
+    val code = code"${ctx.registerComment(str)}" + (if (nullable) {
       code"""
         boolean $isNullVar = $columnVar.isNullAt($ordinal);
         $javaType $valueVar = $isNullVar ? ${CodeGenerator.defaultValue(dataType)} : ($value);

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

@@ -260,7 +260,7 @@ trait CodegenSupport extends SparkPlan {
    * them to be evaluated twice.
    */
   protected def evaluateVariables(variables: Seq[ExprCode]): String = {
-    val evaluate = variables.filter(_.code.toString != "").map(_.code.toString).mkString("\n")
+    val evaluate = variables.filter(_.code.nonEmpty).map(_.code.toString).mkString("\n")
     variables.foreach(_.code = EmptyBlock)
     evaluate
   }