Evaluate common subexpression like lazy variable with a function approach.

Author: viirya

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

@@ -67,7 +67,8 @@ case class SubExprEliminationState(isNull: String, value: String)
 /**
  * Codes and common subexpressions mapping used for subexpression elimination.
  *
- * @param codes Strings representing the codes that evaluate common subexpressions.
+ * @param codes Strings representing the codes that reset the initialization status of
+ *              common subexpression evaluation.
  * @param states Foreach expression that is participating in subexpression elimination,
  *               the state to use.
  */
@@ -680,6 +681,47 @@ class CodegenContext {
     genCodes
   }
 
+  /**
+   * A private helper function used to construct the parameter list for subexpression elimination
+   * evaluation functions
+   *
+   * @param expression The subexpression to evaluate.
+   * @param caller Indicating to construct parameter list for function caller.
+   */
+  private def genFunctionParamsListForSubExprEliminate(
+      expression: Expression,
+      caller: Boolean): String = {
+    val boundRefs = expression.collect {
+      case b: BoundReference => b
+    }.distinct
+    if (currentVars == null) {
+      if (caller) INPUT_ROW else s"InternalRow $INPUT_ROW"
+    } else {
+      val boundRefsInCurrentVars = boundRefs.filter(b => currentVars(b.ordinal) != null)
+      val currentVarsParams = boundRefsInCurrentVars.map { bound =>
+        val paramType = javaType(bound.dataType)
+        val variable = currentVars(bound.ordinal).value
+        val isNull = currentVars(bound.ordinal).isNull
+        if (caller) {
+          if (isNull == "false") variable else s"$variable, $isNull"
+        } else {
+          if (isNull == "false") {
+            s"$paramType $variable"
+          } else {
+            s"$paramType $variable, boolean $isNull"
+          }
+        }
+      }
+
+      if (boundRefsInCurrentVars.size == boundRefs.size) {
+        currentVarsParams.mkString(", ")
+      } else {
+        val rowParam = if (caller) INPUT_ROW else s"InternalRow $INPUT_ROW"
+        (Seq(rowParam) ++ currentVarsParams).mkString(", ")
+      }
+    }
+  }
+
   /**
    * Checks and sets up the state and codegen for subexpression elimination. This finds the
    * common subexpressions, generates the code snippets that evaluate those expressions and
@@ -700,11 +742,63 @@ class CodegenContext {
     val commonExprs = equivalentExpressions.getAllEquivalentExprs.filter(_.size > 1)
     val codes = commonExprs.map { e =>
       val expr = e.head
+      val fnName = freshName("evalSubExpr")
+      val isNull = s"${fnName}IsNull"
+      val value = s"${fnName}Value"
+      val isInitialized = s"${fnName}IsInitialized"
+
+      val functionParams = genFunctionParamsListForSubExprEliminate(expr, false)
+      val callerParams = genFunctionParamsListForSubExprEliminate(expr, true)
+
       // Generate the code for this expression tree.
       val code = expr.genCode(this)
-      val state = SubExprEliminationState(code.isNull, code.value)
+      val returnType = javaType(expr.dataType)
+      val fn =
+        s"""
+           |private void $fnName($functionParams) {
+           |  ${code.code.trim}
+           |  $isNull = ${code.isNull};
+           |  $value = ${code.value};
+           |  $isInitialized = true;
+           |}
+           """.stripMargin
+
+      val valueFnName = s"${fnName}ForValue"
+      val valueFn =
+        s"""
+           |private $returnType $valueFnName($functionParams) {
+           |  if (!$isInitialized) {
+           |    $fnName($callerParams);
+           |  }
+           |  return $value;
+           |}
+           """.stripMargin
+
+      val isNullFnName = s"${fnName}ForIsNull"
+      val isNullFn =
+        s"""
+           |private boolean $isNullFnName($functionParams) {
+           |  if (!$isInitialized) {
+           |    $fnName($callerParams);
+           |  }
+           |  return $isNull;
+           |}
+           """.stripMargin
+
+      addNewFunction(fnName, fn)
+      addNewFunction(valueFnName, valueFn)
+      addNewFunction(isNullFnName, isNullFn)
+
+      addMutableState("boolean", isNull, s"$isNull = false;")
+      addMutableState("boolean", isInitialized, s"$isInitialized = false;")
+      addMutableState(returnType, value, s"$value = ${defaultValue(expr.dataType)};")
+
+      val state = SubExprEliminationState(
+        isNull = s"$isNullFnName($callerParams)",
+        value = s"$valueFnName($callerParams)")
+
       e.foreach(subExprEliminationExprs.put(_, state))
-      code.code.trim
+      s"$isInitialized = false;"
     }
     SubExprCodes(codes, subExprEliminationExprs.toMap)
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala

@@ -253,7 +253,7 @@ case class HashAggregateExec(
     ctx.currentVars = bufVars ++ input
     val boundUpdateExpr = updateExpr.map(BindReferences.bindReference(_, inputAttrs))
     val subExprs = ctx.subexpressionEliminationForWholeStageCodegen(boundUpdateExpr)
-    val effectiveCodes = subExprs.codes.mkString("\n")
+    val resetSubExprEvaluation = subExprs.codes.mkString("\n")
     val aggVals = ctx.withSubExprEliminationExprs(subExprs.states) {
       boundUpdateExpr.map(_.genCode(ctx))
     }
@@ -266,8 +266,8 @@ case class HashAggregateExec(
     }
     s"""
        | // do aggregate
-       | // common sub-expressions
-       | $effectiveCodes
+       | // reset the initialization status for common sub-expressions
+       | $resetSubExprEvaluation
        | // evaluate aggregate function
        | ${evaluateVariables(aggVals)}
        | // update aggregation buffer
@@ -758,7 +758,7 @@ case class HashAggregateExec(
       ctx.INPUT_ROW = fastRowBuffer
       val boundUpdateExpr = updateExpr.map(BindReferences.bindReference(_, inputAttr))
       val subExprs = ctx.subexpressionEliminationForWholeStageCodegen(boundUpdateExpr)
-      val effectiveCodes = subExprs.codes.mkString("\n")
+      val resetSubExprEvaluation = subExprs.codes.mkString("\n")
       val fastRowEvals = ctx.withSubExprEliminationExprs(subExprs.states) {
         boundUpdateExpr.map(_.genCode(ctx))
       }
@@ -768,8 +768,8 @@ case class HashAggregateExec(
       }
       Option(
         s"""
-           |// common sub-expressions
-           |$effectiveCodes
+           |// reset the initialization status for common sub-expressions
+           |$resetSubExprEvaluation
            |// evaluate aggregate function
            |${evaluateVariables(fastRowEvals)}
            |// update fast row
@@ -814,7 +814,7 @@ case class HashAggregateExec(
       ctx.INPUT_ROW = unsafeRowBuffer
       val boundUpdateExpr = updateExpr.map(BindReferences.bindReference(_, inputAttr))
       val subExprs = ctx.subexpressionEliminationForWholeStageCodegen(boundUpdateExpr)
-      val effectiveCodes = subExprs.codes.mkString("\n")
+      val resetSubExprEvaluation = subExprs.codes.mkString("\n")
       val unsafeRowBufferEvals = ctx.withSubExprEliminationExprs(subExprs.states) {
         boundUpdateExpr.map(_.genCode(ctx))
       }
@@ -823,8 +823,8 @@ case class HashAggregateExec(
         ctx.updateColumn(unsafeRowBuffer, dt, i, ev, updateExpr(i).nullable)
       }
       s"""
-         |// common sub-expressions
-         |$effectiveCodes
+         |// reset the initialization status for common sub-expressions
+         |$resetSubExprEvaluation
          |// evaluate aggregate function
          |${evaluateVariables(unsafeRowBufferEvals)}
          |// update unsafe row buffer