[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?

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

micheal-o · 2025-07-29T00:10:13Z

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

+    saveCheckpointMetrics =
+      saveCheckpointMetrics.copy(
+        // Round up to 1ms to reassure that we've logged successfully and avoid flaky tests
+        saveZipFilesTimeMs = Some(Math.max(zipFilesTimeMs, 1L))


...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

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 2 commits July 30, 2025 05:25

PR fixes

76c2c35

Merge branch 'apache:master' into SPARK-52968

5fab7d4

[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!

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!

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!

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!

Uh oh!

dylanwong250 commented Jul 28, 2025 •

edited

Loading