WIP: Remove old method for reference binding, more work on configuration.

Author: marmbrus

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

@@ -172,7 +172,7 @@ package object dsl {
       // Protobuf terminology
       def required = a.withNullability(false)
 
-      def at(ordinal: Int) = BoundReference(ordinal, a)
+      def at(ordinal: Int) = BoundReference(ordinal, a.dataType, a.nullable)
     }
   }
 

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

@@ -17,72 +17,40 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
-import org.apache.spark.sql.catalyst.trees
 import org.apache.spark.sql.catalyst.errors.attachTree
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.types._
+import org.apache.spark.sql.catalyst.trees
+
 import org.apache.spark.sql.Logging
 
 /**
  * A bound reference points to a specific slot in the input tuple, allowing the actual value
  * to be retrieved more efficiently.  However, since operations like column pruning can change
  * the layout of intermediate tuples, BindReferences should be run after all such transformations.
  */
-case class BoundReference(ordinal: Int, baseReference: Attribute)
-  extends Attribute with trees.LeafNode[Expression] {
+case class BoundReference(ordinal: Int, dataType: DataType, nullable: Boolean)
+  extends Expression with trees.LeafNode[Expression] {
 
   type EvaluatedType = Any
 
-  override def nullable = baseReference.nullable
-  override def dataType = baseReference.dataType
-  override def exprId = baseReference.exprId
-  override def qualifiers = baseReference.qualifiers
-  override def name = baseReference.name
+  def references = Set.empty
 
-  override def newInstance = BoundReference(ordinal, baseReference.newInstance)
-  override def withNullability(newNullability: Boolean) =
-    BoundReference(ordinal, baseReference.withNullability(newNullability))
-  override def withQualifiers(newQualifiers: Seq[String]) =
-    BoundReference(ordinal, baseReference.withQualifiers(newQualifiers))
-
-  override def toString = s"$baseReference:$ordinal"
+  override def toString = s"input[$ordinal]"
 
   override def eval(input: Row): Any = input(ordinal)
 }
 
-/**
- * Used to denote operators that do their own binding of attributes internally.
- */
-trait NoBind { self: trees.TreeNode[_] => }
-
-class BindReferences[TreeNode <: QueryPlan[TreeNode]] extends Rule[TreeNode] {
-  import BindReferences._
-
-  def apply(plan: TreeNode): TreeNode = {
-    plan.transform {
-      case n: NoBind => n.asInstanceOf[TreeNode]
-      case leafNode if leafNode.children.isEmpty => leafNode
-      case unaryNode if unaryNode.children.size == 1 => unaryNode.transformExpressions { case e =>
-        bindReference(e, unaryNode.children.head.output)
-      }
-    }
-  }
-}
-
 object BindReferences extends Logging {
   def bindReference[A <: Expression](expression: A, input: Seq[Attribute]): A = {
     expression.transform { case a: AttributeReference =>
       attachTree(a, "Binding attribute") {
         val ordinal = input.indexWhere(_.exprId == a.exprId)
         if (ordinal == -1) {
-          // TODO: This fallback is required because some operators (such as ScriptTransform)
-          // produce new attributes that can't be bound.  Likely the right thing to do is remove
-          // this rule and require all operators to explicitly bind to the input schema that
-          // they specify.
-          logger.debug(s"Couldn't find $a in ${input.mkString("[", ",", "]")}")
-          a
+          sys.error(s"Couldn't find $a in ${input.mkString("[", ",", "]")}")
         } else {
-          BoundReference(ordinal, a)
+          BoundReference(ordinal, a.dataType, a.nullable)
         }
       }
     }.asInstanceOf[A] // Kind of a hack, but safe.  TODO: Tighten return type when possible.

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.expressions.codegen
 
+import com.google.common.cache.{CacheLoader, CacheBuilder}
+
 import scala.language.existentials
 
 import org.apache.spark.Logging
@@ -26,27 +28,53 @@ import org.apache.spark.sql.catalyst.types._
 
 /**
  * A base class for generators of byte code that performs expression evaluation.  Includes helpers
- * for refering to Catalyst types and building trees that perform evaluation of individual
+ * for referring to Catalyst types and building trees that perform evaluation of individual
  * expressions.
  */
-abstract class CodeGenerator extends Logging {
+abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Logging {
   import scala.reflect.runtime.{universe => ru}
   import scala.reflect.runtime.universe._
 
   import scala.tools.reflect.ToolBox
 
-  val toolBox = runtimeMirror(getClass.getClassLoader).mkToolBox()
+  protected val toolBox = runtimeMirror(getClass.getClassLoader).mkToolBox()
 
-  val rowType = typeOf[Row]
-  val mutableRowType = typeOf[MutableRow]
-  val genericRowType = typeOf[GenericRow]
-  val genericMutableRowType = typeOf[GenericMutableRow]
+  protected val rowType = typeOf[Row]
+  protected val mutableRowType = typeOf[MutableRow]
+  protected val genericRowType = typeOf[GenericRow]
+  protected val genericMutableRowType = typeOf[GenericMutableRow]
 
-  val projectionType = typeOf[Projection]
-  val mutableProjectionType = typeOf[MutableProjection]
+  protected val projectionType = typeOf[Projection]
+  protected val mutableProjectionType = typeOf[MutableProjection]
 
   private val curId = new java.util.concurrent.atomic.AtomicInteger()
-  private val javaSeperator = "$"
+  private val javaSeparator = "$"
+
+  /**
+   * Generates a class for a given input expression.  Called when there is not a cached code
+   * already available.
+   */
+  protected def create(in: InType): OutType
+
+  /** Canonicalizes an input expression. */
+  protected def canonicalize(in: InType): InType
+
+  /** Binds an input expression to a given input schema */
+  protected def bind(in: InType, inputSchema: Seq[Attribute]): InType
+
+  protected val cache = CacheBuilder.newBuilder()
+    .maximumSize(1000)
+    .build(
+      new CacheLoader[InType, OutType]() {
+        override def load(in: InType): OutType = globalLock.synchronized {
+           create(in)
+        }
+      })
+
+  def apply(expressions: InType, inputSchema: Seq[Attribute]): OutType=
+    apply(bind(expressions, inputSchema))
+
+  def apply(expressions: InType): OutType = cache.get(canonicalize(expressions))
 
   /**
    * Returns a term name that is unique within this instance of a `CodeGenerator`.
@@ -55,7 +83,7 @@ abstract class CodeGenerator extends Logging {
    * function.)
    */
   protected def freshName(prefix: String): TermName = {
-    newTermName(s"$prefix$javaSeperator${curId.getAndIncrement}")
+    newTermName(s"$prefix$javaSeparator${curId.getAndIncrement}")
   }
 
   /**
@@ -66,7 +94,7 @@ abstract class CodeGenerator extends Logging {
    *                 to null.
    * @param primitiveTerm A term for a possible primitive value of the result of the evaluation. Not
    *                      valid if `nullTerm` is set to `false`.
-   * @param objectTerm An possibly boxed version of the result of evaluating this expression.
+   * @param objectTerm A possibly boxed version of the result of evaluating this expression.
    */
   protected case class EvaluatedExpression(
       code: Seq[Tree],
@@ -87,7 +115,7 @@ abstract class CodeGenerator extends Logging {
       def castOrNull(f: TermName => Tree, dataType: DataType): Seq[Tree] = {
         val eval = expressionEvaluator(e)
         eval.code ++
-          q"""
+        q"""
           val $nullTerm = ${eval.nullTerm}
           val $primitiveTerm =
             if($nullTerm)
@@ -119,7 +147,7 @@ abstract class CodeGenerator extends Logging {
         val resultCode = f(eval1.primitiveTerm, eval2.primitiveTerm)
 
         eval1.code ++ eval2.code ++
-          q"""
+        q"""
           val $nullTerm = ${eval1.nullTerm} || ${eval2.nullTerm}
           val $primitiveTerm: ${termForType(resultType)} =
             if($nullTerm) {
@@ -135,14 +163,15 @@ abstract class CodeGenerator extends Logging {
 
     // TODO: Skip generation of null handling code when expression are not nullable.
     val primitiveEvaluation: PartialFunction[Expression, Seq[Tree]] = {
-      case b @ BoundReference(ordinal, _) =>
+      case b @ BoundReference(ordinal, dataType, nullable) =>
+        val nullValue = if (nullable) q"$inputTuple.isNullAt($ordinal)" else  q"false"
         q"""
-          val $nullTerm: Boolean = $inputTuple.isNullAt($ordinal)
-          val $primitiveTerm: ${termForType(b.dataType)} =
+          val $nullTerm: Boolean = $nullValue
+          val $primitiveTerm: ${termForType(dataType)} =
             if($nullTerm)
-              ${defaultPrimitive(e.dataType)}
+              ${defaultPrimitive(dataType)}
             else
-              ${getColumn(inputTuple, b.dataType, ordinal)}
+              ${getColumn(inputTuple, dataType, ordinal)}
          """.children
 
       case expressions.Literal(null, dataType) =>
@@ -162,11 +191,13 @@ abstract class CodeGenerator extends Logging {
           val $nullTerm = ${value == null}
           val $primitiveTerm: ${termForType(dataType)} = $value
          """.children
+
       case expressions.Literal(value: Int, dataType) =>
         q"""
           val $nullTerm = ${value == null}
           val $primitiveTerm: ${termForType(dataType)} = $value
          """.children
+
       case expressions.Literal(value: Long, dataType) =>
         q"""
           val $nullTerm = ${value == null}
@@ -176,7 +207,7 @@ abstract class CodeGenerator extends Logging {
       case Cast(e @ BinaryType(), StringType) =>
         val eval = expressionEvaluator(e)
         eval.code ++
-          q"""
+        q"""
           val $nullTerm = ${eval.nullTerm}
           val $primitiveTerm =
             if($nullTerm)
@@ -200,7 +231,7 @@ abstract class CodeGenerator extends Logging {
       case Cast(e, StringType) =>
         val eval = expressionEvaluator(e)
         eval.code ++
-          q"""
+        q"""
           val $nullTerm = ${eval.nullTerm}
           val $primitiveTerm =
             if($nullTerm)
@@ -251,7 +282,7 @@ abstract class CodeGenerator extends Logging {
         val eval2 = expressionEvaluator(e2)
 
         eval1.code ++ eval2.code ++
-          q"""
+        q"""
           var $nullTerm = false
           var $primitiveTerm: ${termForType(BooleanType)} = false
 
@@ -272,7 +303,7 @@ abstract class CodeGenerator extends Logging {
         val eval2 = expressionEvaluator(e2)
 
         eval1.code ++ eval2.code ++
-          q"""
+        q"""
           var $nullTerm = false
           var $primitiveTerm: ${termForType(BooleanType)} = false
 
@@ -360,10 +391,10 @@ abstract class CodeGenerator extends Logging {
         log.debug(s"No rules to generate $e")
         val tree = reify { e }
         q"""
-            val $objectTerm = $tree.eval(i)
-            val $nullTerm = $objectTerm == null
-            val $primitiveTerm = $objectTerm.asInstanceOf[${termForType(e.dataType)}]
-          """.children
+          val $objectTerm = $tree.eval(i)
+          val $nullTerm = $objectTerm == null
+          val $primitiveTerm = $objectTerm.asInstanceOf[${termForType(e.dataType)}]
+         """.children
       }
 
     EvaluatedExpression(code, nullTerm, primitiveTerm, objectTerm)
@@ -377,10 +408,10 @@ abstract class CodeGenerator extends Logging {
   }
 
   protected def setColumn(
-                           destinationRow: TermName,
-                           dataType: DataType,
-                           ordinal: Int,
-                           value: TermName) = {
+      destinationRow: TermName,
+      dataType: DataType,
+      ordinal: Int,
+      value: TermName) = {
     dataType match {
       case dt @ NativeType() => q"$destinationRow.${mutatorForType(dt)}($ordinal, $value)"
       case _ => q"$destinationRow.update($ordinal, $value)"

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

@@ -23,31 +23,24 @@ import org.apache.spark.sql.catalyst.expressions._
  * Generates byte code that produces a [[MutableRow]] object that can update itself based on a new
  * input [[Row]] for a fixed set of [[Expression Expressions]].
  */
-object GenerateMutableProjection extends CodeGenerator {
+object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => MutableProjection] {
   import scala.reflect.runtime.{universe => ru}
   import scala.reflect.runtime.universe._
 
-  // TODO: Should be weak references... bounded in size.
-  val projectionCache = new collection.mutable.HashMap[Seq[Expression], () => MutableProjection]
-
-  def apply(expressions: Seq[Expression], inputSchema: Seq[Attribute]): (() => MutableProjection) =
-    apply(expressions.map(BindReferences.bindReference(_, inputSchema)))
+  val mutableRowName = newTermName("mutableRow")
 
-  // TODO: Safe to fire up multiple instances of the compiler?
-  def apply(expressions: Seq[Expression]): () => MutableProjection =
-    globalLock.synchronized {
-      val cleanedExpressions = expressions.map(ExpressionCanonicalizer(_))
-      projectionCache.getOrElseUpdate(cleanedExpressions, createProjection(cleanedExpressions))
-    }
+  protected def canonicalize(in: Seq[Expression]): Seq[Expression] =
+    in.map(ExpressionCanonicalizer(_))
 
-  val mutableRowName = newTermName("mutableRow")
+  protected def bind(in: Seq[Expression], inputSchema: Seq[Attribute]): Seq[Expression] =
+    in.map(BindReferences.bindReference(_, inputSchema))
 
-  def createProjection(expressions: Seq[Expression]): (() => MutableProjection) = {
+  protected def create(expressions: Seq[Expression]): (() => MutableProjection) = {
     val projectionCode = expressions.zipWithIndex.flatMap { case (e, i) =>
       val evaluationCode = expressionEvaluator(e)
 
       evaluationCode.code :+
-        q"""
+      q"""
         if(${evaluationCode.nullTerm})
           mutableRow.setNullAt($i)
         else
@@ -57,25 +50,25 @@ object GenerateMutableProjection extends CodeGenerator {
 
     val code =
       q"""
-      () => { new $mutableProjectionType {
+        () => { new $mutableProjectionType {
 
-        private[this] var $mutableRowName: $mutableRowType =
-          new $genericMutableRowType(${expressions.size})
+          private[this] var $mutableRowName: $mutableRowType =
+            new $genericMutableRowType(${expressions.size})
 
-        def target(row: $mutableRowType): $mutableProjectionType = {
-          $mutableRowName = row
-          this
-        }
+          def target(row: $mutableRowType): $mutableProjectionType = {
+            $mutableRowName = row
+            this
+          }
 
-        /* Provide immutable access to the last projected row. */
-        def currentValue: $rowType = mutableRow
+          /* Provide immutable access to the last projected row. */
+          def currentValue: $rowType = mutableRow
 
-        def apply(i: $rowType): $rowType = {
-          ..$projectionCode
-          mutableRow
-        }
-      } }
-    """
+          def apply(i: $rowType): $rowType = {
+            ..$projectionCode
+            mutableRow
+          }
+        } }
+      """
 
     log.debug(s"code for ${expressions.mkString(",")}:\n$code")
     toolBox.eval(code).asInstanceOf[() => MutableProjection]