Revert "[SPARK-22600][SQL] Fix 64kb limit for deeply nested expressions under wholestage codegen"

This reverts commit c7d0148.

Author: cloud-fan

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

@@ -105,12 +105,6 @@ abstract class Expression extends TreeNode[Expression] {
       val isNull = ctx.freshName("isNull")
       val value = ctx.freshName("value")
       val eval = doGenCode(ctx, ExprCode("", isNull, value))
-      eval.isNull = if (this.nullable) eval.isNull else "false"
-
-      // Records current input row and variables of this expression.
-      eval.inputRow = ctx.INPUT_ROW
-      eval.inputVars = findInputVars(ctx, eval)
-
       reduceCodeSize(ctx, eval)
       if (eval.code.nonEmpty) {
         // Add `this` in the comment.
@@ -121,29 +115,9 @@ abstract class Expression extends TreeNode[Expression] {
     }
   }
 
-  /**
-   * Returns the input variables to this expression.
-   */
-  private def findInputVars(ctx: CodegenContext, eval: ExprCode): Seq[ExprInputVar] = {
-    if (ctx.currentVars != null) {
-      this.collect {
-        case b @ BoundReference(ordinal, _, _) if ctx.currentVars(ordinal) != null =>
-          ExprInputVar(exprCode = ctx.currentVars(ordinal),
-            dataType = b.dataType, nullable = b.nullable)
-      }
-    } else {
-      Seq.empty
-    }
-  }
-
-  /**
-   * In order to prevent 64kb compile error, reducing the size of generated codes by
-   * separating it into a function if the size exceeds a threshold.
-   */
   private def reduceCodeSize(ctx: CodegenContext, eval: ExprCode): Unit = {
-    lazy val funcParams = ExpressionCodegen.getExpressionInputParams(ctx, this)
-
-    if (eval.code.trim.length > 1024 && funcParams.isDefined) {
+    // TODO: support whole stage codegen too
+    if (eval.code.trim.length > 1024 && ctx.INPUT_ROW != null && ctx.currentVars == null) {
       val setIsNull = if (eval.isNull != "false" && eval.isNull != "true") {
         val globalIsNull = ctx.freshName("globalIsNull")
         ctx.addMutableState(ctx.JAVA_BOOLEAN, globalIsNull)
@@ -158,20 +132,17 @@ abstract class Expression extends TreeNode[Expression] {
       val newValue = ctx.freshName("value")
 
       val funcName = ctx.freshName(nodeName)
-      val callParams = funcParams.map(_._1.mkString(", ")).get
-      val declParams = funcParams.map(_._2.mkString(", ")).get
-
       val funcFullName = ctx.addNewFunction(funcName,
         s"""
-           |private $javaType $funcName($declParams) {
+           |private $javaType $funcName(InternalRow ${ctx.INPUT_ROW}) {
            |  ${eval.code.trim}
            |  $setIsNull
            |  return ${eval.value};
            |}
            """.stripMargin)
 
       eval.value = newValue
-      eval.code = s"$javaType $newValue = $funcFullName($callParams);"
+      eval.code = s"$javaType $newValue = $funcFullName(${ctx.INPUT_ROW});"
     }
   }
 

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

@@ -55,24 +55,8 @@ import org.apache.spark.util.{ParentClassLoader, Utils}
  *                 to null.
  * @param value A term for a (possibly primitive) value of the result of the evaluation. Not
  *              valid if `isNull` is set to `true`.
- * @param inputRow A term that holds the input row name when generating this code.
- * @param inputVars A list of [[ExprInputVar]] that holds input variables when generating this code.
  */
-case class ExprCode(
-    var code: String,
-    var isNull: String,
-    var value: String,
-    var inputRow: String = null,
-    var inputVars: Seq[ExprInputVar] = Seq.empty)
-
-/**
- * Represents an input variable [[ExprCode]] to an evaluation of an [[Expression]].
- *
- * @param exprCode The [[ExprCode]] that represents the evaluation result for the input variable.
- * @param dataType The data type of the input variable.
- * @param nullable Whether the input variable can be null or not.
- */
-case class ExprInputVar(exprCode: ExprCode, dataType: DataType, nullable: Boolean)
+case class ExprCode(var code: String, var isNull: String, var value: String)
 
 /**
  * State used for subexpression elimination.
@@ -1028,25 +1012,16 @@ class CodegenContext {
     commonExprs.foreach { e =>
       val expr = e.head
       val fnName = freshName("evalExpr")
-      val isNull = if (expr.nullable) {
-        s"${fnName}IsNull"
-      } else {
-        ""
-      }
+      val isNull = s"${fnName}IsNull"
       val value = s"${fnName}Value"
 
       // Generate the code for this expression tree and wrap it in a function.
       val eval = expr.genCode(this)
-      val assignIsNull = if (expr.nullable) {
-        s"$isNull = ${eval.isNull};"
-      } else {
-        ""
-      }
       val fn =
         s"""
            |private void $fnName(InternalRow $INPUT_ROW) {
            |  ${eval.code.trim}
-           |  $assignIsNull
+           |  $isNull = ${eval.isNull};
            |  $value = ${eval.value};
            |}
            """.stripMargin
@@ -1064,17 +1039,12 @@ class CodegenContext {
       //   2. Less code.
       // Currently, we will do this for all non-leaf only expression trees (i.e. expr trees with
       // at least two nodes) as the cost of doing it is expected to be low.
-      if (expr.nullable) {
-        addMutableState(JAVA_BOOLEAN, isNull)
-      }
-      addMutableState(javaType(expr.dataType), value)
+      addMutableState(JAVA_BOOLEAN, isNull, s"$isNull = false;")
+      addMutableState(javaType(expr.dataType), value,
+        s"$value = ${defaultValue(expr.dataType)};")
 
       subexprFunctions += s"${addNewFunction(fnName, fn)}($INPUT_ROW);"
-      val state = if (expr.nullable) {
-        SubExprEliminationState(isNull, value)
-      } else {
-        SubExprEliminationState("false", value)
-      }
+      val state = SubExprEliminationState(isNull, value)
       e.foreach(subExprEliminationExprs.put(_, state))
     }
   }