Address review comments from @tdas

Author: HeartSaVioR

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

@@ -81,10 +81,10 @@ class SQLMetric(val metricType: String, initValue: Long = 0L) extends Accumulato
 }
 
 object SQLMetrics {
-  val SUM_METRIC = "sum"
-  val SIZE_METRIC = "size"
-  val TIMING_METRIC = "timing"
-  val AVERAGE_METRIC = "average"
+  private val SUM_METRIC = "sum"
+  private val SIZE_METRIC = "size"
+  private val TIMING_METRIC = "timing"
+  private val AVERAGE_METRIC = "average"
 
   private val baseForAvgMetric: Int = 10
 

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala

@@ -90,13 +90,11 @@ trait StateStoreWriter extends StatefulOperator { self: SparkPlan =>
    * the driver after this SparkPlan has been executed and metrics have been updated.
    */
   def getProgress(): StateOperatorProgress = {
-    // average metric is a bit tricky, so hard to aggregate: just exclude them to simplify issue
-    val avgExcludedCustomMetrics = stateStoreCustomMetrics
-      .filterNot(_._2.metricType == SQLMetrics.AVERAGE_METRIC)
+    val customMetrics = stateStoreCustomMetrics
       .map(entry => entry._1 -> longMetric(entry._1).value)
 
     val javaConvertedCustomMetrics: java.util.HashMap[String, java.lang.Long] =
-      new java.util.HashMap(avgExcludedCustomMetrics.mapValues(long2Long).asJava)
+      new java.util.HashMap(customMetrics.mapValues(long2Long).asJava)
 
     new StateOperatorProgress(
       numRowsTotal = longMetric("numTotalStateRows").value,

sql/core/src/main/scala/org/apache/spark/sql/streaming/progress.scala

@@ -52,21 +52,17 @@ class StateOperatorProgress private[sql](
     new StateOperatorProgress(numRowsTotal, newNumRowsUpdated, memoryUsedBytes, customMetrics)
 
   private[sql] def jsonValue: JValue = {
-    def safeMapToJValue[T](map: ju.Map[String, T], valueToJValue: T => JValue): JValue = {
-      if (map.isEmpty) return JNothing
-      val keys = map.keySet.asScala.toSeq.sorted
-      keys.map { k => k -> valueToJValue(map.get(k)) : JObject }.reduce(_ ~ _)
-    }
-
-    val jsonVal = ("numRowsTotal" -> JInt(numRowsTotal)) ~
-      ("numRowsUpdated" -> JInt(numRowsUpdated)) ~
-      ("memoryUsedBytes" -> JInt(memoryUsedBytes))
-
-    if (!customMetrics.isEmpty) {
-      jsonVal ~ ("customMetrics" -> safeMapToJValue[JLong](customMetrics, v => JInt(v.toLong)))
-    } else {
-      jsonVal
-    }
+    ("numRowsTotal" -> JInt(numRowsTotal)) ~
+    ("numRowsUpdated" -> JInt(numRowsUpdated)) ~
+    ("memoryUsedBytes" -> JInt(memoryUsedBytes)) ~
+    ("customMetrics" -> {
+      if (!customMetrics.isEmpty) {
+        val keys = customMetrics.keySet.asScala.toSeq.sorted
+        keys.map { k => k -> JInt(customMetrics.get(k).toLong) : JObject }.reduce(_ ~ _)
+      } else {
+        JNothing
+      }
+    })
   }
 
   override def toString: String = prettyJson