Add basic tests.

Author: viirya

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

@@ -114,15 +114,17 @@ object JavaCode {
   }
 
   def block(code: String): Block = {
-    CodeBlock(codeParts = Seq(code), exprValues = Seq.empty)
+    CodeBlock(codeParts = Seq(code), blockInputs = Seq.empty)
   }
 }
 
 /**
- * A block of java code which involves some expressions represented by `ExprValue`.
+ * A trait representing a block of java code.
  */
 trait Block extends JavaCode {
-  def exprValues: Seq[Any]
+
+  // The expressions to be evaluated inside this block.
+  def exprValues: Seq[ExprValue]
 
   // This will be called during string interpolation.
   override def toString: String = _marginChar match {
@@ -169,16 +171,20 @@ object Block {
 }
 
 /**
- * A block of java code.
+ * 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.
  */
-case class CodeBlock(codeParts: Seq[String], exprValues: Seq[Any]) extends Block {
+case class CodeBlock(codeParts: Seq[String], blockInputs: Seq[Any]) extends Block {
+  override def exprValues: Seq[ExprValue] = blockInputs.filter(_.isInstanceOf[ExprValue])
+    .asInstanceOf[Seq[ExprValue]]
+
   override def code: String = {
     val strings = codeParts.iterator
-    val expressions = exprValues.iterator
+    val inputs = blockInputs.iterator
     var buf = new StringBuffer(strings.next)
     while (strings.hasNext) {
-      if (expressions.hasNext) {
-        buf append expressions.next
+      if (inputs.hasNext) {
+        buf append inputs.next
         buf append strings.next
       }
     }
@@ -193,7 +199,7 @@ case class CodeBlock(codeParts: Seq[String], exprValues: Seq[Any]) extends Block
 }
 
 case class Blocks(blocks: Seq[Block]) extends Block {
-  override def exprValues: Seq[Any] = blocks.flatMap(_.exprValues)
+  override def exprValues: Seq[ExprValue] = blocks.flatMap(_.exprValues)
   override def code: String = blocks.map(_.toString).mkString
 
   override def + (other: Block): Block = other match {
@@ -205,7 +211,7 @@ case class Blocks(blocks: Seq[Block]) extends Block {
 
 object EmptyBlock extends Block with Serializable {
   override def code: String = ""
-  override def exprValues: Seq[Any] = Seq.empty
+  override def exprValues: Seq[ExprValue] = Seq.empty
 
   override def + (other: Block): Block = other
 }

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

@@ -0,0 +1,98 @@
+/*
+ * 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.expressions.codegen
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.types.{BooleanType, IntegerType}
+
+class CodeBlockSuite extends SparkFunSuite {
+
+  test("Block can interpolate string and ExprValue inputs") {
+    val isNull = JavaCode.isNullVariable("expr1_isNull")
+    val code = code"boolean ${isNull} = ${JavaCode.defaultLiteral(BooleanType)};"
+    assert(code.toString == "boolean expr1_isNull = false;")
+  }
+
+  test("Block.stripMargin") {
+    val isNull = JavaCode.isNullVariable("expr1_isNull")
+    val value = JavaCode.variable("expr1", IntegerType)
+    val code1 =
+      code"""
+           |boolean $isNull = false;
+           |int $value = ${JavaCode.defaultLiteral(IntegerType)};""".stripMargin
+    val expected =
+      s"""
+        |boolean expr1_isNull = false;
+        |int expr1 = ${JavaCode.defaultLiteral(IntegerType)};""".stripMargin
+    assert(code1.toString == expected)
+
+    val code2 =
+      code"""
+           >boolean $isNull = false;
+           >int $value = ${JavaCode.defaultLiteral(IntegerType)};""".stripMargin('>')
+    assert(code2.toString == expected)
+  }
+
+  test("Block can capture input expr values") {
+    val isNull = JavaCode.isNullVariable("expr1_isNull")
+    val value = JavaCode.variable("expr1", IntegerType)
+    val code =
+      code"""
+           |boolean $isNull = false;
+           |int $value = -1;
+          """.stripMargin
+    val exprValues = code.exprValues
+    assert(exprValues.length == 2)
+    assert(exprValues(0).isInstanceOf[VariableValue] && exprValues(0).code == "expr1_isNull")
+    assert(exprValues(1).isInstanceOf[VariableValue] && exprValues(1).code == "expr1")
+  }
+
+  test("concatenate blocks") {
+    val isNull1 = JavaCode.isNullVariable("expr1_isNull")
+    val value1 = JavaCode.variable("expr1", IntegerType)
+    val isNull2 = JavaCode.isNullVariable("expr2_isNull")
+    val value2 = JavaCode.variable("expr2", IntegerType)
+    val literal = JavaCode.literal("100", IntegerType)
+
+    val code =
+      code"""
+           |boolean $isNull1 = false;
+           |int $value1 = -1;""".stripMargin +
+      code"""
+           |boolean $isNull2 = true;
+           |int $value2 = $literal;""".stripMargin
+
+    val expected =
+      """
+       |boolean expr1_isNull = false;
+       |int expr1 = -1;
+       |boolean expr2_isNull = true;
+       |int expr2 = 100;""".stripMargin
+
+    assert(code.toString == expected)
+
+    val exprValues = code.exprValues
+    assert(exprValues.length == 5)
+    assert(exprValues(0).isInstanceOf[VariableValue] && exprValues(0).code == "expr1_isNull")
+    assert(exprValues(1).isInstanceOf[VariableValue] && exprValues(1).code == "expr1")
+    assert(exprValues(2).isInstanceOf[VariableValue] && exprValues(2).code == "expr2_isNull")
+    assert(exprValues(3).isInstanceOf[VariableValue] && exprValues(3).code == "expr2")
+    assert(exprValues(4).isInstanceOf[LiteralValue] && exprValues(4).code == "100")
+  }
+}