[SPARK-31115][SQL] Provide config to avoid using switch statement in generated code to avoid Janino bug

Author: HeartSaVioR

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

@@ -454,7 +454,9 @@ case class InSet(child: Expression, hset: Set[Any]) extends UnaryExpression with
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    if (canBeComputedUsingSwitch && hset.size <= SQLConf.get.optimizerInSetSwitchThreshold) {
+    val sqlConf = SQLConf.get
+    if (canBeComputedUsingSwitch && hset.size <= sqlConf.optimizerInSetSwitchThreshold &&
+      sqlConf.codegenUseSwitchStatement) {
       genCodeWithSwitch(ctx, ev)
     } else {
       genCodeWithSet(ctx, ev)

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1130,6 +1130,23 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  val CODEGEN_USE_SWITCH_STATEMENT =
+    buildConf("spark.sql.codegen.useSwitchStatement")
+      .internal()
+      .doc("When true, Spark leverages switch statement while generating code. Otherwise Spark " +
+        "will leverage if ~ else if ~ else statement as an alternative. In normal case, " +
+        "'switch' statement is preferred against if ~ else if ~ else. This configuration is " +
+        "required to avoid Janino bug (https://github.com/janino-compiler/janino/issues/113); " +
+        "If InternalCompilerException has been thrown and following conditions are met, you " +
+        "may want to turn this off and try executing the query again." +
+        "1) The generated code contains 'switch' statement." +
+        "2) Exception message contains 'Operand stack inconsistent at offset xxx: Previous size 1" +
+        ", now 0'." +
+        "The configuration will be no-op and maybe removed once Spark upgrades Janino containing" +
+        " the fix.")
+      .booleanConf
+      .createWithDefault(true)
+
   val FILES_MAX_PARTITION_BYTES = buildConf("spark.sql.files.maxPartitionBytes")
     .doc("The maximum number of bytes to pack into a single partition when reading files. " +
       "This configuration is effective only when using file-based sources such as Parquet, JSON " +
@@ -2764,6 +2781,8 @@ class SQLConf extends Serializable with Logging {
   def wholeStageSplitConsumeFuncByOperator: Boolean =
     getConf(WHOLESTAGE_SPLIT_CONSUME_FUNC_BY_OPERATOR)
 
+  def codegenUseSwitchStatement: Boolean = getConf(CODEGEN_USE_SWITCH_STATEMENT)
+
   def tableRelationCacheSize: Int =
     getConf(StaticSQLConf.FILESOURCE_TABLE_RELATION_CACHE_SIZE)
 

sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala

@@ -54,6 +54,8 @@ case class ExpandExec(
   private[this] val projection =
     (exprs: Seq[Expression]) => UnsafeProjection.create(exprs, child.output)
 
+  private val useSwitchStatement: Boolean = sqlContext.conf.codegenUseSwitchStatement
+
   protected override def doExecute(): RDD[InternalRow] = attachTree(this, "execute") {
     val numOutputRows = longMetric("numOutputRows")
 
@@ -167,8 +169,9 @@ case class ExpandExec(
       }
     }
 
-    // Part 2: switch/case statements
-    val cases = projections.zipWithIndex.map { case (exprs, row) =>
+    // Part 2: switch/case statements, or if/else if statements via configuration
+
+    val updates = projections.map { exprs =>
       var updateCode = ""
       for (col <- exprs.indices) {
         if (!sameOutput(col)) {
@@ -178,27 +181,48 @@ case class ExpandExec(
                |${ev.code}
                |${outputColumns(col).isNull} = ${ev.isNull};
                |${outputColumns(col).value} = ${ev.value};
-            """.stripMargin
+             """.stripMargin
         }
       }
+      updateCode.trim
+    }
+
+    // the name needs to be known to build conditions
+    val i = ctx.freshName("i")
+    val loopContent = if (useSwitchStatement) {
+      val cases = updates.zipWithIndex.map { case (updateCode, row) =>
+        s"""
+           |case $row:
+           |  ${updateCode.trim}
+           |  break;
+         """.stripMargin
+      }
 
       s"""
-         |case $row:
-         |  ${updateCode.trim}
-         |  break;
+         |switch ($i) {
+         |  ${cases.mkString("\n").trim}
+         |}
        """.stripMargin
+    } else {
+      val conditions = updates.zipWithIndex.map { case (updateCode, row) =>
+        (if (row > 0) "else " else "") +
+          s"""
+             |if ($i == $row) {
+             |  ${updateCode.trim}
+             |}
+           """.stripMargin
+      }
+
+      conditions.mkString("\n").trim
     }
 
     val numOutput = metricTerm(ctx, "numOutputRows")
-    val i = ctx.freshName("i")
     // these column have to declared before the loop.
     val evaluate = evaluateVariables(outputColumns)
     s"""
        |$evaluate
        |for (int $i = 0; $i < ${projections.length}; $i ++) {
-       |  switch ($i) {
-       |    ${cases.mkString("\n").trim}
-       |  }
+       |  $loopContent
        |  $numOutput.add(1);
        |  ${consume(ctx, outputColumns)}
        |}

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

@@ -31,7 +31,6 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.test.SQLTestData.DecimalData
 import org.apache.spark.sql.types._
-import org.apache.spark.unsafe.types.CalendarInterval
 
 case class Fact(date: Int, hour: Int, minute: Int, room_name: String, temp: Double)
 
@@ -957,4 +956,61 @@ class DataFrameAggregateSuite extends QueryTest
       assert(error.message.contains("function count_if requires boolean type"))
     }
   }
+
+  /**
+   * NOTE: The test code tries to control the size of for/switch statement in expand_doConsume,
+   * as well as the overall size of expand_doConsume, so that the query triggers known Janino
+   * bug - https://github.com/janino-compiler/janino/issues/113.
+   *
+   * The expected exception message from Janino when we use switch statement for "ExpandExec":
+   * - "Operand stack inconsistent at offset xxx: Previous size 1, now 0"
+   * which will not happen when we use if-else-if statement for "ExpandExec".
+   *
+   * "The number of fields" and "The number of distinct aggregation functions" are the major
+   * factors to increase the size of generated code: while these values should be large enough
+   * to trigger the Janino bug, these values should not also too big; otherwise one of below
+   * exceptions might be thrown:
+   * - "expand_doConsume would be beyond 64KB"
+   * - "java.lang.ClassFormatError: Too many arguments in method signature in class file"
+   */
+  test("SPARK-31115 Lots of columns and distinct aggregations shouldn't break code generation") {
+    withSQLConf(
+      (SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true"),
+      (SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key, "10000"),
+      (SQLConf.CODEGEN_FALLBACK.key, "false"),
+      (SQLConf.CODEGEN_LOGGING_MAX_LINES.key, "-1"),
+      (SQLConf.CODEGEN_USE_SWITCH_STATEMENT.key, "false")
+    ) {
+      var df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c")
+
+      // The value is tested under commit "e807118eef9e0214170ff62c828524d237bd58e3":
+      // the query fails with switch statement, whereas it passes with if-else statement.
+      // Note that the value depends on the Spark logic as well - different Spark versions may
+      // require different value to ensure the test failing with switch statement.
+      val numNewFields = 100
+
+      df = df.withColumns(
+        (1 to numNewFields).map { idx => s"a$idx" },
+        (1 to numNewFields).map { idx =>
+          when(col("c").mod(lit(2)).===(lit(0)), lit(idx)).otherwise(col("c"))
+        }
+      )
+
+      val aggExprs: Array[Column] = Range(1, numNewFields).map { idx =>
+        if (idx % 2 == 0) {
+          coalesce(countDistinct(s"a$idx"), lit(0))
+        } else {
+          coalesce(count(s"a$idx"), lit(0))
+        }
+      }.toArray
+
+      val aggDf = df
+        .groupBy("a", "b")
+        .agg(aggExprs.head, aggExprs.tail: _*)
+
+      // We are only interested in whether the code compilation fails or not, so skipping
+      // verificaion on outputs.
+      aggDf.collect()
+    }
+  }
 }