address comment

Author: cloud-fan

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

@@ -136,7 +136,7 @@ private[sql] trait ColumnarBatchScan extends CodegenSupport {
        |if ($batch == null) {
        |  $nextBatchFuncName();
        |}
-       |while ($batch != null$keepProducingDataCond) {
+       |while ($batch != null$limitNotReachedCond) {
        |  int $numRows = $batch.numRows();
        |  int $localEnd = $numRows - $idx;
        |  for (int $localIdx = 0; $localIdx < $localEnd; $localIdx++) {
@@ -166,7 +166,7 @@ private[sql] trait ColumnarBatchScan extends CodegenSupport {
     }
     val inputRow = if (needsUnsafeRowConversion) null else row
     s"""
-       |while ($input.hasNext()$keepProducingDataCond) {
+       |while ($input.hasNext()$limitNotReachedCond) {
        |  InternalRow $row = (InternalRow) $input.next();
        |  $numOutputRows.add(1);
        |  ${consume(ctx, outputVars, inputRow).trim}

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

@@ -132,11 +132,12 @@ case class SortExec(
   // a stop check before sorting.
   override def needStopCheck: Boolean = false
 
-  // Sort is a blocking operator. It needs to consume all the inputs before producing any
-  // output. This means, Limit after Sort has no effect to Sort's upstream operators.
-  // Here we override this method to return Nil, so that upstream operators will not generate
-  // unnecessary conditions (which is always evaluated to false) for the Limit after Sort.
-  override def conditionsOfKeepProducingData: Seq[String] = Nil
+  // Sort is a blocking operator. It needs to consume all the inputs before producing any output.
+  // This means, Limit operator after Sort will never reach its limit during the execution of Sort's
+  // upstream operators. Here we override this method to return Nil, so that upstream operators will
+  // not generate useless conditions (which are always evaluated to false) for the Limit operators
+  // after Sort.
+  override def limitNotReachedChecks: Seq[String] = Nil
 
   override protected def doProduce(ctx: CodegenContext): String = {
     val needToSort =
@@ -178,7 +179,7 @@ case class SortExec(
        |   $needToSort = false;
        | }
        |
-       | while ($sortedIterator.hasNext()$keepProducingDataCond) {
+       | while ($sortedIterator.hasNext()$limitNotReachedCond) {
        |   UnsafeRow $outputRow = (UnsafeRow)$sortedIterator.next();
        |   ${consume(ctx, null, outputRow)}
        |   if (shouldStop()) return;

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

@@ -346,13 +346,24 @@ trait CodegenSupport extends SparkPlan {
    */
   def needStopCheck: Boolean = parent.needStopCheck
 
-  def conditionsOfKeepProducingData: Seq[String] = parent.conditionsOfKeepProducingData
+  /**
+   * A sequence of checks which evaluate to true if the downstream Limit operators have not received
+   * enough records and reached the limit. If current node is a data producing node, it can leverage
+   * this information to stop producing data and complete the data flow earlier. Common data
+   * producing nodes are leaf nodes like Range and Scan, and blocking nodes like Sort and Aggregate.
+   * These checks should be put into the loop condition of the data producing loop.
+   */
+  def limitNotReachedChecks: Seq[String] = parent.limitNotReachedChecks
 
-  final protected def keepProducingDataCond: String = {
-    if (parent.conditionsOfKeepProducingData.isEmpty) {
+  /**
+   * A helper method to generate the data producing loop condition according to the
+   * limit-not-reached checks.
+   */
+  final def limitNotReachedCond: String = {
+    if (parent.limitNotReachedChecks.isEmpty) {
       ""
     } else {
-      parent.conditionsOfKeepProducingData.mkString(" && ", " && ", "")
+      parent.limitNotReachedChecks.mkString(" && ", " && ", "")
     }
   }
 }
@@ -391,7 +402,7 @@ case class InputAdapter(child: SparkPlan) extends UnaryExecNode with CodegenSupp
       forceInline = true)
     val row = ctx.freshName("row")
     s"""
-       | while ($input.hasNext()$keepProducingDataCond) {
+       | while ($input.hasNext()$limitNotReachedCond) {
        |   InternalRow $row = (InternalRow) $input.next();
        |   ${consume(ctx, null, row).trim}
        |   if (shouldStop()) return;
@@ -687,7 +698,7 @@ case class WholeStageCodegenExec(child: SparkPlan)(val codegenStageId: Int)
 
   override def needStopCheck: Boolean = true
 
-  override def conditionsOfKeepProducingData: Seq[String] = Nil
+  override def limitNotReachedChecks: Seq[String] = Nil
 
   override protected def otherCopyArgs: Seq[AnyRef] = Seq(codegenStageId.asInstanceOf[Integer])
 }

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

@@ -160,10 +160,11 @@ case class HashAggregateExec(
   override def needStopCheck: Boolean = false
 
   // Aggregate is a blocking operator. It needs to consume all the inputs before producing any
-  // output. This means, Limit after Aggregate has no effect to Aggregate's upstream operators.
-  // Here we override this method to return Nil, so that upstream operators will not generate
-  // unnecessary conditions (which is always evaluated to false) for the Limit after Aggregate.
-  override def conditionsOfKeepProducingData: Seq[String] = Nil
+  // output. This means, Limit operator after Aggregate will never reach its limit during the
+  // execution of Aggregate's upstream operators. Here we override this method to return Nil, so
+  // that upstream operators will not generate useless conditions (which are always evaluated to
+  // true) for the Limit operators after Aggregate.
+  override def limitNotReachedChecks: Seq[String] = Nil
 
   protected override def doProduce(ctx: CodegenContext): String = {
     if (groupingExpressions.isEmpty) {
@@ -711,7 +712,7 @@ case class HashAggregateExec(
 
     def outputFromRegularHashMap: String = {
       s"""
-         |while ($iterTerm.next()$keepProducingDataCond) {
+         |while ($iterTerm.next()$limitNotReachedCond) {
          |  UnsafeRow $keyTerm = (UnsafeRow) $iterTerm.getKey();
          |  UnsafeRow $bufferTerm = (UnsafeRow) $iterTerm.getValue();
          |  $outputFunc($keyTerm, $bufferTerm);

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

@@ -490,7 +490,7 @@ case class RangeExec(range: org.apache.spark.sql.catalyst.plans.logical.Range)
       |   $initRangeFuncName(partitionIndex);
       | }
       |
-      | while (true$keepProducingDataCond) {
+      | while (true$limitNotReachedCond) {
       |   if ($nextIndex == $batchEnd) {
       |     long $nextBatchTodo;
       |     if ($numElementsTodo > ${batchSize}L) {

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

@@ -623,7 +623,7 @@ case class SortMergeJoinExec(
     }
 
     s"""
-       |while (findNextInnerJoinRows($leftInput, $rightInput)$keepProducingDataCond) {
+       |while (findNextInnerJoinRows($leftInput, $rightInput)) {
        |  ${leftVarDecl.mkString("\n")}
        |  ${beforeLoop.trim}
        |  scala.collection.Iterator<UnsafeRow> $iterator = $matches.generateIterator();

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

@@ -77,8 +77,8 @@ trait BaseLimitExec extends UnaryExecNode with CodegenSupport {
 
   private lazy val countTerm = BaseLimitExec.newLimitCountTerm()
 
-  override lazy val conditionsOfKeepProducingData: Seq[String] = {
-    s"$countTerm < $limit" +: super.conditionsOfKeepProducingData
+  override lazy val limitNotReachedChecks: Seq[String] = {
+    s"$countTerm < $limit" +: super.limitNotReachedChecks
   }
 
   protected override def doProduce(ctx: CodegenContext): String = {