make sure StateStoreChangeData is used everywhere

eason-yuchen-liu · Jun 26, 2024 · 271b98e · 271b98e
1 parent ff5bff2
commit 271b98e
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 15 deletions.
diff --git a/...main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala b/...main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala
@@ -41,7 +41,7 @@ class StatePartitionReaderFactory(
   override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
     val stateStoreInputPartition = partition.asInstanceOf[StateStoreInputPartition]
     if (stateStoreInputPartition.sourceOptions.modeType == StateDataSourceModeType.CDC) {
-      new StateCDCPartitionReader(storeConf, hadoopConf,
+      new StateStoreChangeDataPartitionReader(storeConf, hadoopConf,
         stateStoreInputPartition, schema)
     } else {
       new StatePartitionReader(storeConf, hadoopConf,
@@ -151,16 +151,17 @@ class StatePartitionReader(
   }
 }
 
-class StateCDCPartitionReader(
+class StateStoreChangeDataPartitionReader(
   storeConf: StateStoreConf,
   hadoopConf: SerializableConfiguration,
   partition: StateStoreInputPartition,
   schema: StructType) extends StatePartitionReader(storeConf, hadoopConf, partition, schema) {
 
-  private lazy val cdcReader: StateChangeDataReader = {
-    provider.getStateChangeDataReader(
-      partition.sourceOptions.cdcStartBatchID.get + 1,
-      partition.sourceOptions.cdcEndBatchId.get + 1)
+  private lazy val cdcReader: StateStoreChangeDataReader = {
+    provider.asInstanceOf[SupportsStateStoreChangeDataFeed]
+      .getStateStoreChangeDataReader(
+        partition.sourceOptions.cdcStartBatchID.get + 1,
+        partition.sourceOptions.cdcEndBatchId.get + 1)
   }
 
   override protected lazy val iter: Iterator[InternalRow] = {

diff --git a/...n/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/...n/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -72,7 +72,7 @@ import org.apache.spark.util.ArrayImplicits._
  * store.
  */
 private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with Logging
-  with SupportsFineGrainedReplayFromSnapshot {
+  with SupportsFineGrainedReplayFromSnapshot with SupportsStateStoreChangeDataFeed {
 
   private val providerName = "HDFSBackedStateStoreProvider"
 
@@ -976,7 +976,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     result
   }
 
-  override def getStateChangeDataReader(startVersion: Long, endVersion: Long): StateChangeDataReader = {
+  override def getStateStoreChangeDataReader(startVersion: Long, endVersion: Long):
+    StateStoreChangeDataReader = {
     new HDFSBackedStateStoreCDCReader(fm, baseDir, startVersion, endVersion,
       CompressionCodec.createCodec(sparkConf, storeConf.compressionCodec),
       keySchema, valueSchema)

diff --git a/...main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/...main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
@@ -35,7 +35,7 @@ import org.apache.spark.util.Utils
 
 private[sql] class RocksDBStateStoreProvider
   extends StateStoreProvider with Logging with Closeable
-    with SupportsFineGrainedReplayFromSnapshot {
+    with SupportsFineGrainedReplayFromSnapshot with SupportsStateStoreChangeDataFeed {
   import RocksDBStateStoreProvider._
 
   class RocksDBStateStore(lastVersion: Long) extends StateStore {
@@ -403,7 +403,8 @@ private[sql] class RocksDBStateStoreProvider
     }
   }
 
-  override def getStateChangeDataReader(startVersion: Long, endVersion: Long): StateChangeDataReader = {
+  override def getStateStoreChangeDataReader(startVersion: Long, endVersion: Long):
+    StateStoreChangeDataReader = {
     val statePath = stateStoreId.storeCheckpointLocation()
     val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf)
     new RocksDBStateStoreCDCReader(

diff --git a/...streaming/state/StateStoreCDCReader.scala → ...ng/state/StateStoreChangeDataReader.scala b/...streaming/state/StateStoreCDCReader.scala → ...ng/state/StateStoreChangeDataReader.scala
@@ -30,8 +30,9 @@ import org.apache.spark.util.NextIterator
  * This is an optional trait for [[StateStoreProvider]]s to mix in if they support reading state
  * change data. It is used by the readChangeFeed option of State Data Source.
  */
-trait SupportsStateChangeDataFeed {
-  def getStateChangeDataReader(startVersion: Long, endVersion: Long): StateChangeDataReader
+trait SupportsStateStoreChangeDataFeed {
+  def getStateStoreChangeDataReader(startVersion: Long, endVersion: Long):
+    StateStoreChangeDataReader
 }
 
 /**
@@ -40,7 +41,7 @@ trait SupportsStateChangeDataFeed {
  * @param fileToRead - name of file to use to read changelog
  * @param compressionCodec - de-compression method using for reading changelog file
  */
-abstract class StateChangeDataReader(
+abstract class StateStoreChangeDataReader(
     fm: CheckpointFileManager,
     stateLocation: Path,
     startVersion: Long,
@@ -93,7 +94,7 @@ class HDFSBackedStateStoreCDCReader(
     keySchema: StructType,
     valueSchema: StructType
   )
-  extends StateChangeDataReader(
+  extends StateStoreChangeDataReader(
     fm, stateLocation, startVersion, endVersion, compressionCodec) {
   override protected var changelogSuffix: String = "delta"
 
@@ -140,7 +141,7 @@ class RocksDBStateStoreCDCReader(
   keySchema: StructType,
   valueSchema: StructType
 )
-  extends StateChangeDataReader(
+  extends StateStoreChangeDataReader(
     fm, stateLocation, startVersion, endVersion, compressionCodec) {
   override protected var changelogSuffix: String = "changelog"