Unify with UnsafeProjectionCreator.

Author: viirya

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

@@ -31,49 +31,37 @@ object CodegenError {
   }
 }
 
+trait CodegenObjectFactoryBase[IN, OUT] {
+  protected def createObject(in: IN): OUT
+  protected def createCodeGeneratedObject(in: IN): OUT
+  protected def createInterpretedObject(in: IN): OUT
+}
+
 /**
- * A factory class which can be used to create objects that have both codegen and interpreted
+ * A factory which can be used to create objects that have both codegen and interpreted
  * implementations. This tries to create codegen object first, if any compile error happens,
  * it fallbacks to interpreted version.
  */
-abstract class CodegenObjectFactory[IN, OUT] {
-
-  def createObject(in: IN): OUT = try {
+trait CodegenObjectFactory[IN, OUT] extends CodegenObjectFactoryBase[IN, OUT] {
+  override protected def createObject(in: IN): OUT = try {
     createCodeGeneratedObject(in)
   } catch {
     case CodegenError(_) => createInterpretedObject(in)
   }
-
-  protected def createCodeGeneratedObject(in: IN): OUT
-  protected def createInterpretedObject(in: IN): OUT
 }
 
-object UnsafeProjectionFactory extends CodegenObjectFactory[Seq[Expression], UnsafeProjection]
-    with UnsafeProjectionCreator {
-
-  override protected def createCodeGeneratedObject(in: Seq[Expression]): UnsafeProjection = {
-    UnsafeProjection.createProjection(in)
-  }
-
-  override protected def createInterpretedObject(in: Seq[Expression]): UnsafeProjection = {
-    InterpretedUnsafeProjection.createProjection(in)
-  }
-
-  override protected[sql] def createProjection(exprs: Seq[Expression]): UnsafeProjection =
-    createObject(exprs)
-
-  /**
-   * Same as other create()'s but allowing enabling/disabling subexpression elimination.
-   * The param `subexpressionEliminationEnabled` doesn't guarantee to work. For example,
-   * when fallbacking to interpreted execution, it is not supported.
-   */
-  def create(
-      exprs: Seq[Expression],
-      inputSchema: Seq[Attribute],
-      subexpressionEliminationEnabled: Boolean): UnsafeProjection = try {
-    UnsafeProjection.create(exprs, inputSchema, subexpressionEliminationEnabled)
-  } catch {
-    case CodegenError(_) => InterpretedUnsafeProjection.create(exprs, inputSchema)
-  }
+/**
+ * A factory which can be used to create codegen objects without fallback to interpreted version.
+ */
+trait CodegenObjectFactoryWithoutFallback[IN, OUT] extends CodegenObjectFactoryBase[IN, OUT] {
+  override protected def createObject(in: IN): OUT =
+    createCodeGeneratedObject(in)
 }
 
+/**
+ * A factory which can be used to create objects with interpreted implementation.
+ */
+trait InterpretedCodegenObjectFactory[IN, OUT] extends CodegenObjectFactoryBase[IN, OUT] {
+  override protected def createObject(in: IN): OUT =
+    createInterpretedObject(in)
+}

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

@@ -87,12 +87,11 @@ class InterpretedUnsafeProjection(expressions: Array[Expression]) extends Unsafe
 /**
  * Helper functions for creating an [[InterpretedUnsafeProjection]].
  */
-object InterpretedUnsafeProjection extends UnsafeProjectionCreator {
-
+object InterpretedUnsafeProjection {
   /**
    * Returns an [[UnsafeProjection]] for given sequence of bound Expressions.
    */
-  override protected[sql] def createProjection(exprs: Seq[Expression]): UnsafeProjection = {
+  protected[sql] def createProjection(exprs: Seq[Expression]): UnsafeProjection = {
     // We need to make sure that we do not reuse stateful expressions.
     val cleanedExpressions = exprs.map(_.transform {
       case s: Stateful => s.freshCopy()

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

@@ -108,7 +108,20 @@ abstract class UnsafeProjection extends Projection {
   override def apply(row: InternalRow): UnsafeRow
 }
 
-trait UnsafeProjectionCreator {
+/**
+ * The factory base class for `UnsafeProjection`.
+ */
+abstract class UnsafeProjectionCreator extends
+    CodegenObjectFactoryBase[Seq[Expression], UnsafeProjection] {
+
+  override protected def createCodeGeneratedObject(in: Seq[Expression]): UnsafeProjection = {
+    GenerateUnsafeProjection.generate(in)
+  }
+
+  override protected def createInterpretedObject(in: Seq[Expression]): UnsafeProjection = {
+    InterpretedUnsafeProjection.createProjection(in)
+  }
+
   protected def toBoundExprs(
       exprs: Seq[Expression],
       inputSchema: Seq[Attribute]): Seq[Expression] = {
@@ -157,18 +170,36 @@ trait UnsafeProjectionCreator {
   /**
    * Returns an [[UnsafeProjection]] for given sequence of bound Expressions.
    */
-  protected[sql] def createProjection(exprs: Seq[Expression]): UnsafeProjection
+  protected[sql] def createProjection(exprs: Seq[Expression]): UnsafeProjection =
+    createObject(exprs)
 }
 
-object UnsafeProjection extends UnsafeProjectionCreator {
-
-  override protected[sql] def createProjection(exprs: Seq[Expression]): UnsafeProjection = {
-    GenerateUnsafeProjection.generate(exprs)
+// A `UnsafeProjectionCreator` can fallback to interpreted version if codegen compile error happens.
+object UnsafeProjection extends UnsafeProjectionCreator
+    with CodegenObjectFactory[Seq[Expression], UnsafeProjection] {
+  /**
+   * Same as other create()'s but allowing enabling/disabling subexpression elimination.
+   * The param `subexpressionEliminationEnabled` doesn't guarantee to work. For example,
+   * when fallbacking to interpreted execution, it is not supported.
+   */
+  def create(
+      exprs: Seq[Expression],
+      inputSchema: Seq[Attribute],
+      subexpressionEliminationEnabled: Boolean): UnsafeProjection = {
+    val unsafeExprs = toUnsafeExprs(toBoundExprs(exprs, inputSchema))
+    try {
+      GenerateUnsafeProjection.generate(unsafeExprs, subexpressionEliminationEnabled)
+    } catch {
+      case CodegenError(_) => InterpretedUnsafeProjection.createProjection(unsafeExprs)
+    }
   }
+}
 
+// Codegen version `UnsafeProjectionCreator`.
+object CodegenUnsafeProjectionCreator extends UnsafeProjectionCreator
+    with CodegenObjectFactoryWithoutFallback[Seq[Expression], UnsafeProjection] {
   /**
    * Same as other create()'s but allowing enabling/disabling subexpression elimination.
-   * TODO: refactor the plumbing and clean this up.
    */
   def create(
       exprs: Seq[Expression],
@@ -179,6 +210,10 @@ object UnsafeProjection extends UnsafeProjectionCreator {
   }
 }
 
+// Interpreted version `UnsafeProjectionCreator`.
+object InterpretedUnsafeProjectionCreator extends UnsafeProjectionCreator
+    with InterpretedCodegenObjectFactory[Seq[Expression], UnsafeProjection]
+
 /**
  * A projection that could turn UnsafeRow into GenericInternalRow
  */

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala

@@ -196,8 +196,10 @@ trait ExpressionEvalHelper extends GeneratorDrivenPropertyChecks {
       expression: Expression,
       expected: Any,
       inputRow: InternalRow = EmptyRow): Unit = {
-    checkEvaluationWithUnsafeProjection(expression, expected, inputRow, UnsafeProjection)
-    checkEvaluationWithUnsafeProjection(expression, expected, inputRow, InterpretedUnsafeProjection)
+    checkEvaluationWithUnsafeProjection(expression, expected, inputRow,
+      CodegenUnsafeProjectionCreator)
+    checkEvaluationWithUnsafeProjection(expression, expected, inputRow,
+      InterpretedUnsafeProjectionCreator)
   }
 
   protected def checkEvaluationWithUnsafeProjection(
@@ -228,7 +230,7 @@ trait ExpressionEvalHelper extends GeneratorDrivenPropertyChecks {
   protected def evaluateWithUnsafeProjection(
       expression: Expression,
       inputRow: InternalRow = EmptyRow,
-      factory: UnsafeProjectionCreator = UnsafeProjection): InternalRow = {
+      factory: UnsafeProjectionCreator = CodegenUnsafeProjectionCreator): InternalRow = {
     // SPARK-16489 Explicitly doing code generation twice so code gen will fail if
     // some expression is reusing variable names across different instances.
     // This behavior is tested in ExpressionEvalHelperSuite.

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/UnsafeRowConverterSuite.scala

@@ -37,8 +37,8 @@ class UnsafeRowConverterSuite extends SparkFunSuite with Matchers {
     name: String)(
     f: UnsafeProjectionCreator => Unit): Unit = {
     test(name) {
-      f(UnsafeProjection)
-      f(InterpretedUnsafeProjection)
+      f(CodegenUnsafeProjectionCreator)
+      f(InterpretedUnsafeProjectionCreator)
     }
   }