[SPARK-25710][SQL] range should report metrics correctly

sql/core/benchmarks/RangeBenchmark-results.txt

@@ -0,0 +1,16 @@
+================================================================================================
+range
+================================================================================================
+
+Java HotSpot(TM) 64-Bit Server VM 1.8.0_161-b12 on Mac OS X 10.13.6
+Intel(R) Core(TM) i7-6920HQ CPU @ 2.90GHz
+
+range:                                   Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------
+full scan                                   12674 / 12840         41.4          24.2       1.0X
+limit after range                               33 /   37      15900.2           0.1     384.4X
+filter after range                             969 /  985        541.0           1.8      13.1X
+count after range                               42 /   42      12510.5           0.1     302.4X
+count after limit after range                   32 /   33      16337.0           0.1     394.9X
+
+

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

@@ -452,8 +452,15 @@ case class RangeExec(range: org.apache.spark.sql.catalyst.plans.logical.Range)
 
     val localIdx = ctx.freshName("localIdx")
     val localEnd = ctx.freshName("localEnd")
-    val shouldStop = if (parent.needStopCheck) {
-      s"if (shouldStop()) { $nextIndex = $value + ${step}L; return; }"
+    val stopCheck = if (parent.needStopCheck) {
+      s"""
+         |if (shouldStop()) {
+         |  $nextIndex = $value + ${step}L;
+         |  $numOutput.add($localIdx + 1);
+         |  $inputMetrics.incRecordsRead($localIdx + 1);
+         |  return;
+         |}
+       """.stripMargin
     } else {
       "// shouldStop check is eliminated"
     }
@@ -506,18 +513,18 @@ case class RangeExec(range: org.apache.spark.sql.catalyst.plans.logical.Range)
       |       $numElementsTodo = 0;
       |       if ($nextBatchTodo == 0) break;
       |     }
-      |     $numOutput.add($nextBatchTodo);
-      |     $inputMetrics.incRecordsRead($nextBatchTodo);
       |     $batchEnd += $nextBatchTodo * ${step}L;
       |   }
       |
       |   int $localEnd = (int)(($batchEnd - $nextIndex) / ${step}L);
       |   for (int $localIdx = 0; $localIdx < $localEnd; $localIdx++) {
       |     long $value = ((long)$localIdx * ${step}L) + $nextIndex;
       |     ${consume(ctx, Seq(ev))}
-      |     $shouldStop
+      |     $stopCheck
       |   }
       |   $nextIndex = $batchEnd;
+      |   $numOutput.add($localEnd);
+      |   $inputMetrics.incRecordsRead($localEnd);
       |   $taskContext.killTaskIfInterrupted();
       | }
      """.stripMargin

sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RangeBenchmark.scala

@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+
+/**
+ * Benchmark to measure performance for range operator.
+ * To run this benchmark:
+ * {{{
+ *   1. without sbt:
+ *      bin/spark-submit --class <this class> --jars <spark core test jar> <spark sql test jar>
+ *   2. build/sbt "sql/test:runMain <this class>"
+ *   3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/test:runMain <this class>"
+ *      Results will be written to "benchmarks/RangeBenchmark-results.txt".
+ * }}}
+ */
+object RangeBenchmark extends SqlBasedBenchmark {
+
+  override def runBenchmarkSuite(): Unit = {
+    import spark.implicits._
+
+    runBenchmark("range") {
+      val N = 500L << 20
+      val benchmark = new Benchmark("range", N, output = output)
+
+      benchmark.addCase("full scan", numIters = 4) { _ =>
+        spark.range(N).queryExecution.toRdd.foreach(_ => ())
+      }
+
+      benchmark.addCase("limit after range", numIters = 4) { _ =>
+        spark.range(N).limit(100).queryExecution.toRdd.foreach(_ => ())
+      }
+
+      benchmark.addCase("filter after range", numIters = 4) { _ =>
+        spark.range(N).filter('id % 100 === 0).queryExecution.toRdd.foreach(_ => ())
+      }
+
+      benchmark.addCase("count after range", numIters = 4) { _ =>
+        spark.range(N).count()
+      }
+
+      benchmark.addCase("count after limit after range", numIters = 4) { _ =>
+        spark.range(N).limit(100).count()
+      }
+
+      benchmark.run()
+    }
+  }
+}

sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala

@@ -556,16 +556,16 @@ class SQLMetricsSuite extends SparkFunSuite with SQLMetricsTestUtils with Shared
 
       df.queryExecution.executedPlan.foreach(_.resetMetrics())
       // For each partition, we get 2 rows. Then the Filter should produce 2 rows per-partition,
-      // and Range should produce 1000 rows (one batch) per-partition. Totally Filter produces
-      // 4 rows, and Range produces 2000 rows.
+      // and Range should produce 4 rows per-partition ([0, 1, 2, 3] and [15, 16, 17, 18]). Totally
+      // Filter produces 4 rows, and Range produces 8 rows.
       df.queryExecution.toRdd.mapPartitions(_.take(2)).collect()
-      checkFilterAndRangeMetrics(df, filterNumOutputs = 4, rangeNumOutputs = 2000)
+      checkFilterAndRangeMetrics(df, filterNumOutputs = 4, rangeNumOutputs = 8)
 
       // Top-most limit will call `CollectLimitExec.executeCollect`, which will only run the first
-      // task, so totally the Filter produces 2 rows, and Range produces 1000 rows (one batch).
+      // task, so totally the Filter produces 2 rows, and Range produces 4 rows ([0, 1, 2, 3]).
       val df2 = df.limit(2)
       df2.collect()
-      checkFilterAndRangeMetrics(df2, filterNumOutputs = 2, rangeNumOutputs = 1000)
+      checkFilterAndRangeMetrics(df2, filterNumOutputs = 2, rangeNumOutputs = 4)
     }
   }