[WIP] [SPARK-51097] [SS] Split apart SparkPlan metrics and instance metrics

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

@@ -140,11 +140,17 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
    */
   def metrics: Map[String, SQLMetric] = Map.empty
 
+  /**
+   * @return All instance metrics of this SparkPlan.
+   */
+  def instanceMetrics: Map[String, SQLMetric] = Map.empty
+
   /**
    * Resets all the metrics.
    */
   def resetMetrics(): Unit = {
     metrics.valuesIterator.foreach(_.reset())
+    instanceMetrics.valuesIterator.foreach(_.reset())
     children.foreach(_.resetMetrics())
   }
 
@@ -153,6 +159,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
    */
   def longMetric(name: String): SQLMetric = metrics(name)
 
+  /**
+   * @return The instance-specific [[SQLMetric]] for the given `name`.
+   */
+  def longInstanceMetric(name: String): SQLMetric = instanceMetrics(name)
+
   // TODO: Move to `DistributedPlan`
   /**
    * Specifies how data is partitioned across different nodes in the cluster.

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

@@ -216,7 +216,9 @@ trait StateStoreWriter
     "stateMemory" -> SQLMetrics.createSizeMetric(sparkContext, "memory used by state"),
     "numStateStoreInstances" -> SQLMetrics.createMetric(sparkContext,
       "number of state store instances")
-  ) ++ stateStoreCustomMetrics ++ pythonMetrics ++ stateStoreInstanceMetrics
+  ) ++ stateStoreCustomMetrics ++ pythonMetrics
+
+  override lazy val instanceMetrics = stateStoreInstanceMetrics
 
   val stateStoreNames: Seq[String] = Seq(StateStoreId.DEFAULT_STORE_NAME)
 
@@ -332,7 +334,7 @@ trait StateStoreWriter
         case (name, metricConfig) =>
           // Keep instance metrics that are updated or aren't marked to be ignored,
           // as their initial value could still be important.
-          !metricConfig.ignoreIfUnchanged || !longMetric(name).isZero
+          !metricConfig.ignoreIfUnchanged || !longInstanceMetric(name).isZero
       }
       .groupBy {
         // Group all instance metrics underneath their common metric prefix
@@ -347,8 +349,8 @@ trait StateStoreWriter
           metrics
             .map {
               case (_, metric) =>
-                metric.name -> (if (longMetric(metric.name).isZero) metricConf.initValue
-                                else longMetric(metric.name).value)
+                metric.name -> (if (longInstanceMetric(metric.name).isZero) metricConf.initValue
+                                else longInstanceMetric(metric.name).value)
             }
             .toSeq
             .sortBy(_._2)(metricConf.ordering)
@@ -439,7 +441,8 @@ trait StateStoreWriter
       case (metric, value) =>
         val metricConfig = instanceMetricConfiguration(metric.name)
         // Update the metric's value based on the defined combine method
-        longMetric(metric.name).set(metricConfig.combine(longMetric(metric.name), value))
+        longInstanceMetric(metric.name)
+          .set(metricConfig.combine(longInstanceMetric(metric.name), value))
     }
   }