[SPARK-31253][SQL][FOLLOW-UP] Improve the partition data size metrics in CustomShuffleReaderExec

### What changes were proposed in this pull request?
Currently the partition data size metrics contain three entries (min/max/avg)  in Spark UI, which is not user friendly. This PR lets the metrics with min/max/avg in one entry by calling SQLMetrics.postDriverMetricUpdates multiple times.
Before this PR, the spark UI is shown in the following:
![image](https://user-images.githubusercontent.com/11972570/78980137-da1a2200-7b4f-11ea-81ee-76858e887bde.png)

After this PR. the spark UI is shown in the following:
![image](https://user-images.githubusercontent.com/11972570/78980192-fae27780-7b4f-11ea-9faa-07f58699acfd.png)

### Why are the changes needed?
Improving UI

### Does this PR introduce any user-facing change?
No

### How was this patch tested?
existing ut

Closes apache#28175 from JkSelf/improveAqeMetrics.

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

Author: JkSelf

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CustomShuffleReaderExec.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.execution.adaptive
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
@@ -95,54 +97,68 @@ case class CustomShuffleReaderExec private(
     case _ => None
   }
 
-  private def partitionDataSizeMetrics = {
-    val maxSize = SQLMetrics.createSizeMetric(sparkContext, "maximum partition data size")
-    val minSize = SQLMetrics.createSizeMetric(sparkContext, "minimum partition data size")
-    val avgSize = SQLMetrics.createSizeMetric(sparkContext, "average partition data size")
-    val mapStatsOpt = shuffleStage.get.mapStats
-    val sizes = mapStatsOpt.map { mapStats =>
-      val mapSizes = mapStats.bytesByPartitionId
-      partitionSpecs.map {
-        case CoalescedPartitionSpec(startReducerIndex, endReducerIndex) =>
-          startReducerIndex.until(endReducerIndex).map(mapSizes).sum
-        case p: PartialReducerPartitionSpec => p.dataSize
-        case p => throw new IllegalStateException("unexpected " + p)
+  private def sendDriverMetrics(): Unit = {
+    val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+    var driverAccumUpdates: Seq[(Long, Long)] = Seq.empty
+
+    val numPartitionsMetric = metrics("numPartitions")
+    numPartitionsMetric.set(partitionSpecs.length)
+    driverAccumUpdates = driverAccumUpdates :+
+      (numPartitionsMetric.id, partitionSpecs.length.toLong)
+
+    if (hasSkewedPartition) {
+      val skewedMetric = metrics("numSkewedPartitions")
+      val numSkewedPartitions = partitionSpecs.collect {
+        case p: PartialReducerPartitionSpec => p.reducerIndex
+      }.distinct.length
+      skewedMetric.set(numSkewedPartitions)
+      driverAccumUpdates = driverAccumUpdates :+ (skewedMetric.id, numSkewedPartitions.toLong)
+    }
+
+    if(!isLocalReader) {
+      val partitionMetrics = metrics("partitionDataSize")
+      val mapStats = shuffleStage.get.mapStats
+
+      if (mapStats.isEmpty) {
+        partitionMetrics.set(0)
+        driverAccumUpdates = driverAccumUpdates :+ (partitionMetrics.id, 0L)
+      } else {
+        var sum = 0L
+        partitionSpecs.foreach {
+          case CoalescedPartitionSpec(startReducerIndex, endReducerIndex) =>
+            val dataSize = startReducerIndex.until(endReducerIndex).map(
+              mapStats.get.bytesByPartitionId(_)).sum
+            driverAccumUpdates = driverAccumUpdates :+ (partitionMetrics.id, dataSize)
+            sum += dataSize
+          case p: PartialReducerPartitionSpec =>
+            driverAccumUpdates = driverAccumUpdates :+ (partitionMetrics.id, p.dataSize)
+            sum += p.dataSize
+          case p => throw new IllegalStateException("unexpected " + p)
+        }
+
+        // Set sum value to "partitionDataSize" metric.
+        partitionMetrics.set(sum)
       }
-    }.getOrElse(Seq(0L))
-
-    maxSize.set(sizes.max)
-    minSize.set(sizes.min)
-    avgSize.set(sizes.sum / sizes.length)
-    Map(
-      "maxPartitionDataSize" -> maxSize,
-      "minPartitionDataSize" -> minSize,
-      "avgPartitionDataSize" -> avgSize)
-  }
+    }
 
-  private def skewedPartitionMetrics = {
-    val metrics = SQLMetrics.createMetric(sparkContext, "number of skewed partitions")
-    val numSkewedPartitions = partitionSpecs.collect {
-      case p: PartialReducerPartitionSpec => p.reducerIndex
-    }.distinct.length
-    metrics.set(numSkewedPartitions)
-    Map("numSkewedPartitions" -> metrics)
+    SQLMetrics.postDriverMetricsUpdatedByValue(sparkContext, executionId, driverAccumUpdates)
   }
 
   @transient override lazy val metrics: Map[String, SQLMetric] = {
     if (shuffleStage.isDefined) {
-      val numPartitions = SQLMetrics.createMetric(sparkContext, "number of partitions")
-      numPartitions.set(partitionSpecs.length)
-      Map("numPartitions" -> numPartitions) ++ {
+      Map("numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions")) ++ {
         if (isLocalReader) {
           // We split the mapper partition evenly when creating local shuffle reader, so no
           // data size info is available.
           Map.empty
         } else {
-          partitionDataSizeMetrics
+          Map("partitionDataSize" ->
+            SQLMetrics.createSizeMetric(sparkContext, "partition data size"))
         }
       } ++ {
         if (hasSkewedPartition) {
-          skewedPartitionMetrics
+          Map("numSkewedPartitions" ->
+            SQLMetrics.createMetric(sparkContext, "number of skewed partitions"))
         } else {
           Map.empty
         }
@@ -154,8 +170,8 @@ case class CustomShuffleReaderExec private(
   }
 
   private lazy val cachedShuffleRDD: RDD[InternalRow] = {
-    val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
-    SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq)
+    sendDriverMetrics()
+
     shuffleStage.map { stage =>
       new ShuffledRowRDD(
         stage.shuffle.shuffleDependency, stage.shuffle.readMetrics, partitionSpecs.toArray)

sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala

@@ -222,6 +222,16 @@ object SQLMetrics {
     }
   }
 
+  def postDriverMetricsUpdatedByValue(
+      sc: SparkContext,
+      executionId: String,
+      accumUpdates: Seq[(Long, Long)]): Unit = {
+    if (executionId != null) {
+      sc.listenerBus.post(
+        SparkListenerDriverAccumUpdates(executionId.toLong, accumUpdates))
+    }
+  }
+
   /**
    * Updates metrics based on the driver side value. This is useful for certain metrics that
    * are only updated on the driver, e.g. subquery execution time, or number of files.

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListener.scala

@@ -450,7 +450,7 @@ private class LiveExecutionData(val executionId: Long) extends LiveEntity {
 
   var jobs = Map[Int, JobExecutionStatus]()
   var stages = Set[Int]()
-  var driverAccumUpdates = Map[Long, Long]()
+  var driverAccumUpdates = Seq[(Long, Long)]()
 
   @volatile var metricsValues: Map[Long, String] = null
 

sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala

@@ -809,11 +809,9 @@ class AdaptiveQueryExecSuite
       assert(!reader.hasSkewedPartition)
       assert(reader.hasCoalescedPartition)
       assert(reader.metrics.keys.toSeq.sorted == Seq(
-        "avgPartitionDataSize", "maxPartitionDataSize", "minPartitionDataSize", "numPartitions"))
+        "numPartitions", "partitionDataSize"))
       assert(reader.metrics("numPartitions").value == reader.partitionSpecs.length)
-      assert(reader.metrics("avgPartitionDataSize").value > 0)
-      assert(reader.metrics("maxPartitionDataSize").value > 0)
-      assert(reader.metrics("minPartitionDataSize").value > 0)
+      assert(reader.metrics("partitionDataSize").value > 0)
 
       withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
         val (_, adaptivePlan) = runAdaptiveAndVerifyResult(