[SPARK-23375][SQL][FOLLOWUP][TEST] Test Sort metrics while Sort is missing

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

@@ -26,7 +26,7 @@ import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
 import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
-import org.apache.spark.sql.execution.{FilterExec, RangeExec, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.{FilterExec, RangeExec, SortExec, SparkPlan, WholeStageCodegenExec}
 import org.apache.spark.sql.execution.aggregate.HashAggregateExec
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
@@ -194,10 +194,22 @@ class SQLMetricsSuite extends SparkFunSuite with SQLMetricsTestUtils with Shared
   }
 
   test("Sort metrics") {
-    // Assume the execution plan is
-    // WholeStageCodegen(nodeId = 0, Range(nodeId = 2) -> Sort(nodeId = 1))
-    val ds = spark.range(10).sort('id)
-    testSparkPlanMetrics(ds.toDF(), 2, Map.empty)
+    // Assume the execution plan with node id is
+    // Sort(nodeId = 0)
+    //   Exchange(nodeId = 1)
+    //     Project(nodeId = 2)
+    //       LocalTableScan(nodeId = 3)
+    // Because of SPARK-25267, ConvertToLocalRelation is disabled in the test cases of sql/core,
+    // so Project here is not collapsed into LocalTableScan.
+    val df = Seq(1, 3, 2).toDF("id").sort('id)
+    testSparkPlanMetricsWithPredicates(df, 2, Map(
+      0L -> (("Sort", Map(
+        // In SortExec, sort time is collected as nanoseconds, but it is converted and stored as
+        // milliseconds. So sort time may be 0 if sort is executed very fast.
+        "sort time total (min, med, max)" -> timingMetricAllStatsShould(_ >= 0),
+        "peak memory total (min, med, max)" -> sizeMetricAllStatsShould(_ > 0),
+        "spill size total (min, med, max)" -> sizeMetricAllStatsShould(_ >= 0))))
+    ))
   }
 
   test("SortMergeJoin metrics") {

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

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.execution.metric
 
 import java.io.File
+import java.util.regex.Pattern
 
 import scala.collection.mutable.HashMap
 
@@ -40,6 +41,10 @@ trait SQLMetricsTestUtils extends SQLTestUtils {
 
   protected def statusStore: SQLAppStatusStore = spark.sharedState.statusStore
 
+  protected val bytesPattern = Pattern.compile("([0-9]+(\\.[0-9]+)?) (EiB|PiB|TiB|GiB|MiB|KiB|B)")
+
+  protected val durationPattern = Pattern.compile("([0-9]+(\\.[0-9]+)?) (ms|s|m|h)")
+
   /**
    * Get execution metrics for the SQL execution and verify metrics values.
    *
@@ -185,19 +190,105 @@ trait SQLMetricsTestUtils extends SQLTestUtils {
       df: DataFrame,
       expectedNumOfJobs: Int,
       expectedMetrics: Map[Long, (String, Map[String, Any])]): Unit = {
-    val optActualMetrics = getSparkPlanMetrics(df, expectedNumOfJobs, expectedMetrics.keySet)
+    val expectedMetricsPredicates = expectedMetrics.mapValues { case (nodeName, nodeMetrics) =>
+      (nodeName, nodeMetrics.mapValues(expectedMetricValue =>
+        (actualMetricValue: Any) => expectedMetricValue.toString === actualMetricValue)
+    )}
+    testSparkPlanMetricsWithPredicates(df, expectedNumOfJobs, expectedMetricsPredicates)
+  }
+
+  /**
+   * Call `df.collect()` and verify if the collected metrics satisfy the specified predicates.
+   * @param df `DataFrame` to run
+   * @param expectedNumOfJobs number of jobs that will run
+   * @param expectedMetricsPredicates the expected metrics predicates. The format is
+   *                        `nodeId -> (operatorName, metric name -> metric value predicate)`.
+   */
+  protected def testSparkPlanMetricsWithPredicates(
+      df: DataFrame,
+      expectedNumOfJobs: Int,
+      expectedMetricsPredicates: Map[Long, (String, Map[String, Any => Boolean])]): Unit = {
+    val optActualMetrics =
+      getSparkPlanMetrics(df, expectedNumOfJobs, expectedMetricsPredicates.keySet)
     optActualMetrics.foreach { actualMetrics =>
-      assert(expectedMetrics.keySet === actualMetrics.keySet)
-      for (nodeId <- expectedMetrics.keySet) {
-        val (expectedNodeName, expectedMetricsMap) = expectedMetrics(nodeId)
+      assert(expectedMetricsPredicates.keySet === actualMetrics.keySet)
+      for (nodeId <- expectedMetricsPredicates.keySet) {
+        val (expectedNodeName, expectedMetricsPredicatesMap) = expectedMetricsPredicates(nodeId)
         val (actualNodeName, actualMetricsMap) = actualMetrics(nodeId)
         assert(expectedNodeName === actualNodeName)
-        for (metricName <- expectedMetricsMap.keySet) {
-          assert(expectedMetricsMap(metricName).toString === actualMetricsMap(metricName))
+        for (metricName <- expectedMetricsPredicatesMap.keySet) {
+          assert(expectedMetricsPredicatesMap(metricName)(actualMetricsMap(metricName)))
         }
       }
     }
   }
+
+  private def metricStats(metricStr: String): Seq[String] = {
+    val sum = metricStr.substring(0, metricStr.indexOf("(")).stripPrefix("\n").stripSuffix(" ")
+    val minMedMax = metricStr.substring(metricStr.indexOf("(") + 1, metricStr.indexOf(")"))
+      .split(", ").toSeq
+    (sum +: minMedMax)
+  }
+
+  private def stringToBytes(str: String): (Float, String) = {
+    val matcher = bytesPattern.matcher(str)
+    if (matcher.matches()) {
+      (matcher.group(1).toFloat, matcher.group(3))
+    } else {
+      throw new NumberFormatException("Failed to parse byte string: " + str)
+    }
+  }
+
+  private def stringToDuration(str: String): (Float, String) = {
+    val matcher = durationPattern.matcher(str)
+    if (matcher.matches()) {
+      (matcher.group(1).toFloat, matcher.group(3))
+    } else {
+      throw new NumberFormatException("Failed to parse time string: " + str)
+    }
+  }
+
+  /**
+   * Convert a size metric string to a sequence of stats, including sum, min, med and max in order,
+   * each a tuple of (value, unit).
+   * @param metricStr size metric string, e.g. "\n96.2 MB (32.1 MB, 32.1 MB, 32.1 MB)"
+   * @return A sequence of stats, e.g. ((96.2,MB), (32.1,MB), (32.1,MB), (32.1,MB))
+   */
+  protected def sizeMetricStats(metricStr: String): Seq[(Float, String)] = {
+    metricStats(metricStr).map(stringToBytes)
+  }
+
+  /**
+   * Convert a timing metric string to a sequence of stats, including sum, min, med and max in
+   * order, each a tuple of (value, unit).
+   * @param metricStr timing metric string, e.g. "\n2.0 ms (1.0 ms, 1.0 ms, 1.0 ms)"
+   * @return A sequence of stats, e.g. ((2.0,ms), (1.0,ms), (1.0,ms), (1.0,ms))
+   */
+  protected def timingMetricStats(metricStr: String): Seq[(Float, String)] = {
+    metricStats(metricStr).map(stringToDuration)
+  }
+
+  /**
+   * Returns a function to check whether all stats (sum, min, med and max) of a timing metric
+   * satisfy the specified predicate.
+   * @param predicate predicate to check stats
+   * @return function to check all stats of a timing metric
+   */
+  protected def timingMetricAllStatsShould(predicate: Float => Boolean): Any => Boolean = {
+    (timingMetric: Any) =>
+      timingMetricStats(timingMetric.toString).forall { case (duration, _) => predicate(duration) }
+  }
+
+  /**
+   * Returns a function to check whether all stats (sum, min, med and max) of a size metric satisfy
+   * the specified predicate.
+   * @param predicate predicate to check stats
+   * @return function to check all stats of a size metric
+   */
+  protected def sizeMetricAllStatsShould(predicate: Float => Boolean): Any => Boolean = {
+    (sizeMetric: Any) =>
+      sizeMetricStats(sizeMetric.toString).forall { case (bytes, _) => predicate(bytes)}
+  }
 }