(spark) branch branch-4.x updated: [SPARK-57658][SS] Report available snapshot versions when a RocksDB state snapshot load fails

[gu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new 83784d4e64c3 [SPARK-57658][SS] Report available snapshot versions when a RocksDB state snapshot load fails
83784d4e64c3 is described below

commit 83784d4e64c3eac9ca3a911686bbf44d9cb8e188
Author: Hyukjin Kwon <gu...@apache.org>
AuthorDate: Sat Jun 27 11:10:21 2026 +0900

    [SPARK-57658][SS] Report available snapshot versions when a RocksDB state snapshot load fails
    
    ### What changes were proposed in this pull request?
    
    When loading state from a snapshot (e.g. reading with the `snapshotStartBatchId` option) the snapshot zip for the requested version can be missing — most often because the asynchronous maintenance thread has not uploaded it yet. Today that surfaces only as:
    
    ```
    [CANNOT_LOAD_STATE_STORE.UNCATEGORIZED] ...
    Caused by: java.io.FileNotFoundException: .../state/0/1/2.zip does not exist
    ```
    
    with no indication of whether the snapshot was never created or merely not uploaded in time. This PR enriches the `FileNotFoundException` thrown from `RocksDBFileManager` with the snapshot (`.zip`) and changelog (`.changelog`) files that **are** present in the DFS checkpoint root, so the situation is self-diagnosing from logs (e.g. *asked for 2.zip, only `snapshots=[1.zip]` present* clearly indicates the async-upload race). The listing is best-effort and never throws.
    
    ### Why are the changes needed?
    
    The `snapshotStartBatchId ... transformWithState` tests have failed intermittently in scheduled Maven jobs with exactly this opaque `FileNotFoundException`, in different suite variants per run:
    - master, `sql#core - slow tests`: https://github.com/apache/spark/actions/runs/28048347820 (plain `StateDataSourceTransformWithStateSuite`)
    - branch-4.2, `sql#core - slow tests`: https://github.com/apache/spark/actions/runs/28045133937/job/83029837937 (`StateDataSourceTransformWithStateSuiteWithRowChecksum`)
    
    The deterministic test-side wait was handled separately (SPARK-57652); this change makes any *future* recurrence — in any suite or scheduled job — immediately diagnosable instead of requiring artifact spelunking.
    
    ### Does this PR introduce any user-facing change?
    
    No (improves an error message on an already-failing path).
    
    ### How was this patch tested?
    
    - New unit test `RocksDBFileManager: missing snapshot during load reports the available versions` in `RocksDBSuite`, covering both the no-`checkpointUniqueId` (`fs.open`) and the `checkpointUniqueId` state-checkpoint-v2 (`fm.open`) load paths.
    - Focused fork run (RocksDBSuite once + the snapshot suites 5× to confirm stability): **passing** — https://github.com/HyukjinKwon/spark/actions/runs/28075060090/job/83117473449
    
    This pull request and its description were written by Isaac.
    
    Closes #56718 from HyukjinKwon/ci-fix/snapshot-load-diagnostics.
    
    Authored-by: Hyukjin Kwon <gu...@apache.org>
    Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
    (cherry picked from commit a57de55e6aef36ac76a0ca42c3007f32b27eb915)
    Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
---
 .../streaming/state/RocksDBFileManager.scala       | 50 +++++++++++++++++++---
 .../execution/streaming/state/RocksDBSuite.scala   | 43 +++++++++++++++++++
 2 files changed, 87 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBFileManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBFileManager.scala
index 593b2ef7951f..374cfc5ed0cd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBFileManager.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBFileManager.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
-import java.io.{File, FileInputStream, InputStream}
+import java.io.{File, FileInputStream, FileNotFoundException, InputStream}
 import java.nio.charset.StandardCharsets.UTF_8
 import java.nio.file.Files
 import java.util.concurrent.ConcurrentHashMap
@@ -25,6 +25,7 @@ import java.util.zip.{ZipEntry, ZipOutputStream}
 
 import scala.collection.{mutable, Map}
 import scala.math._
+import scala.util.control.NonFatal
 
 import com.fasterxml.jackson.annotation.JsonInclude.Include
 import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
@@ -976,11 +977,48 @@ class RocksDBFileManager(
    */
   private def unzipBatchZipFileFromDfs(
       version: Long, checkpointUniqueId: Option[String], localDir: File): Seq[File] = {
-    if (checkpointUniqueId.isDefined) {
-      Utils.unzipFilesFromInputStream(
-        fm.open(dfsBatchZipFile(version, checkpointUniqueId)), localDir)
-    } else {
-      Utils.unzipFilesFromFile(fs, dfsBatchZipFile(version, checkpointUniqueId), localDir)
+    try {
+      if (checkpointUniqueId.isDefined) {
+        Utils.unzipFilesFromInputStream(
+          fm.open(dfsBatchZipFile(version, checkpointUniqueId)), localDir)
+      } else {
+        Utils.unzipFilesFromFile(fs, dfsBatchZipFile(version, checkpointUniqueId), localDir)
+      }
+    } catch {
+      case e: FileNotFoundException =>
+        // The snapshot zip for this version is missing. The most common cause is that the
+        // asynchronous maintenance thread has not uploaded the snapshot yet (e.g. reading state
+        // with the snapshotStartBatchId option immediately after writing). Enrich the error with
+        // the snapshot/changelog files that ARE present in DFS, so the situation is diagnosable
+        // straight from the logs (in particular for intermittent failures in scheduled jobs).
+        // e.getMessage already names the missing zip path, so we do not repeat it; the original
+        // exception is preserved as the cause so its stack trace is not lost.
+        val enriched = new FileNotFoundException(
+          s"${e.getMessage}\nFailed to load the snapshot file for version $version" +
+          checkpointUniqueId.map(id => s" (checkpointUniqueId=$id)").getOrElse("") +
+          s". Files currently present in $dfsRootDir: ${listDfsFilesForDiagnostics()}")
+        enriched.initCause(e)
+        throw enriched
+    }
+  }
+
+  /**
+   * Best-effort listing of the snapshot (.zip) and changelog (.changelog) files present in the
+   * DFS checkpoint root. Used only to enrich diagnostics when a snapshot load fails; never throws.
+   */
+  private def listDfsFilesForDiagnostics(): String = {
+    try {
+      val path = new Path(dfsRootDir)
+      if (!fm.exists(path)) {
+        "<DFS checkpoint root does not exist>"
+      } else {
+        val names = fm.list(path).map(_.getPath.getName)
+        val snapshots = names.filter(_.endsWith(".zip")).sorted
+        val changelogs = names.filter(_.endsWith(".changelog")).sorted
+        s"snapshots=[${snapshots.mkString(", ")}], changelogs=[${changelogs.mkString(", ")}]"
+      }
+    } catch {
+      case NonFatal(t) => s"<failed to list DFS files for diagnostics: ${t.getMessage}>"
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
index dc697f5b99dc..df29d94ec32d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
@@ -1826,6 +1826,49 @@ class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession
     }
   }
 
+  test("RocksDBFileManager: missing snapshot during load reports the available versions") {
+    // Loading a snapshot version that has not been uploaded yet (e.g. the asynchronous
+    // maintenance thread has not finished uploading it when reading state with
+    // snapshotStartBatchId) should fail with a FileNotFoundException whose message lists the
+    // snapshot/changelog files that ARE present, so intermittent failures in scheduled jobs are
+    // diagnosable straight from the logs.
+    val hadoopConf = new Configuration()
+    val remoteDir = Utils.createTempDir().toString
+    val fileManager = new RocksDBFileManager(remoteDir, Utils.createTempDir(), hadoopConf)
+    val fileMapping = new RocksDBFileMapping()
+    // Upload only snapshot version 1, leaving version 2 absent.
+    saveCheckpointFiles(
+      fileManager, Seq("001.sst" -> 10, "002.sst" -> 20), version = 1, numKeys = 10, fileMapping)
+
+    val ex = intercept[FileNotFoundException] {
+      fileManager.loadCheckpointFromDfs(2, Utils.createTempDir(), fileMapping)
+    }
+    assert(ex.getMessage.contains("Failed to load the snapshot file for version 2"))
+    assert(ex.getMessage.contains("Files currently present"))
+    // The version-1 snapshot that does exist must be surfaced in the diagnostic.
+    assert(ex.getMessage.contains("snapshots=[1.zip]"))
+
+    // Also cover the checkpointUniqueId (state checkpoint v2) path, which loads the snapshot via
+    // fm.open instead of fs.open. The diagnostic should still fire and name the unique id.
+    val v2RemoteDir = Utils.createTempDir().toString
+    val v2FileManager = new RocksDBFileManager(v2RemoteDir, Utils.createTempDir(), hadoopConf)
+    val v2FileMapping = new RocksDBFileMapping()
+    val checkpointUniqueId = Some(java.util.UUID.randomUUID.toString)
+    saveCheckpointFiles(
+      v2FileManager, Seq("001.sst" -> 10), version = 1, numKeys = 10, v2FileMapping,
+      checkpointUniqueId = checkpointUniqueId)
+
+    val v2Ex = intercept[FileNotFoundException] {
+      v2FileManager.loadCheckpointFromDfs(
+        2, Utils.createTempDir(), v2FileMapping, checkpointUniqueId = checkpointUniqueId)
+    }
+    assert(v2Ex.getMessage.contains("Failed to load the snapshot file for version 2"))
+    assert(v2Ex.getMessage.contains(s"checkpointUniqueId=${checkpointUniqueId.get}"))
+    assert(v2Ex.getMessage.contains("Files currently present"))
+    // The version-1 snapshot that does exist (named with the unique id) must be surfaced.
+    assert(v2Ex.getMessage.contains(s"1_${checkpointUniqueId.get}.zip"))
+  }
+
   testWithChangelogCheckpointingEnabled("RocksDBFileManager: read and write changelog") {
     val dfsRootDir = new File(Utils.createTempDir().getAbsolutePath + "/state/1/1")
     val fileManager = new RocksDBFileManager(


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org