[SPARK-59174][SS] State-store instance metrics are shipped in every task result

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

@@ -47,7 +47,9 @@ import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.{OutputMode, StateOperatorProgress}
 import org.apache.spark.sql.types._
-import org.apache.spark.util.{CollectionAccumulator, CompletionIterator, NextIterator, Utils}
+import org.apache.spark.sql.util.PartitionKeyedAccumulator
+import org.apache.spark.util.{AccumulatorV2, CollectionAccumulator, CompletionIterator}
+import org.apache.spark.util.{NextIterator, Utils}
 
 
 /** Used to identify the state store for a given operator.
@@ -169,6 +171,54 @@ case class StatefulOpStateStoreCheckpointInfo(
     // to validate the batch is processed based on the correct checkpoint.
     baseStateStoreCkptId: Option[Array[String]])
 
+/**
+ * An accumulator that records state store instance metrics per partition.
+ * Extends [[PartitionKeyedAccumulator]] to bound driver-side state to O(numPartitions).
+ * When duplicate updates for the same partition are received (e.g. from speculative execution,
+ * retries, or multiple stores within the same task), metrics are merged using each
+ * [[StateStoreInstanceMetric]]'s combine semantics rather than default last-write-wins.
+ */
+class StateStoreInstanceMetricAccumulator
+  extends PartitionKeyedAccumulator[Map[StateStoreInstanceMetric, Long]] {
+
+  override def copyAndReset(): StateStoreInstanceMetricAccumulator =
+    new StateStoreInstanceMetricAccumulator
+
+  override def copy(): StateStoreInstanceMetricAccumulator = synchronized {
+    val newAcc = new StateStoreInstanceMetricAccumulator
+    newAcc.byPartition.putAll(byPartition)
+    newAcc
+  }
+
+  override def add(v: (Int, Map[StateStoreInstanceMetric, Long])): Unit = synchronized {
+    byPartition.merge(v._1, v._2, (m1, m2) => combineMetrics(m1, m2))
+  }
+
+  override def merge(
+      other: AccumulatorV2[(Int, Map[StateStoreInstanceMetric, Long]),
+        java.util.Map[Int, Map[StateStoreInstanceMetric, Long]]]): Unit = synchronized {
+    other match {
+      case o: StateStoreInstanceMetricAccumulator =>
+        o.byPartition.forEach { (k, v) =>
+          byPartition.merge(k, v, (m1, m2) => combineMetrics(m1, m2))
+        }
+      case _ => throw new UnsupportedOperationException(
+        s"Cannot merge ${this.getClass.getName} with ${other.getClass.getName}")
+    }
+  }
+
+  private def combineMetrics(
+      m1: Map[StateStoreInstanceMetric, Long],
+      m2: Map[StateStoreInstanceMetric, Long]): Map[StateStoreInstanceMetric, Long] = {
+    m2.foldLeft(m1) { case (acc, (metric, v2)) =>
+      acc.get(metric) match {
+        case Some(v1) => acc.updated(metric, metric.combine(v1, v2))
+        case None => acc.updated(metric, v2)
+      }
+    }
+  }
+}
+
 /** An operator that writes to a StateStore. */
 trait StateStoreWriter
   extends StatefulOperator with PythonSQLMetrics with Logging { self: SparkPlan =>
@@ -209,8 +259,6 @@ trait StateStoreWriter
   def operatorStateMetadataVersion: Int = 1
 
   override lazy val metrics = {
-    // Lazy initialize instance metrics, but do not include these with regular metrics
-    instanceMetrics
     statefulOperatorCustomMetrics ++ Map(
       "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
       "numRowsDroppedByWatermark" -> SQLMetrics
@@ -230,21 +278,20 @@ trait StateStoreWriter
   }
 
   /**
-   * Map of all instance metrics (including partition ID and store names) to
-   * their SQLMetric counterpart.
-   *
-   * The instance metric objects hold additional information on how to report these metrics,
-   * while the SQLMetric objects store the metric values.
-   *
-   * This map is similar to the metrics map, but needs to be kept separate to prevent propagating
-   * all initialized instance metrics to SparkUI.
+   * Aggregator used for executors to pass instance metrics (per partition/store) back to driver.
+   * Extends PartitionKeyedAccumulator to bound driver-side state to O(numPartitions) while
+   * preserving StateStoreInstanceMetric.combine semantics when duplicate updates are merged
+   * (e.g. from retries, speculative execution, or multiple stores within the same task).
    */
-  lazy val instanceMetrics: Map[StateStoreInstanceMetric, SQLMetric] =
-    stateStoreInstanceMetrics
+  val instanceMetricsAccumulator: StateStoreInstanceMetricAccumulator = {
+    val acc = new StateStoreInstanceMetricAccumulator
+    SparkContext.getActive.foreach(_.register(acc))
+    acc
+  }
 
   override def resetMetrics(): Unit = {
     super.resetMetrics()
-    instanceMetrics.valuesIterator.foreach(_.reset())
+    instanceMetricsAccumulator.reset()
   }
 
   val stateStoreNames: Seq[String] = Seq(StateStoreId.DEFAULT_STORE_NAME)
@@ -354,17 +401,25 @@ trait StateStoreWriter
    * the driver after this SparkPlan has been executed and metrics have been updated.
    */
   def getProgress(): StateOperatorProgress = {
-    val instanceMetricsToReport = instanceMetrics
+    // StateStoreInstanceMetricAccumulator holds one Map[StateStoreInstanceMetric, Long] per
+    // partition, using combine semantics to deduplicate task retries and speculative execution.
+    // Folding the per-partition maps together produces a flat metric map for progress reporting.
+    val combinedMetrics: Map[StateStoreInstanceMetric, Long] =
+      instanceMetricsAccumulator.foldValues(Map.empty[StateStoreInstanceMetric, Long]) {
+        (acc, partitionMap) => acc ++ partitionMap
+      }
+
+    val instanceMetricsToReport = combinedMetrics
       .filter {
-        case (metricConf, sqlMetric) =>
+        case (metricConf, value) =>
           // Keep instance metrics that are updated or aren't marked to be ignored,
           // as their initial value could still be important.
-          !metricConf.ignoreIfUnchanged || !sqlMetric.isZero
+          !metricConf.ignoreIfUnchanged || value != metricConf.initValue
       }
       .groupBy {
         // Group all instance metrics underneath their common metric prefix
         // to ignore partition and store names.
-        case (metricConf, sqlMetric) => metricConf.metricPrefix
+        case (metricConf, _) => metricConf.metricPrefix
       }
       .flatMap {
         case (_, metrics) =>
@@ -373,12 +428,9 @@ trait StateStoreWriter
           val metricConf = metrics.head._1
           metrics
             .map {
-              case (metricConf, sqlMetric) =>
+              case (metricConf, value) =>
                 // Use metric name as it will be combined with custom metrics in progress reports.
-                // All metrics that are at their initial value at this stage should not be ignored
-                // and should show their real initial value.
-                metricConf.name -> (if (sqlMetric.isZero) metricConf.initValue
-                                    else sqlMetric.value)
+                metricConf.name -> value
             }
             .toSeq
             .sortBy(_._2)(metricConf.ordering)
@@ -407,7 +459,7 @@ trait StateStoreWriter
       numShufflePartitions = stateInfo.map(_.numPartitions.toLong).getOrElse(-1L),
       numStateStoreInstances = longMetric("numStateStoreInstances").value,
       javaConvertedCustomMetrics,
-      snapshotCustomMetricNames
+      snapshotCustomMetricNames ++ instanceMetricsToReport.keySet
     )
   }
 
@@ -466,10 +518,12 @@ trait StateStoreWriter
 
   protected def setStoreInstanceMetrics(
       otherStoreInstanceMetrics: Map[StateStoreInstanceMetric, Long]): Unit = {
-    otherStoreInstanceMetrics.foreach {
-      case (metric, value) =>
-        // Update the metric's value based on the defined combine method
-        instanceMetrics(metric).set(metric.combine(instanceMetrics(metric), value))
+    if (otherStoreInstanceMetrics.nonEmpty) {
+      // All instance metrics for a given store share the same partitionId.
+      val partitionId = otherStoreInstanceMetrics.keys.head.partitionId.getOrElse(
+        throw new IllegalStateException(
+          "StateStoreInstanceMetric must have a partitionId when reporting metrics"))
+      instanceMetricsAccumulator.add((partitionId, otherStoreInstanceMetrics))
     }
   }
 
@@ -480,29 +534,9 @@ trait StateStoreWriter
     }.toMap
   }
 
-  // All instance metrics with their (partitionId, storeName) bindings; consumed by
-  // both `stateStoreInstanceMetrics` (for SQLMetric registration) and
-  // `snapshotCustomMetricNames` (for the snapshot-name set). The result is a
-  // serializable Seq so storing it as a lazy val on this trait is safe even when
-  // the enclosing SparkPlan is shipped to executors. The provider itself is NOT
-  // stored as a field (it is non-serializable), so each consumer below recreates
-  // it locally.
-  private lazy val stateStoreInstanceMetricsWithIds: Seq[StateStoreInstanceMetric] = {
-    val provider = StateStoreProvider.create(conf.stateStoreProviderClass)
-    val maxPartitions =
-      stateInfo.map(_.numPartitions).getOrElse(conf.defaultNumShufflePartitions)
-    (0 until maxPartitions).flatMap { partitionId =>
-      provider.supportedInstanceMetrics.flatMap { metric =>
-        stateStoreNames.map(metric.withNewId(partitionId, _))
-      }
-    }
-  }
-
   // Names of customMetrics entries treated as snapshots; preserved by
   // StateOperatorProgress.copyForNoExecution() on no-data trigger events. Includes
-  // provider- and operator-level metrics with isSnapshot = true, and all instance
-  // metric names (instance metrics use sentinel inits like -1 with monotonic
-  // combine, so they are always snapshot-style).
+  // provider- and operator-level metrics with isSnapshot = true.
   private lazy val snapshotCustomMetricNames: Set[String] = {
     val provider = StateStoreProvider.create(conf.stateStoreProviderClass)
     val customSnapshots = provider.supportedCustomMetrics.collect {
@@ -511,13 +545,7 @@ trait StateStoreWriter
     val operatorSnapshots = customStatefulOperatorMetrics.collect {
       case m if m.isSnapshot => m.name
     }.toSet
-    customSnapshots ++ operatorSnapshots ++ stateStoreInstanceMetricsWithIds.map(_.name).toSet
-  }
-
-  private def stateStoreInstanceMetrics: Map[StateStoreInstanceMetric, SQLMetric] = {
-    stateStoreInstanceMetricsWithIds.map { metric =>
-      (metric, metric.createSQLMetric(sparkContext))
-    }.toMap
+    customSnapshots ++ operatorSnapshots
   }
 
   /**

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

@@ -655,7 +655,7 @@ case class StateStoreCustomTimingMetric(name: String, desc: String) extends Stat
     SQLMetrics.createTimingMetric(sparkContext, desc)
 }
 
-trait StateStoreInstanceMetric {
+trait StateStoreInstanceMetric extends Serializable {
   def metricPrefix: String
   def descPrefix: String
   def partitionId: Option[Int]
@@ -681,6 +681,14 @@ trait StateStoreInstanceMetric {
    */
   def combine(originalMetric: SQLMetric, value: Long): Long
 
+  def combine(originalValue: Long, value: Long): Long = {
+    if (originalValue == initValue) {
+      value
+    } else {
+      Math.max(originalValue, value)
+    }
+  }
+
   def name: String = {
     assert(partitionId.isDefined, "Partition ID must be defined for instance metric name")
     s"$metricPrefix.partition_${partitionId.get}_$storeName"
@@ -724,7 +732,15 @@ case class StateStoreSnapshotLastUploadInstanceMetric(
     } else {
       // Use max to grab the most recent snapshot version across all executors
       // of the same store instance
-      Math.max(originalMetric.value, value)
+      combine(originalMetric.value, value)
+    }
+  }
+
+  override def combine(originalValue: Long, value: Long): Long = {
+    if (originalValue == initValue) {
+      value
+    } else {
+      Math.max(originalValue, value)
     }
   }
 

sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala

@@ -43,7 +43,7 @@ import org.apache.spark.util.AccumulatorV2
 class PartitionKeyedAccumulator[T] extends AccumulatorV2[(Int, T), java.util.Map[Int, T]] {
 
   // partition id -> value.
-  private val byPartition = new ConcurrentHashMap[Int, T]()
+  protected val byPartition = new ConcurrentHashMap[Int, T]()
 
   override def isZero: Boolean = byPartition.isEmpty
 

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala

@@ -18,8 +18,11 @@
 package org.apache.spark.sql.execution.streaming.state
 
 import scala.concurrent.duration.DurationInt
-import scala.jdk.CollectionConverters.MapHasAsScala
+import scala.jdk.CollectionConverters._
 
+import org.apache.spark.sql.execution.streaming.operators.stateful.{
+  StateStoreInstanceMetricAccumulator, StateStoreWriter
+}
 import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
 import org.apache.spark.sql.functions.expr
 import org.apache.spark.sql.internal.SQLConf
@@ -486,6 +489,85 @@ class StateStoreInstanceMetricSuite extends StreamTest with AlsoTestWithRocksDBF
         }
       }
   }
+
+  testWithChangelogCheckpointingEnabled(
+    "SPARK-59174: StateStoreWriter uses single accumulator for instance metrics"
+  ) {
+    withSQLConf(
+      SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName,
+      SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.key -> "1",
+      SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> "100"
+    ) {
+      withTempDir { checkpointDir =>
+        val inputData = MemoryStream[String]
+        val result = inputData.toDS().dropDuplicates()
+
+        testStream(result, outputMode = OutputMode.Update)(
+          StartStream(checkpointLocation = checkpointDir.getCanonicalPath),
+          AddData(inputData, "a", "b", "c"),
+          ProcessAllAvailable(),
+          Execute { q =>
+            val stateOp = q.lastExecution.executedPlan.collectFirst {
+              case s: StateStoreWriter => s
+            }.get
+            // Verify the accumulator is a PartitionKeyedAccumulator (not a CollectionAccumulator),
+            // has entries for the executed partitions, and does not allocate individual
+            // per-partition SQLMetrics on the plan.
+            val accValue = stateOp.instanceMetricsAccumulator.value
+            assert(!accValue.isEmpty, "accumulator should have entries after processing data")
+            // Each partition's metrics are stored as a Map; flatten all metric keys.
+            val allMetricKeys = accValue.values().asScala.flatMap(_.keys)
+            assert(
+              allMetricKeys.forall(_.name.startsWith(SNAPSHOT_LAG_METRIC_PREFIX)),
+              s"unexpected metric keys: ${allMetricKeys.map(_.name).mkString(", ")}")
+          },
+          StopStream
+        )
+      }
+    }
+  }
+
+  test("SPARK-59174: StateStoreInstanceMetricAccumulator preserves combine semantics") {
+    val metric0 = StateStoreSnapshotLastUploadInstanceMetric(Some(0), "default")
+    val metric0Store2 = StateStoreSnapshotLastUploadInstanceMetric(Some(0), "other")
+    val metric1 = StateStoreSnapshotLastUploadInstanceMetric(Some(1), "default")
+
+    // 1. Add updates to the same partition: commutative combine (max version wins)
+    val acc1 = new StateStoreInstanceMetricAccumulator
+    acc1.add((0, Map(metric0 -> 100L)))
+    acc1.add((0, Map(metric0 -> 105L)))
+    assert(acc1.value.get(0).get(metric0) === Some(105L))
+
+    val acc2 = new StateStoreInstanceMetricAccumulator
+    acc2.add((0, Map(metric0 -> 105L)))
+    acc2.add((0, Map(metric0 -> 100L)))
+    assert(acc2.value.get(0).get(metric0) === Some(105L))
+
+    // Initial value (-1) does not overwrite an existing valid snapshot version
+    acc1.add((0, Map(metric0 -> -1L)))
+    assert(acc1.value.get(0).get(metric0) === Some(105L))
+
+    // 2. Multiple stores within the same partition merge cleanly
+    acc1.add((0, Map(metric0Store2 -> 50L)))
+    assert(acc1.value.get(0).size == 2)
+    assert(acc1.value.get(0).get(metric0) === Some(105L))
+    assert(acc1.value.get(0).get(metric0Store2) === Some(50L))
+
+    // 3. Merge between accumulators: preserves combine semantics across attempts/retries
+    val accA = new StateStoreInstanceMetricAccumulator
+    accA.add((0, Map(metric0 -> 100L)))
+    accA.add((1, Map(metric1 -> 200L)))
+
+    val accB = new StateStoreInstanceMetricAccumulator
+    accB.add((0, Map(metric0 -> 105L)))
+    accB.add((1, Map(metric1 -> 150L)))
+
+    accA.merge(accB)
+    assert(accA.accumulatedNumPartitions == 2)
+    assert(accA.value.get(0).get(metric0) === Some(105L))
+    assert(accA.value.get(1).get(metric1) === Some(200L))
+  }
 }
 
 /**