[SS][SPARK-52637] Fix version ID mismatch issue for RocksDB compaction leading to incorrect file mapping

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

@@ -1705,6 +1705,24 @@ class RocksDBFileMapping {
     }.getOrElse(None)
   }
 
+  /**
+   * Remove all local file mappings that are incompatible with the current version we are
+   * trying to load.
+   *
+   * @return seq of purged mappings
+   */
+  def purgeIncompatibleMappingsForLoad(versionToLoad: Long):
+  Seq[(String, (Long, RocksDBImmutableFile))] = {
+    val filesToRemove = localFileMappings.filter {
+      case (_, (dfsFileMappedVersion, _)) =>
+        dfsFileMappedVersion >= versionToLoad
+    }.toSeq
+    filesToRemove.foreach { case (localFileName, _) =>
+      remove(localFileName)
+    }
+    filesToRemove
+  }
+
   def mapToDfsFile(
       localFileName: String,
       dfsFile: RocksDBImmutableFile,

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

@@ -38,6 +38,7 @@ import org.json4s.jackson.Serialization
 
 import org.apache.spark.{SparkConf, SparkEnv, SparkException}
 import org.apache.spark.internal.{Logging, LogKeys, MDC, MessageWithContext}
+import org.apache.spark.internal.LogKeys.{DFS_FILE, VERSION_NUM}
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.sql.errors.QueryExecutionErrors
 import org.apache.spark.sql.execution.streaming.CheckpointFileManager
@@ -785,6 +786,17 @@ class RocksDBFileManager(
       }
     }
 
+    // Delete remaining unnecessary local immutable file mappings.
+    // Files present in the file mapping but not the filesystem may lead to
+    // versionID mismatch error (SPARK-52637), so we should explicitly delete
+    // them.
+    rocksDBFileMapping.purgeIncompatibleMappingsForLoad(version).foreach {
+      case (_, (dfsFileMappedVersion, dfsFile)) =>
+        logInfo(log"Deleted local fileMapping to ${MDC(DFS_FILE, dfsFile)} because " +
+          log"mapped file version ${MDC(VERSION_NUM, dfsFileMappedVersion)} was " +
+          log"incompatible with versionToLoad ${MDC(VERSION_NUM, version)}")
+    }
+
     var filesCopied = 0L
     var bytesCopied = 0L
     var filesReused = 0L

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

@@ -3566,6 +3566,55 @@ class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession
     }
   }
 
+  test("SPARK-52637: RocksDB compaction leading to incorrect file mapping during load " +
+    "does not lead to versionID mismatch") {
+    val sqlConf = new SQLConf
+    sqlConf.setConf(
+      SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT,
+      1)
+    val dbConf = RocksDBConf(StateStoreConf(sqlConf))
+
+    withTempDir { remoteDir => withTempDir { localDir =>
+      withDB(remoteDir.toString, localDir = localDir, conf = dbConf) { db =>
+        db.load(0)
+        db.commit()
+
+        val workingDir = localDir.listFiles().filter(_.getName.startsWith("workingDir")).head
+
+        logInfo(s"files: ${db.fileManager.listRocksDBFiles(workingDir)}")
+
+        db.load(1)
+        db.put("0", "0")
+        db.commit()
+
+        db.doMaintenance() // upload snapshot to remoteDir
+
+        // confirm that sst files exist
+        assert(db.fileManager.listRocksDBFiles(workingDir)._1.nonEmpty)
+        db.fileManager.listRocksDBFiles(workingDir)._1
+          .foreach(file => file.delete()) // simulate rocksdb compaction by removing SST files
+
+        // confirm that there are entries in the mapping
+        val fileMapping = PrivateMethod[RocksDBFileMapping](Symbol("rocksDBFileMapping"))
+        val localFileMappings = PrivateMethod[mutable.Map[String, (Long, RocksDBImmutableFile)]](
+          Symbol("localFileMappings"))
+        val fileMappingObj = db invokePrivate fileMapping()
+        val localFileMappingsObj = fileMappingObj invokePrivate localFileMappings()
+        assert(localFileMappingsObj.exists { case (_, (version, _)) =>
+          version >= 1
+        })
+
+        // reload version 1
+        db.load(1)
+
+        // ensure that there are no leftover fileMappings from the first load of version 1
+        assert(!localFileMappingsObj.exists { case (_, (version, _)) =>
+          version >= 1
+        })
+      }
+    }}
+  }
+
   private def assertAcquiredThreadIsCurrentThread(db: RocksDB): Unit = {
     val threadInfo = db.getAcquiredThreadInfo()
     assert(threadInfo != None,