[SPARK-52968][SS] Emit additional state store metrics #51679

dylanwong250 · 2025-07-28T02:16:07Z

What changes were proposed in this pull request?

Add additional metrics in structured streaming:

State Store Commit Metrics

"rocksdbChangeLogWriterCommitLatencyMs"

"rocksdbSaveZipFilesLatencyMs"

State Store Load Metrics

"rocksdbLoadFromSnapshotLatencyMs"

"rocksdbLoadLatencyMs"

"rocksdbReplayChangeLogLatencyMs"

"rocksdbNumReplayChangelogFiles"

Why are the changes needed?

Currently there are no metrics emitted related to loading the state store and replaying the change log.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Added unit tests to verify that the new metrics are populated and that query progress contains the correct metrics.

Was this patch authored or co-authored using generative AI tooling?

No

liviazhu · 2025-07-28T16:58:58Z

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

+    )
+
+    // Refresh the recorded metrics after loading from snapshot
+    recordedMetrics = Some(metrics)


Why do we set recordedMetrics here but not in load?

I missed that load can be read only. For more context, metricsOpt is how the RocksDBMetrics is accessed. It uses recordedMetrics which was previously only updated in commit(). This caused metrics from loadFromSnapshot to not be up to date since it is currently only used as read-only which does not commit().

I added the metrics refresh in load().

@dylanwong250 please add comment to say why you added this

I added a comment in both load() and loadFromSnapshot().

When I added recordedMetrics = Some(metrics) in load() it caused some tests to fail. This was due to cases when abort() was called and then after that metrics were expected to be empty. I fixed this by clearing the metrics in rollback. I will see if that fixes the tests.

@dylanwong250 Why do we want the metrics to be empty after rollback?

I was running into a test case where in TransformWithStateExec there is the statement:

if (isStreaming) { store.commit() } else { store.abort() }

The abort path would be taken and then the metrics would be set. But doing the set I was getting NPEs because the metrics did not exist in the metrics Map. I think this has something to do with how metrics are being initialized when isStreaming = false for TransformWithStateExec. The test that was failing was testTransformWithState.

I discussed with Micheal to remove the refresh of the metrics on load() since these metrics are only currently reported on commit and the additional calls could increase latency of loading the store.

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

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

liviazhu · 2025-07-28T17:35:33Z

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

@@ -1350,6 +1392,7 @@ class RocksDB(
      totalSSTFilesBytes,
      nativeOpsLatencyMicros,
      commitLatencyMs,
+      loadMetrics,


Is there any concern about race conditions if we're not creating a copy of the maps here? I know that this is a pre-existing pattern, but could you just verify with a quick test that if we get the metrics, reload the state store (resetting the maps), and then read the metrics, that the metrics are correct?

Good catch. It seems that the existing commitLatencyMs also has the issue where the metrics in the maps from the call to metricsOpt can change. I added a clone for these maps and added a test.

micheal-o · 2025-07-28T22:41:16Z

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

+    )
+
+    // Refresh the recorded metrics after loading from snapshot
+    recordedMetrics = Some(metrics)


why are you doing this? This pulls the rocksdb engine metrics right

Response in the other comment thread related to this.

micheal-o · 2025-07-28T22:46:13Z

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

+    )
+
+    // Refresh the recorded metrics after loading from snapshot
+    recordedMetrics = Some(metrics)


@dylanwong250 please add comment to say why you added this

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

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

...rc/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala

micheal-o · 2025-07-29T00:33:24Z

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

@@ -1101,7 +1144,8 @@ class RocksDB(
      numInternalKeysOnLoadedVersion = numInternalKeysOnWritingVersion
      loadedVersion = newVersion
      commitLatencyMs ++= Map(
-        "fileSync" -> fileSyncTimeMs
+        "fileSync" -> fileSyncTimeMs,
+        "saveZipFiles" -> fileManagerMetrics.saveZipFilesTimeMs.getOrElse(0L)


there should be a better way to do this. Also even if there was no zip upload in this commit, this would return the metrics from when maintenance thread did upload. That can confuse the user or eng debugging

Changed. I added a flag to only set this metric in commit when uploadSnapshot has been called which clears and recreates the fileManagerMetrics.

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

liviazhu · 2025-07-29T22:16:31Z

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

+    )
+
+    // Refresh the recorded metrics after loading from snapshot
+    recordedMetrics = Some(metrics)


@dylanwong250 Why do we want the metrics to be empty after rollback?

liviazhu · 2025-07-29T22:19:55Z

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

@@ -1312,7 +1376,8 @@ class RocksDB(
      pinnedBlocksMemUsage,
      totalSSTFilesBytes,
      nativeOpsLatencyMicros,
-      commitLatencyMs,
+      commitLatencyMs.clone(),


Thanks for making this change, can you just add a quick comment explaining why we need to clone?

liviazhu · 2025-07-29T22:21:07Z

Change looks good mostly. Besides nits, just had 1 question about rollback

micheal-o · 2025-07-30T17:41:23Z

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

@@ -1034,14 +1078,22 @@ class RocksDB(
          if (shouldForceSnapshot.get()) {
            assert(snapshot.isDefined)
            uploadSnapshot(snapshot.get)
+            uploadSnapshotCalled = true


no need to introduce new flag here. We already have the isUploaded flag here to detect if we uploaded snapshot

Changed to use a single flag.

micheal-o

Thanks for making this change. Stamped with one last comment. PTAL

micheal-o · 2025-07-31T00:03:47Z

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

+  /**
+   * Refresh the recorded metrics with the latest metrics.
+   */
+  def refreshRecordedMetrics(): Unit = {


is this for test only? if so make it private[state] and name refreshRecordedMetricsForTest

micheal-o · 2025-07-31T19:32:06Z

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

@@ -156,6 +156,8 @@ class RocksDB(
  private val byteArrayPair = new ByteArrayPair()
  private val commitLatencyMs = new mutable.HashMap[String, Long]()


@dylanwong250 please can we reset this commitLatencyMs map too just like you do for loadMetrics. We currently don't reset it and might be returning metrics from previous batch (existing bug). Thanks.

Emit additional state store metrics

4018608

github-actions bot added SQL STRUCTURED STREAMING labels Jul 28, 2025

liviazhu reviewed Jul 28, 2025

View reviewed changes

Dylan Wong and others added 3 commits July 28, 2025 20:26

PR fixes

d9cdeaf

Fix

86dc65c

Merge branch 'apache:master' into SPARK-52968

4878bf0

micheal-o reviewed Jul 29, 2025

View reviewed changes

Dylan Wong added 4 commits July 29, 2025 16:34

Some PR fixes

97bf9ac

Merge branch 'master' into SPARK-52968

dc51917

Fix flakey test

c4f256b

Only snapshot commit metrics when there is a upload in the commit

3b99287

dylanwong250 requested review from micheal-o and liviazhu July 29, 2025 21:25

liviazhu reviewed Jul 29, 2025

View reviewed changes

Dylan Wong and others added 3 commits July 30, 2025 05:25

PR fixes

76c2c35

Merge branch 'apache:master' into SPARK-52968

5fab7d4

Remove clearing metrics

f9d76b4

micheal-o reviewed Jul 30, 2025

View reviewed changes

Remove metrics refresh

72ba026

dylanwong250 requested review from liviazhu and micheal-o July 30, 2025 22:20

micheal-o approved these changes Jul 31, 2025

View reviewed changes

Fix test method

a91d851

liviazhu approved these changes Jul 31, 2025

View reviewed changes

micheal-o reviewed Jul 31, 2025

View reviewed changes

Dylan Wong added 3 commits July 31, 2025 22:53

Fix commit metrics bug

8590307

style fix

913b6c1

Merge branch 'master' into SPARK-52968

bae64c4

		@@ -156,6 +156,8 @@ class RocksDB(
		private val byteArrayPair = new ByteArrayPair()
		private val commitLatencyMs = new mutable.HashMap[String, Long]()

[SPARK-52968][SS] Emit additional state store metrics #51679

Are you sure you want to change the base?

[SPARK-52968][SS] Emit additional state store metrics #51679

Conversation

dylanwong250 commented Jul 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dylanwong250 Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liviazhu commented Jul 29, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

micheal-o left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dylanwong250 commented Jul 28, 2025 •

edited

Loading

dylanwong250 Jul 30, 2025 •

edited

Loading