fallback in case-when

Author: Davies Liu

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

@@ -86,7 +86,7 @@ case class If(predicate: Expression, trueValue: Expression, falseValue: Expressi
  * @param elseValue optional value for the else branch
  */
 case class CaseWhen(branches: Seq[(Expression, Expression)], elseValue: Option[Expression] = None)
-  extends Expression {
+  extends Expression with CodegenFallback {
 
   override def children: Seq[Expression] = branches.flatMap(b => b._1 :: b._2 :: Nil) ++ elseValue
 
@@ -136,7 +136,16 @@ case class CaseWhen(branches: Seq[(Expression, Expression)], elseValue: Option[E
     }
   }
 
+  def shouldCodegen: Boolean = {
+    branches.length < CaseWhen.MAX_NUMBER_OF_SWITCHES
+  }
+
   override def genCode(ctx: CodegenContext, ev: ExprCode): String = {
+    if (!shouldCodegen) {
+      // Fallback to interpreted mode if there are too many branches, or it may reach the
+      // 64K limit (number of bytecode for single Java method).
+      return super.genCode(ctx, ev)
+    }
     // Generate code that looks like:
     //
     // condA = ...
@@ -205,6 +214,9 @@ case class CaseWhen(branches: Seq[(Expression, Expression)], elseValue: Option[E
 /** Factory methods for CaseWhen. */
 object CaseWhen {
 
+  // The maxium number of switches supported with codegen.
+  val MAX_NUMBER_OF_SWITCHES = 20
+
   def apply(branches: Seq[(Expression, Expression)], elseValue: Expression): CaseWhen = {
     CaseWhen(branches, Option(elseValue))
   }

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

@@ -58,6 +58,27 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("SPARK-13242: complicated case-when expressions") {
+    val cases = 50
+    val clauses = 20
+
+    // Generate an individual case
+    def generateCase(n: Int): (Expression, Expression) = {
+      val condition = (1 to clauses)
+        .map(c => EqualTo(BoundReference(0, StringType, false), Literal(s"$c:$n")))
+        .reduceLeft[Expression]((l, r) => Or(l, r))
+      (condition, Literal(n))
+    }
+
+    val expression = CaseWhen((1 to cases).map(generateCase(_)))
+
+    val plan = GenerateMutableProjection.generate(Seq(expression))()
+    val input = new GenericMutableRow(Array[Any](UTF8String.fromString(s"${clauses}:${cases}")))
+    val actual = plan(input).toSeq(Seq(expression.dataType))
+
+    assert(actual(0) == cases)
+  }
+
   test("test generated safe and unsafe projection") {
     val schema = new StructType(Array(
       StructField("a", StringType, true),

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

@@ -416,6 +416,7 @@ private[sql] case class CollapseCodegenStages(sqlContext: SQLContext) extends Ru
 
   private def supportCodegen(e: Expression): Boolean = e match {
     case e: LeafExpression => true
+    case e: CaseWhen => e.shouldCodegen
     // CodegenFallback requires the input to be an InternalRow
     case e: CodegenFallback => false
     case _ => true