[SPARK-21717][SQL] Decouple consume functions of physical operators in whole-stage codegen

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

@@ -89,6 +89,8 @@ case class ExpandExec(
     child.asInstanceOf[CodegenSupport].inputRDDs()
   }
 
+  override protected def doConsumeInChainOfFunc: Boolean = false
+
   protected override def doProduce(ctx: CodegenContext): String = {
     child.asInstanceOf[CodegenSupport].produce(ctx, this)
   }

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

@@ -177,6 +177,8 @@ case class SortExec(
      """.stripMargin.trim
   }
 
+  override protected def doConsumeInChainOfFunc: Boolean = false
+
   protected override val shouldStopRequired = false
 
   override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = {

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

@@ -149,13 +149,143 @@ trait CodegenSupport extends SparkPlan {
 
     ctx.freshNamePrefix = parent.variablePrefix
     val evaluated = evaluateRequiredVariables(output, inputVars, parent.usedInputs)
+
+    // Under certain conditions, we can put the logic to consume the rows of this operator into
+    // another function. So we can prevent a generated function too long to be optimized by JIT.
+    // The conditions:
+    // 1. The parent uses all variables in output. we can't defer variable evaluation when consume
+    //    in another function.
+    // 2. The output variables are not empty. If it's empty, we don't bother to do that.
+    // 3. We don't use row variable. The construction of row uses deferred variable evaluation. We
+    //    can't do it.
+    val requireAllOutput = output.forall(parent.usedInputs.contains(_))
+    val consumeFunc =
+      if (row == null && outputVars.nonEmpty && requireAllOutput) {
+        constructDoConsumeFunction(ctx, inputVars)
+      } else {
+        parent.doConsume(ctx, inputVars, rowVar)
+      }
     s"""
        |${ctx.registerComment(s"CONSUME: ${parent.simpleString}")}
        |$evaluated
-       |${parent.doConsume(ctx, inputVars, rowVar)}
+       |$consumeFunc
+     """.stripMargin
+  }
+
+  /**
+   * To prevent concatenated function growing too long to be optimized by JIT. Instead of inlining,
+   * we may put the consume logic of parent operator into a function and set this flag to `true`.
+   * The parent operator can know if its consume logic is inlined or in separated function.
+   */
+  private var doConsumeInFunc: Boolean = false
+
+  /**
+   * Returning true means we have at least one consume logic from child operator or this operator is
+   * separated in a function. If this is `true`, this operator shouldn't use `continue` statement to
+   * continue on next row, because its generated codes aren't enclosed in main while-loop.
+   *
+   * For example, we have generated codes for a query plan like:
+   *   Op1Exec
+   *     Op2Exec
+   *       Op3Exec
+   *
+   * If we put the consume code of Op2Exec into a separated function, the generated codes are like:
+   *   while (...) {
+   *     ... // logic of Op3Exec.
+   *     Op2Exec_doConsume(...);
+   *   }
+   *   private boolean Op2Exec_doConsume(...) {
+   *     ... // logic of Op2Exec to consume rows.
+   *   }
+   * For now, `doConsumeInChainOfFunc` of Op2Exec will be `true`.
+   *
+   * Notice for some operators like `HashAggregateExec`, it doesn't chain previous consume functions
+   * but begins with its produce framework. We should override `doConsumeInChainOfFunc` to return
+   * `false`.
+   */
+  protected def doConsumeInChainOfFunc: Boolean = {
+    val codegenChildren = children.map(_.asInstanceOf[CodegenSupport])
+    doConsumeInFunc || codegenChildren.exists(_.doConsumeInChainOfFunc)
+  }
+
+  /**
+   * The actual java statement this operator should use if there is a need to continue on next row
+   * in its `doConsume` codes.
+   *
+   *   while (...) {
+   *     ...       // logic of Op3Exec.
+   *     Op2Exec_doConsume(...);
+   *   }
+   *   private boolean Op2Exec_doConsume(...) {
+   *     ...       // logic of Op2Exec to consume rows.
+   *     continue; // Wrong. We can't use continue with the while-loop.
+   *   }
+   * In above code, we can't use `continue` in `Op2Exec_doConsume`.
+   *
+   * Instead, we do something like:
+   *   while (...) {
+   *     ...          // logic of Op3Exec.
+   *     boolean continueForLoop = Op2Exec_doConsume(...);
+   *     if (continueForLoop) continue;
+   *   }
+   *   private boolean Op2Exec_doConsume(...) {
+   *     ...          // logic of Op2Exec to consume rows.
+   *     return true; // When we need to do continue, we return true.
+   *   }
+   */
+  protected def continueStatementInDoConsume: String = if (doConsumeInChainOfFunc) {
+    "return true;";
+  } else {
+    "continue;"
+  }
+
+  /**
+   * To prevent concatenated function growing too long to be optimized by JIT. We can separate the
+   * parent's `doConsume` codes of a `CodegenSupport` operator into a function to call.
+   */
+  protected def constructDoConsumeFunction(
+      ctx: CodegenContext,
+      inputVars: Seq[ExprCode]): String = {
+    val (callingParams, arguList, inputVarsInFunc) =
+      constructConsumeParameters(ctx, output, inputVars)
+    parent.doConsumeInFunc = true
+    val rowVar = ExprCode("", "false", "unsafeRow")
+    val doConsume = ctx.freshName("doConsume")
+    val doConsumeFuncName = ctx.addNewFunction(doConsume,
+      s"""
+         | private boolean $doConsume($arguList) throws java.io.IOException {
+         |   ${parent.doConsume(ctx, inputVarsInFunc, rowVar)}
+         |   return false;
+         | }
+       """.stripMargin)
+
+    s"""
+       | boolean continueForLoop = $doConsumeFuncName($callingParams);
+       | if (continueForLoop) $continueStatementInDoConsume
      """.stripMargin
   }
 
+  /**
+   * Returns source code for calling consume function and the argument list of the consume function
+   * and also the `ExprCode` for the argument list.
+   */
+  protected def constructConsumeParameters(
+      ctx: CodegenContext,
+      attributes: Seq[Attribute],
+      variables: Seq[ExprCode]): (String, String, Seq[ExprCode]) = {
+    val params = variables.zipWithIndex.map { case (ev, i) =>
+      val callingParam = ev.value + ", " + ev.isNull
+      val arguName = ctx.freshName(s"expr_$i")
+      val arguIsNull = ctx.freshName(s"exprIsNull_$i")
+      (callingParam,
+        ctx.javaType(attributes(i).dataType) + " " + arguName + ", boolean " + arguIsNull,
+        ExprCode("", arguIsNull, arguName))
+    }.unzip3
+    (params._1.mkString(", "),
+      params._2.mkString(", "),
+      params._3)
+  }
+
   /**
    * Returns source code to evaluate all the variables, and clear the code of them, to prevent
    * them to be evaluated twice.
@@ -252,6 +382,8 @@ case class InputAdapter(child: SparkPlan) extends UnaryExecNode with CodegenSupp
     child.execute() :: Nil
   }
 
+  override protected def doConsumeInChainOfFunc: Boolean = false
+
   override def doProduce(ctx: CodegenContext): String = {
     val input = ctx.freshName("input")
     // Right now, InputAdapter is only used when there is one input RDD.

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

@@ -163,6 +163,8 @@ case class HashAggregateExec(
   // The variables used as aggregation buffer. Only used for aggregation without keys.
   private var bufVars: Seq[ExprCode] = _
 
+  override protected def doConsumeInChainOfFunc: Boolean = false
+
   private def doProduceWithoutKeys(ctx: CodegenContext): String = {
     val initAgg = ctx.freshName("initAgg")
     ctx.addMutableState("boolean", initAgg, s"$initAgg = false;")

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

@@ -149,7 +149,7 @@ case class FilterExec(condition: Expression, child: SparkPlan)
       s"""
          |$evaluated
          |${ev.code}
-         |if (${nullCheck}!${ev.value}) continue;
+         |if (${nullCheck}!${ev.value}) $continueStatementInDoConsume
        """.stripMargin
     }
 
@@ -317,7 +317,7 @@ case class SampleExec(
          """.stripMargin.trim)
 
       s"""
-         | if ($sampler.sample() == 0) continue;
+         | if ($sampler.sample() == 0) $continueStatementInDoConsume
          | $numOutput.add(1);
          | ${consume(ctx, input)}
        """.stripMargin.trim

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala

@@ -187,9 +187,15 @@ case class BroadcastHashJoinExec(
   private def getJoinCondition(
       ctx: CodegenContext,
       input: Seq[ExprCode],
+      uniqueKeyCodePath: Boolean,
       anti: Boolean = false): (String, String, Seq[ExprCode]) = {
     val matched = ctx.freshName("matched")
     val buildVars = genBuildSideVars(ctx, matched)
+    val continueStatement = if (uniqueKeyCodePath) {
+      continueStatementInDoConsume
+    } else {
+      "continue;"
+    }
     val checkCondition = if (condition.isDefined) {
       val expr = condition.get
       // evaluate the variables from build side that used by condition
@@ -198,18 +204,18 @@ case class BroadcastHashJoinExec(
       ctx.currentVars = input ++ buildVars
       val ev =
         BindReferences.bindReference(expr, streamedPlan.output ++ buildPlan.output).genCode(ctx)
-      val skipRow = if (!anti) {
+      val skipRow = if (!(anti && uniqueKeyCodePath)) {
         s"${ev.isNull} || !${ev.value}"
       } else {
         s"!${ev.isNull} && ${ev.value}"
       }
       s"""
          |$eval
          |${ev.code}
-         |if ($skipRow) continue;
+         |if ($skipRow) $continueStatement
        """.stripMargin
-    } else if (anti) {
-      "continue;"
+    } else if (anti && uniqueKeyCodePath) {
+      continueStatement
     } else {
       ""
     }
@@ -221,21 +227,22 @@ case class BroadcastHashJoinExec(
    */
   private def codegenInner(ctx: CodegenContext, input: Seq[ExprCode]): String = {
     val (broadcastRelation, relationTerm) = prepareBroadcast(ctx)
+    val uniqueKeyCodePath = broadcastRelation.value.keyIsUnique
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
-    val (matched, checkCondition, buildVars) = getJoinCondition(ctx, input)
+    val (matched, checkCondition, buildVars) = getJoinCondition(ctx, input, uniqueKeyCodePath)
     val numOutput = metricTerm(ctx, "numOutputRows")
 
     val resultVars = buildSide match {
       case BuildLeft => buildVars ++ input
       case BuildRight => input ++ buildVars
     }
-    if (broadcastRelation.value.keyIsUnique) {
+    if (uniqueKeyCodePath) {
       s"""
          |// generate join key for stream side
          |${keyEv.code}
          |// find matches from HashedRelation
          |UnsafeRow $matched = $anyNull ? null: (UnsafeRow)$relationTerm.getValue(${keyEv.value});
-         |if ($matched == null) continue;
+         |if ($matched == null) $continueStatementInDoConsume
          |$checkCondition
          |$numOutput.add(1);
          |${consume(ctx, resultVars)}
@@ -250,7 +257,7 @@ case class BroadcastHashJoinExec(
          |${keyEv.code}
          |// find matches from HashRelation
          |$iteratorCls $matches = $anyNull ? null : ($iteratorCls)$relationTerm.get(${keyEv.value});
-         |if ($matches == null) continue;
+         |if ($matches == null) $continueStatementInDoConsume
          |while ($matches.hasNext()) {
          |  UnsafeRow $matched = (UnsafeRow) $matches.next();
          |  $checkCondition
@@ -342,16 +349,17 @@ case class BroadcastHashJoinExec(
    */
   private def codegenSemi(ctx: CodegenContext, input: Seq[ExprCode]): String = {
     val (broadcastRelation, relationTerm) = prepareBroadcast(ctx)
+    val uniqueKeyCodePath = broadcastRelation.value.keyIsUnique
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
-    val (matched, checkCondition, _) = getJoinCondition(ctx, input)
+    val (matched, checkCondition, _) = getJoinCondition(ctx, input, uniqueKeyCodePath)
     val numOutput = metricTerm(ctx, "numOutputRows")
-    if (broadcastRelation.value.keyIsUnique) {
+    if (uniqueKeyCodePath) {
       s"""
          |// generate join key for stream side
          |${keyEv.code}
          |// find matches from HashedRelation
          |UnsafeRow $matched = $anyNull ? null: (UnsafeRow)$relationTerm.getValue(${keyEv.value});
-         |if ($matched == null) continue;
+         |if ($matched == null) $continueStatementInDoConsume
          |$checkCondition
          |$numOutput.add(1);
          |${consume(ctx, input)}
@@ -365,14 +373,14 @@ case class BroadcastHashJoinExec(
          |${keyEv.code}
          |// find matches from HashRelation
          |$iteratorCls $matches = $anyNull ? null : ($iteratorCls)$relationTerm.get(${keyEv.value});
-         |if ($matches == null) continue;
+         |if ($matches == null) $continueStatementInDoConsume
          |boolean $found = false;
          |while (!$found && $matches.hasNext()) {
          |  UnsafeRow $matched = (UnsafeRow) $matches.next();
          |  $checkCondition
          |  $found = true;
          |}
-         |if (!$found) continue;
+         |if (!$found) $continueStatementInDoConsume
          |$numOutput.add(1);
          |${consume(ctx, input)}
        """.stripMargin
@@ -386,7 +394,7 @@ case class BroadcastHashJoinExec(
     val (broadcastRelation, relationTerm) = prepareBroadcast(ctx)
     val uniqueKeyCodePath = broadcastRelation.value.keyIsUnique
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
-    val (matched, checkCondition, _) = getJoinCondition(ctx, input, uniqueKeyCodePath)
+    val (matched, checkCondition, _) = getJoinCondition(ctx, input, uniqueKeyCodePath, true)
     val numOutput = metricTerm(ctx, "numOutputRows")
 
     if (uniqueKeyCodePath) {
@@ -424,7 +432,7 @@ case class BroadcastHashJoinExec(
          |      $checkCondition
          |      $found = true;
          |    }
-         |    if ($found) continue;
+         |    if ($found) $continueStatementInDoConsume
          |  }
          |}
          |$numOutput.add(1);

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala

@@ -547,6 +547,8 @@ case class SortMergeJoinExec(
     }
   }
 
+  override protected def doConsumeInChainOfFunc: Boolean = false
+
   override def doProduce(ctx: CodegenContext): String = {
     ctx.copyResult = true
     val leftInput = ctx.freshName("leftInput")