diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala index 799f6326f7fd3..acf9ed544e55c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala +++ b/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 } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala index 9c969efe93191..ad942d718b84e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala +++ b/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) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala b/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala index f56a6f882eda0..8cac9562d62bf 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala +++ b/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 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala index 2015b7eaed1ac..4fd8425530f27 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala +++ b/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)) + } } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala index 17aa7de7f86d5..a235ff1067531 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala @@ -18,6 +18,10 @@ package org.apache.spark.sql.util import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.execution.streaming.operators.stateful.{ + StateStoreInstanceMetricAccumulator +} +import org.apache.spark.sql.execution.streaming.state.StateStoreSnapshotLastUploadInstanceMetric class PartitionKeyedAccumulatorSuite extends SparkFunSuite { @@ -126,4 +130,45 @@ class PartitionKeyedAccumulatorSuite extends SparkFunSuite { (rows + partitionRows, bytes + partitionBytes) }.contains((17L, 170L))) } + + 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)) + } }