address review comment

Author: kiszk

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

@@ -279,33 +279,38 @@ case class CaseWhenCodegen(
       """
     }
 
-    var isGlobalVariable = false
-    var generatedCode = ""
     var numIfthen = 0
-    cases.foreach { ifthen =>
-      generatedCode += ifthen + "\nelse {\n"
-      numIfthen += 1
-
-      if (generatedCode.length > 1024 &&
-        // Split these expressions only if they are created from a row object
-        (ctx.INPUT_ROW != null && ctx.currentVars == null)) {
-        val flag = "flag"
-        generatedCode += s" $flag = false;\n" + "}\n" * numIfthen
-        val funcName = ctx.freshName("caseWhenNestedIf")
-        val funcBody =
-          s"""
-           |private boolean $funcName(InternalRow ${ctx.INPUT_ROW}) {
-           |  boolean $flag = true;
-           |  $generatedCode
-           |  return $flag;
-           |}
-         """.stripMargin
-        val fullFuncName = ctx.addNewFunction(funcName, funcBody)
-        isGlobalVariable = true
-
-        generatedCode = s"if ($fullFuncName(${ctx.INPUT_ROW})) {\n// do nothing\n} else {\n"
-        numIfthen = 1
+    var isGlobalVariable = false
+    var generatedCode = if (cases.map(s => s.length).sum <= 1024) {
+      cases.mkString("\nelse {\n")
+    } else {
+      var code = ""
+      cases.foreach { ifthen =>
+        code += ifthen + "\nelse {\n"
+        numIfthen += 1
+
+        if (code.length > 1024 &&
+          // Split these expressions only if they are created from a row object
+          (ctx.INPUT_ROW != null && ctx.currentVars == null)) {
+          val flag = "flag"
+          code += s" $flag = false;\n" + "}\n" * numIfthen
+          val funcName = ctx.freshName("caseWhenNestedIf")
+          val funcBody =
+            s"""
+             |private boolean $funcName(InternalRow ${ctx.INPUT_ROW}) {
+             |  boolean $flag = true;
+             |  $code
+             |  return $flag;
+             |}
+           """.stripMargin
+          val fullFuncName = ctx.addNewFunction(funcName, funcBody)
+          isGlobalVariable = true
+
+          code = s"if ($fullFuncName(${ctx.INPUT_ROW})) {\n// do nothing\n} else {\n"
+          numIfthen = 1
+        }
       }
+      code
     }
 
     elseValue.foreach { elseExpr =>

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

@@ -2159,20 +2159,13 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
     assert(mean.collect().toSet === Set(Row("0.0345678900000000000000000000000000000")))
   }
 
-  testQuietly("SPARK-21413: Multiple projections with CASE WHEN fails") {
+  // ignore end-to-end test since sbt test does not go to fallback path in whole-stage codegen
+  ignore("SPARK-21413: Multiple projections with CASE WHEN fails") {
     val schema = StructType(StructField("a", IntegerType) :: Nil)
-    val df = spark.createDataFrame(sparkContext.parallelize(Seq(Row(1))), schema)
-    val df1 =
-      df.withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-        .withColumn("a", when($"a" === 0, null).otherwise($"a"))
-    checkAnswer(df1, Row(1))
+    var df = spark.createDataFrame(sparkContext.parallelize(Seq(Row(1))), schema)
+    for (i <- 1 to 10) {
+      df = df.withColumn("a", when($"a" === 0, null).otherwise($"a"))
+    }
+    checkAnswer(df, Row(1))
   }
 }