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

## What changes were proposed in this pull request?

Currently `Range` reports metrics in batch granularity. This is acceptable, but it's better if we can make it row granularity without performance penalty.

Before this PR,  the metrics are updated when preparing the batch, which is before we actually consume data. In this PR, the metrics are updated after the data are consumed. There are 2 different cases:
1. The data processing loop has a stop check. The metrics are updated when we need to stop.
2. no stop check. The metrics are updated after the loop.

## How was this patch tested?

existing tests and a new benchmark

Closes #22698 from cloud-fan/range.

Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

Author: cloud-fan

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)
     }
   }