Address reviewer comments.

Author: viirya

sql/core/src/main/scala/org/apache/spark/sql/execution/BufferedRowIterator.java

@@ -35,7 +35,7 @@ public abstract class BufferedRowIterator {
   // used when there is no column in output
   protected UnsafeRow unsafeRow = new UnsafeRow(0);
 
-  private boolean stopEarly = false;
+  protected boolean stopEarly = false;
 
   public boolean hasNext() throws IOException {
     if (currentRows.isEmpty()) {
@@ -66,7 +66,7 @@ protected void append(InternalRow row) {
    * If it returns true, the caller should exit the loop (return from processNext()).
    */
   protected boolean shouldStop() {
-    return !currentRows.isEmpty() || stopEarly;
+    return !currentRows.isEmpty();
   }
 
   /**

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

@@ -67,13 +67,18 @@ trait BaseLimit extends UnaryNode with CodegenSupport {
   }
 
   override def doConsume(ctx: CodegenContext, input: Seq[ExprCode]): String = {
+    ctx.addNewFunction("shouldStop", """
+      @Override
+      protected boolean shouldStop() {
+        return !currentRows.isEmpty() || stopEarly;
+      }
+    """)
     val countTerm = ctx.freshName("count")
     ctx.addMutableState("int", countTerm, s"$countTerm = 0;")
-    ctx.currentVars = input
     s"""
        | if ($countTerm < $limit) {
        |   $countTerm += 1;
-       |   ${consume(ctx, ctx.currentVars)}
+       |   ${consume(ctx, input)}
        | } else {
        |   setStopEarly(true);
        | }

sql/core/src/test/scala/org/apache/spark/sql/execution/BenchmarkWholeStageCodegen.scala

@@ -73,14 +73,14 @@ class BenchmarkWholeStageCodegen extends SparkFunSuite {
   ignore("range/limit/sum") {
     val N = 500 << 20
     runBenchmark("range/limit/sum", N) {
-      sqlContext.range(N).limit(100).groupBy().sum().collect()
+      sqlContext.range(N).limit(1000000).groupBy().sum().collect()
     }
     /*
     Westmere E56xx/L56xx/X56xx (Nehalem-C)
     range/limit/sum:                    Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
     -------------------------------------------------------------------------------------------
-    range/limit/sum codegen=false             154 /  201       3402.0           0.3       1.0X
-    range/limit/sum codegen=true              110 /  129       4760.9           0.2       1.4X
+    range/limit/sum codegen=false             609 /  672        861.6           1.2       1.0X
+    range/limit/sum codegen=true              561 /  621        935.3           1.1       1.1X
     */
   }