[SPARK-53333][SS] Enable StateDataSource with state checkpoint v2 (only readChangeFeed)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala

@@ -371,7 +371,8 @@ case class StateSourceOptions(
     stateVarName: Option[String],
     readRegisteredTimers: Boolean,
     flattenCollectionTypes: Boolean,
-    operatorStateUniqueIds: Option[Array[Array[String]]] = None) {
+    startOperatorStateUniqueIds: Option[Array[Array[String]]] = None,
+    endOperatorStateUniqueIds: Option[Array[Array[String]]] = None) {
   def stateCheckpointLocation: Path = new Path(resolvedCpLocation, DIR_NAME_STATE)
 
   override def toString: String = {
@@ -576,29 +577,46 @@ object StateSourceOptions extends DataSourceOptions {
       batchId.get
     }
 
-    val operatorStateUniqueIds = getOperatorStateUniqueIds(
+    val endBatchId = if (readChangeFeedOptions.isDefined) {
+      readChangeFeedOptions.get.changeEndBatchId
+    } else {
+      batchId.get
+    }
+
+    val startOperatorStateUniqueIds = getOperatorStateUniqueIds(
       sparkSession,
       startBatchId,
       operatorId,
       resolvedCpLocation)
 
-    if (operatorStateUniqueIds.isDefined) {
+    val endOperatorStateUniqueIds = if (startBatchId == endBatchId) {
+      startOperatorStateUniqueIds
+    } else {
+      getOperatorStateUniqueIds(
+        sparkSession,
+        endBatchId,
+        operatorId,
+        resolvedCpLocation)
+    }
+
+    if (startOperatorStateUniqueIds.isDefined != endOperatorStateUniqueIds.isDefined) {
+      throw StateDataSourceErrors.internalError(
+        "Reading source across different checkpoint format versions is not supported.")
+    }
+
+    if (startOperatorStateUniqueIds.isDefined) {
       if (fromSnapshotOptions.isDefined) {
         throw StateDataSourceErrors.invalidOptionValue(
           SNAPSHOT_START_BATCH_ID,
           "Snapshot reading is currently not supported with checkpoint v2.")
       }
-      if (readChangeFeedOptions.isDefined) {
-        throw StateDataSourceErrors.invalidOptionValue(
-          READ_CHANGE_FEED,
-          "Read change feed is currently not supported with checkpoint v2.")
-      }
     }
 
     StateSourceOptions(
       resolvedCpLocation, batchId.get, operatorId, storeName, joinSide,
       readChangeFeed, fromSnapshotOptions, readChangeFeedOptions,
-      stateVarName, readRegisteredTimers, flattenCollectionTypes, operatorStateUniqueIds)
+      stateVarName, readRegisteredTimers, flattenCollectionTypes,
+      startOperatorStateUniqueIds, endOperatorStateUniqueIds)
   }
 
   private def resolvedCheckpointLocation(

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala

@@ -96,11 +96,20 @@ abstract class StatePartitionReaderBase(
       schema, "value").asInstanceOf[StructType]
   }
 
-  protected val getStoreUniqueId : Option[String] = {
+  protected def getStoreUniqueId(
+    operatorStateUniqueIds: Option[Array[Array[String]]]) : Option[String] = {
     SymmetricHashJoinStateManager.getStateStoreCheckpointId(
       storeName = partition.sourceOptions.storeName,
       partitionId = partition.partition,
-      stateStoreCkptIds = partition.sourceOptions.operatorStateUniqueIds)
+      stateStoreCkptIds = operatorStateUniqueIds)
+  }
+
+  protected def getStartStoreUniqueId: Option[String] = {
+    getStoreUniqueId(partition.sourceOptions.startOperatorStateUniqueIds)
+  }
+
+  protected def getEndStoreUniqueId: Option[String] = {
+    getStoreUniqueId(partition.sourceOptions.endOperatorStateUniqueIds)
   }
 
   protected lazy val provider: StateStoreProvider = {
@@ -123,7 +132,7 @@ abstract class StatePartitionReaderBase(
     if (useColFamilies) {
       val store = provider.getStore(
         partition.sourceOptions.batchId + 1,
-        getStoreUniqueId)
+        getEndStoreUniqueId)
       require(stateStoreColFamilySchemaOpt.isDefined)
       val stateStoreColFamilySchema = stateStoreColFamilySchemaOpt.get
       require(stateStoreColFamilySchema.keyStateEncoderSpec.isDefined)
@@ -182,9 +191,11 @@ class StatePartitionReader(
   private lazy val store: ReadStateStore = {
     partition.sourceOptions.fromSnapshotOptions match {
       case None =>
+        assert(getStartStoreUniqueId == getEndStoreUniqueId,
+          "Start and end store unique IDs must be the same when not reading from snapshot")
         provider.getReadStore(
           partition.sourceOptions.batchId + 1,
-          getStoreUniqueId
+          getStartStoreUniqueId
         )
 
       case Some(fromSnapshotOptions) =>
@@ -261,7 +272,8 @@ class StateStoreChangeDataPartitionReader(
       .getStateStoreChangeDataReader(
         partition.sourceOptions.readChangeFeedOptions.get.changeStartBatchId + 1,
         partition.sourceOptions.readChangeFeedOptions.get.changeEndBatchId + 1,
-        colFamilyNameOpt)
+        colFamilyNameOpt,
+        getEndStoreUniqueId)
   }
 
   override lazy val iter: Iterator[InternalRow] = {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StreamStreamJoinStatePartitionReader.scala

@@ -76,21 +76,22 @@ class StreamStreamJoinStatePartitionReader(
     partition.sourceOptions.stateCheckpointLocation.toString,
     partition.sourceOptions.operatorId)
 
-  private val stateStoreCheckpointIds = SymmetricHashJoinStateManager.getStateStoreCheckpointIds(
+  private val startStateStoreCheckpointIds =
+    SymmetricHashJoinStateManager.getStateStoreCheckpointIds(
     partition.partition,
-    partition.sourceOptions.operatorStateUniqueIds,
+    partition.sourceOptions.startOperatorStateUniqueIds,
     usesVirtualColumnFamilies)
 
   private val keyToNumValuesStateStoreCkptId = if (joinSide == LeftSide) {
-    stateStoreCheckpointIds.left.keyToNumValues
+    startStateStoreCheckpointIds.left.keyToNumValues
   } else {
-    stateStoreCheckpointIds.right.keyToNumValues
+    startStateStoreCheckpointIds.right.keyToNumValues
   }
 
   private val keyWithIndexToValueStateStoreCkptId = if (joinSide == LeftSide) {
-    stateStoreCheckpointIds.left.keyWithIndexToValue
+    startStateStoreCheckpointIds.left.keyWithIndexToValue
   } else {
-    stateStoreCheckpointIds.right.keyWithIndexToValue
+    startStateStoreCheckpointIds.right.keyWithIndexToValue
   }
 
   /*

sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala

@@ -389,7 +389,7 @@ case class TransformWithStateInPySparkExec(
           store.abort()
         }
       }
-      setStoreMetrics(store)
+      setStoreMetrics(store, isStreaming)
       setOperatorMetrics()
     }).map { row =>
       numOutputRows += 1

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala

@@ -430,14 +430,14 @@ trait StateStoreWriter
    * Set the SQL metrics related to the state store.
    * This should be called in that task after the store has been updated.
    */
-  protected def setStoreMetrics(store: StateStore): Unit = {
+  protected def setStoreMetrics(store: StateStore, setCheckpointInfo: Boolean = true): Unit = {
     val storeMetrics = store.metrics
     longMetric("numTotalStateRows") += storeMetrics.numKeys
     longMetric("stateMemory") += storeMetrics.memoryUsedBytes
     setStoreCustomMetrics(storeMetrics.customMetrics)
     setStoreInstanceMetrics(storeMetrics.instanceMetrics)
 
-    if (StatefulOperatorStateInfo.enableStateStoreCheckpointIds(conf)) {
+    if (StatefulOperatorStateInfo.enableStateStoreCheckpointIds(conf) && setCheckpointInfo) {
       // Set the state store checkpoint information for the driver to collect
       val ssInfo = store.getStateStoreCheckpointInfo()
       setStateStoreCheckpointInfo(

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExec.scala

@@ -346,7 +346,7 @@ case class TransformWithStateExec(
           store.abort()
         }
       }
-      setStoreMetrics(store)
+      setStoreMetrics(store, isStreaming)
       setOperatorMetrics()
       closeStatefulProcessor()
       statefulProcessor.setHandle(null)

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

@@ -1064,8 +1064,15 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
   override def getStateStoreChangeDataReader(
       startVersion: Long,
       endVersion: Long,
-      colFamilyNameOpt: Option[String] = None):
+      colFamilyNameOpt: Option[String] = None,
+      endVersionStateStoreCkptId: Option[String] = None):
     StateStoreChangeDataReader = {
+
+    if (endVersionStateStoreCkptId.isDefined) {
+      throw QueryExecutionErrors.cannotLoadStore(new SparkException(
+        "HDFSBackedStateStoreProvider does not support endVersionStateStoreCkptId"))
+    }
+
     // Multiple column families are not supported with HDFSBackedStateStoreProvider
     if (colFamilyNameOpt.isDefined) {
       throw StateStoreErrors.multipleColumnFamiliesNotSupported(providerName)
@@ -1099,7 +1106,7 @@ class HDFSBackedStateStoreChangeDataReader(
   extends StateStoreChangeDataReader(
     fm, stateLocation, startVersion, endVersion, compressionCodec) {
 
-  override protected var changelogSuffix: String = "delta"
+  override protected val changelogSuffix: String = "delta"
 
   override def getNext(): (RecordType.Value, UnsafeRow, UnsafeRow, Long) = {
     val reader = currentChangelogReader()

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

@@ -342,6 +342,51 @@ class RocksDB(
     currLineage
   }
 
+  /**
+   * Construct the full lineage from startVersion to endVersion (inclusive) by
+   * walking backwards using lineage information embedded in changelog files.
+   */
+  def getFullLineage(
+      startVersion: Long,
+      endVersion: Long,
+      endVersionStateStoreCkptId: Option[String]): Array[LineageItem] = {
+    assert(startVersion <= endVersion,
+      s"startVersion $startVersion should be less than or equal to endVersion $endVersion")
+
+    // A buffer to collect the lineage information, the entries should be decreasing in version
+    val buf = mutable.ArrayBuffer[LineageItem]()
+    buf.append(LineageItem(endVersion, endVersionStateStoreCkptId.get))
+
+    while (buf.last.version > startVersion) {
+      val prevSmallestVersion = buf.last.version
+      val lineage = getLineageFromChangelogFile(buf.last.version, Some(buf.last.checkpointUniqueId))
+      // lineage array is sorted in increasing order, we need to reverse it
+      val lineageSorted = lineage.filter(_.version >= startVersion).sortBy(_.version).reverse
+      // append to the buffer in reverse order, so the buffer is always decreasing in version
+      buf.appendAll(lineageSorted)
+
+      // to prevent infinite loop if we make no progress, throw an exception
+      if (buf.last.version == prevSmallestVersion) {
+        throw new IllegalStateException(s"Lineage is not complete")
+      }
+    }
+
+    // we return the lineage in increasing order
+    val ret = buf.reverse.toArray
+
+    // Sanity checks
+    assert(ret.head.version == startVersion,
+      s"Expected first lineage version to be $startVersion, but got ${ret.head.version}")
+    assert(ret.last.version == endVersion,
+      s"Expected last lineage version to be $endVersion, but got ${ret.last.version}")
+    // Assert that the lineage array is strictly increasing in version
+    assert(ret.sliding(2).forall {
+      case Array(prev, next) => prev.version + 1 == next.version
+      case _ => true
+    }, s"Lineage array is not strictly increasing in version")
+
+    ret
+  }
 
   /**
    * Load the given version of data in a native RocksDB instance.

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

@@ -877,15 +877,18 @@ private[sql] class RocksDBStateStoreProvider
   override def getStateStoreChangeDataReader(
       startVersion: Long,
       endVersion: Long,
-      colFamilyNameOpt: Option[String] = None):
+      colFamilyNameOpt: Option[String] = None,
+      endVersionStateStoreCkptId: Option[String] = None):
     StateStoreChangeDataReader = {
     val statePath = stateStoreId.storeCheckpointLocation()
     val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf)
     new RocksDBStateStoreChangeDataReader(
       CheckpointFileManager.create(statePath, hadoopConf),
+      rocksDB,
       statePath,
       startVersion,
       endVersion,
+      endVersionStateStoreCkptId,
       CompressionCodec.createCodec(sparkConf, storeConf.compressionCodec),
       keyValueEncoderMap,
       colFamilyNameOpt)
@@ -1224,17 +1227,32 @@ object RocksDBStateStoreProvider {
 /** [[StateStoreChangeDataReader]] implementation for [[RocksDBStateStoreProvider]] */
 class RocksDBStateStoreChangeDataReader(
     fm: CheckpointFileManager,
+    rocksDB: RocksDB,
     stateLocation: Path,
     startVersion: Long,
     endVersion: Long,
+    endVersionStateStoreCkptId: Option[String],
     compressionCodec: CompressionCodec,
     keyValueEncoderMap:
       ConcurrentHashMap[String, (RocksDBKeyStateEncoder, RocksDBValueStateEncoder, Short)],
     colFamilyNameOpt: Option[String] = None)
   extends StateStoreChangeDataReader(
     fm, stateLocation, startVersion, endVersion, compressionCodec, colFamilyNameOpt) {
 
-  override protected var changelogSuffix: String = "changelog"
+  override protected val versionsAndUniqueIds: Array[(Long, Option[String])] =
+    if (endVersionStateStoreCkptId.isDefined) {
+      val fullVersionLineage = rocksDB.getFullLineage(
+        startVersion,
+        endVersion,
+        endVersionStateStoreCkptId)
+      fullVersionLineage
+        .sortBy(_.version)
+        .map(item => (item.version, Some(item.checkpointUniqueId)))
+    } else {
+      (startVersion to endVersion).map((_, None)).toArray
+    }
+
+  override protected val changelogSuffix: String = "changelog"
 
   override def getNext(): (RecordType.Value, UnsafeRow, UnsafeRow, Long) = {
     var currRecord: (RecordType.Value, Array[Byte], Array[Byte]) = null

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

@@ -834,7 +834,8 @@ trait SupportsFineGrainedReplay {
   def getStateStoreChangeDataReader(
       startVersion: Long,
       endVersion: Long,
-      colFamilyNameOpt: Option[String] = None):
+      colFamilyNameOpt: Option[String] = None,
+      endVersionStateStoreCkptId: Option[String] = None):
     NextIterator[(RecordType.Value, UnsafeRow, UnsafeRow, Long)]
 }
 

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

@@ -447,6 +447,7 @@ abstract class StateStoreChangelogReader(
     Serialization.read[Array[LineageItem]](lineageStr)
   }
 
+  // The array contains lineage information from [snapShotVersion, version]
   lazy val lineage: Array[LineageItem] = readLineage()
 
   def version: Short
@@ -632,27 +633,41 @@ abstract class StateStoreChangeDataReader(
    * Iterator that iterates over the changelog files in the state store.
    */
   private class ChangeLogFileIterator extends Iterator[Path] {
+    val versionsAndUniqueIds: Iterator[(Long, Option[String])] =
+      StateStoreChangeDataReader.this.versionsAndUniqueIds.iterator
 
     private var currentVersion = StateStoreChangeDataReader.this.startVersion - 1
+    private var currentUniqueId: Option[String] = None
 
     /** returns the version of the changelog returned by the latest [[next]] function call */
     def getVersion: Long = currentVersion
 
-    override def hasNext: Boolean = currentVersion < StateStoreChangeDataReader.this.endVersion
+    override def hasNext: Boolean = versionsAndUniqueIds.hasNext
 
     override def next(): Path = {
-      currentVersion += 1
-      getChangelogPath(currentVersion)
+      val nextTuple = versionsAndUniqueIds.next()
+      currentVersion = nextTuple._1
+      currentUniqueId = nextTuple._2
+      getChangelogPath(currentVersion, currentUniqueId)
     }
 
-    private def getChangelogPath(version: Long): Path =
-      new Path(
-        StateStoreChangeDataReader.this.stateLocation,
-        s"$version.${StateStoreChangeDataReader.this.changelogSuffix}")
+    private def getChangelogPath(version: Long, checkpointUniqueId: Option[String]): Path =
+      if (checkpointUniqueId.isDefined) {
+        new Path(
+          StateStoreChangeDataReader.this.stateLocation,
+          s"${version}_${checkpointUniqueId.get}." +
+            s"${StateStoreChangeDataReader.this.changelogSuffix}")
+      } else {
+        new Path(
+          StateStoreChangeDataReader.this.stateLocation,
+          s"$version.${StateStoreChangeDataReader.this.changelogSuffix}")
+      }
   }
 
   /** file format of the changelog files */
-  protected var changelogSuffix: String
+  protected val changelogSuffix: String
+  protected val versionsAndUniqueIds: Array[(Long, Option[String])] =
+    (startVersion to endVersion).map((_, None)).toArray
   private lazy val fileIterator = new ChangeLogFileIterator
   private var changelogReader: StateStoreChangelogReader = null
 
@@ -671,11 +686,10 @@ abstract class StateStoreChangeDataReader(
         return null
       }
 
-      changelogReader = if (colFamilyNameOpt.isDefined) {
-        new StateStoreChangelogReaderV2(fm, fileIterator.next(), compressionCodec)
-      } else {
-        new StateStoreChangelogReaderV1(fm, fileIterator.next(), compressionCodec)
-      }
+      val changelogFile = fileIterator.next()
+      changelogReader =
+        new StateStoreChangelogReaderFactory(fm, changelogFile, compressionCodec)
+          .constructChangelogReader()
     }
     changelogReader
   }