move shouldStop() to end of while loop

Author: Davies Liu

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

@@ -255,11 +255,11 @@ private[sql] case class DataSourceScan(
       |       $numOutputRows.add(numRows);
       |     }
       |
-      |     while (!shouldStop() && $idx < numRows) {
+      |     while ($idx < numRows) {
       |       int $rowidx = $idx++;
       |       ${consume(ctx, columns1).trim}
+      |       if (shouldStop()) return;
       |     }
-      |     if (shouldStop()) return;
       |
       |     if (!$input.hasNext()) {
       |       $batch = null;
@@ -280,14 +280,15 @@ private[sql] case class DataSourceScan(
       s"""
        | private void $scanRows(InternalRow $row) throws java.io.IOException {
        |   boolean firstRow = true;
-       |   while (!shouldStop() && (firstRow || $input.hasNext())) {
+       |   while (firstRow || $input.hasNext()) {
        |     if (firstRow) {
        |       firstRow = false;
        |     } else {
        |       $row = (InternalRow) $input.next();
        |     }
        |     $numOutputRows.add(1);
        |     ${consume(ctx, columns2, inputRow).trim}
+       |     if (shouldStop()) return;
        |   }
        | }""".stripMargin)
 

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

@@ -104,11 +104,12 @@ trait CodegenSupport extends SparkPlan {
     *     # call child.produce()
     *     initialized = true;
     *   }
-    *   while (!shouldStop() && hashmap.hasNext()) {
+    *   while (hashmap.hasNext()) {
     *     row = hashmap.next();
     *     # build the aggregation results
     *     # create variables for results
     *     # call consume(), which will call parent.doConsume()
+   *      if (shouldStop()) return;
     *   }
     */
   protected def doProduce(ctx: CodegenContext): String
@@ -252,9 +253,10 @@ case class InputAdapter(child: SparkPlan) extends UnaryNode with CodegenSupport
     ctx.currentVars = null
     val columns = exprs.map(_.gen(ctx))
     s"""
-       | while (!shouldStop() && $input.hasNext()) {
+       | while ($input.hasNext()) {
        |   InternalRow $row = (InternalRow) $input.next();
        |   ${consume(ctx, columns, row).trim}
+       |   if (shouldStop()) return;
        | }
      """.stripMargin
   }
@@ -321,7 +323,7 @@ case class WholeStageCodegen(child: SparkPlan) extends UnaryNode with CodegenSup
       /** Codegened pipeline for:
         * ${toCommentSafeString(child.treeString.trim)}
         */
-      class GeneratedIterator extends org.apache.spark.sql.execution.BufferedRowIterator {
+      final class GeneratedIterator extends org.apache.spark.sql.execution.BufferedRowIterator {
 
         private Object[] references;
         ${ctx.declareMutableStates()}

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

@@ -282,13 +282,14 @@ case class Range(
       |   }
       | }
       |
-      | while (!$overflow && $checkEnd && !shouldStop()) {
+      | while (!$overflow && $checkEnd) {
       |  long $value = $number;
       |  $number += ${step}L;
       |  if ($number < $value ^ ${step}L < 0) {
       |    $overflow = true;
       |  }
       |  ${consume(ctx, Seq(ev))}
+      |  if (shouldStop()) return;
       | }
      """.stripMargin
   }