address commenet

Author: zifeif2

common/utils/src/main/resources/error/error-conditions.json

@@ -5480,6 +5480,11 @@
         "message" : [
           "Unsupported offset sequence version <version>. Please make sure the checkpoint is from a supported Spark version (Spark 4.0+)."
         ]
+      },
+      "UNSUPPORTED_PROVIDER" : {
+        "message" : [
+          "<provider> is not supported"
+        ]
       }
     },
     "sqlState" : "55019"
@@ -5507,19 +5512,6 @@
     },
     "sqlState" : "42616"
   },
-  "STATE_REPARTITION_INVALID_STATE_STORE_CONFIG" : {
-    "message" : [
-      "StateStoreConfig <configName> is invalid:"
-    ],
-    "subClass" : {
-      "UNSUPPORTED_PROVIDER" : {
-        "message" : [
-          "<provider> is not supported"
-        ]
-      }
-    },
-    "sqlState" : "42617"
-  },
   "STATE_STORE_CANNOT_CREATE_COLUMN_FAMILY_WITH_RESERVED_CHARS" : {
     "message" : [
       "Failed to create column family with unsupported starting character and name=<colFamilyName>."

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

@@ -67,11 +67,12 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
     val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
       StateSourceOptions.apply(session, hadoopConf, properties))
     val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
-    // We only support RocksDB because the repartition work that this reader
+    // We only support RocksDB because the repartition work that this option
     // is built for only supports RocksDB
     if (sourceOptions.internalOnlyReadAllColumnFamilies
       && stateConf.providerClass != classOf[RocksDBStateStoreProvider].getName) {
       throw OfflineStateRepartitionErrors.unsupportedStateStoreProviderError(
+        sourceOptions.resolvedCpLocation,
         stateConf.providerClass)
     }
     val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(

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

@@ -16,6 +16,8 @@
  */
 package org.apache.spark.sql.execution.datasources.v2.state
 
+import scala.collection.mutable
+
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow}
@@ -50,17 +52,8 @@ class StatePartitionReaderFactory(
   override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
     val stateStoreInputPartition = partition.asInstanceOf[StateStoreInputPartition]
     if (stateStoreInputPartition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-      // Disable format validation because the schema returned by
-      // StatePartitionAllColumnFamiliesReader does not contain the corresponding
-      // keySchema or valueSchema.
-      // It's safe to do so we also don't expect the caller of StatePartitionAllColumnFamiliesReader
-      // to extract specific fields out of the returning row.
-      val modifiedStoreConf = storeConf.withExtraOptions(Map(
-        StateStoreConf.FORMAT_VALIDATION_ENABLED_CONFIG -> "false",
-        StateStoreConf.FORMAT_VALIDATION_CHECK_VALUE_CONFIG -> "false"
-      ))
-      new StatePartitionAllColumnFamiliesReader(modifiedStoreConf, hadoopConf,
-        stateStoreInputPartition, schema, keyStateEncoderSpec)
+      new StatePartitionAllColumnFamiliesReader(storeConf, hadoopConf,
+        stateStoreInputPartition, schema, keyStateEncoderSpec, joinColFamilyOpt)
     } else if (stateStoreInputPartition.sourceOptions.readChangeFeed) {
       new StateStoreChangeDataPartitionReader(storeConf, hadoopConf,
         stateStoreInputPartition, schema, keyStateEncoderSpec, stateVariableInfoOpt,
@@ -99,12 +92,14 @@ abstract class StatePartitionReaderBase(
     if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
       SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
     } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+      // StatePartitionAllFamiliesReader will have its own provider won't use this keySchema
       placeholderSchema
     } else {
       SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
     }
   }
 
+  // StatePartitionAllFamiliesReader will have its own provider won't use this valueSchema
   protected val valueSchema = if (stateVariableInfoOpt.isDefined ||
       partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
     placeholderSchema
@@ -265,7 +260,8 @@ class StatePartitionAllColumnFamiliesReader(
     hadoopConf: SerializableConfiguration,
     partition: StateStoreInputPartition,
     schema: StructType,
-    keyStateEncoderSpec: KeyStateEncoderSpec)
+    keyStateEncoderSpec: KeyStateEncoderSpec,
+    joinColFamilyOpt: Option[String])
   extends StatePartitionReaderBase(
     storeConf,
     hadoopConf, partition, schema,
@@ -280,6 +276,37 @@ class StatePartitionAllColumnFamiliesReader(
     )
   }
 
+  private val colFamilyToSchema : mutable.HashMap[String, StateStoreColFamilySchema] = {
+    val stateStoreId = StateStoreId(
+      partition.sourceOptions.stateCheckpointLocation.toString,
+      partition.sourceOptions.operatorId,
+      StateStore.PARTITION_ID_TO_CHECK_SCHEMA,
+      partition.sourceOptions.storeName)
+    val stateStoreProviderId = StateStoreProviderId(stateStoreId, partition.queryId)
+    val manager = new StateSchemaCompatibilityChecker(stateStoreProviderId, hadoopConf.value)
+    val schemaFile = manager.readSchemaFile()
+    val schemaMap = mutable.HashMap[String, StateStoreColFamilySchema]()
+    schemaFile.foreach { schema => schemaMap.put(schema.colFamilyName, schema)}
+    schemaMap
+  }
+
+  override lazy val provider: StateStoreProvider = {
+    val stateStoreId = StateStoreId(
+      partition.sourceOptions.stateCheckpointLocation.toString,
+      partition.sourceOptions.operatorId,
+      partition.partition,
+      partition.sourceOptions.storeName)
+    val stateStoreProviderId = StateStoreProviderId(stateStoreId, partition.queryId)
+    val provider = StateStoreProvider.createAndInit(
+      stateStoreProviderId,
+      colFamilyToSchema(StateStore.DEFAULT_COL_FAMILY_NAME).keySchema,
+      colFamilyToSchema(StateStore.DEFAULT_COL_FAMILY_NAME).valueSchema,
+      keyStateEncoderSpec,
+      useColumnFamilies = false, storeConf, hadoopConf.value,
+      useMultipleValuesPerKey = false, None)
+    provider
+  }
+
   override lazy val iter: Iterator[InternalRow] = {
     store
       .iterator()

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

@@ -17,8 +17,7 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
-import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, SparkUnsupportedOperationException}
-import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException}
 
 /**
  * Errors thrown by Offline state repartitioning.
@@ -88,9 +87,10 @@ object OfflineStateRepartitionErrors {
   }
 
   def unsupportedStateStoreProviderError(
+      checkpointLocation: String,
       providerClass: String
-    ): StateRepartitionInvalidStateStoreConfigUnsupportedProviderError = {
-    new StateRepartitionInvalidStateStoreConfigUnsupportedProviderError(providerClass)
+    ): StateRepartitionUnsupportedProviderError = {
+    new StateRepartitionUnsupportedProviderError(checkpointLocation, providerClass)
   }
 }
 
@@ -209,16 +209,9 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
     subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
     messageParameters = Map("version" -> version.toString))
 
-abstract class StateRepartitionInvalidStateStoreConfigError(
-    configName: String,
-    subClass: String,
-    messageParameters: Map[String, String] = Map.empty)
-  extends SparkUnsupportedOperationException(
-    errorClass = s"STATE_REPARTITION_INVALID_STATE_STORE_CONFIG.$subClass",
-    messageParameters = Map("configName" -> configName) ++ messageParameters)
-
-class StateRepartitionInvalidStateStoreConfigUnsupportedProviderError(
-    provider: String) extends StateRepartitionInvalidStateStoreConfigError(
-    SQLConf.STATE_STORE_PROVIDER_CLASS.key,
+class StateRepartitionUnsupportedProviderError(
+    checkpointLocation: String,
+    provider: String) extends StateRepartitionInvalidCheckpointError(
+    checkpointLocation,
     subClass = "UNSUPPORTED_PROVIDER",
     messageParameters = Map("provider" -> provider))

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

@@ -30,18 +30,6 @@ class StateStoreConf(
 
   def this() = this(new SQLConf)
 
-  def withExtraOptions(additionalOptions: Map[String, String]): StateStoreConf = {
-    val reconstructedSqlConf = {
-      // Reconstruct a SQLConf with the all settings preserved because sqlConf is transient
-      val conf = new SQLConf()
-      // Restore all state store related settings
-      sqlConfs.foreach { case (key, value) =>
-        conf.setConfString(key, value)
-      }
-      conf
-    }
-    new StateStoreConf(reconstructedSqlConf, extraOptions ++ additionalOptions)
-  }
   /**
    * Size of MaintenanceThreadPool to perform maintenance tasks for StateStore
    */
@@ -95,10 +83,7 @@ class StateStoreConf(
   val providerClass: String = sqlConf.stateStoreProviderClass
 
   /** Whether validate the underlying format or not. */
-  val formatValidationEnabled: Boolean = extraOptions.get(
-      StateStoreConf.FORMAT_VALIDATION_ENABLED_CONFIG)
-    .map(_ == "true")
-    .getOrElse(sqlConf.stateStoreFormatValidationEnabled)
+  val formatValidationEnabled: Boolean = sqlConf.stateStoreFormatValidationEnabled
 
   /**
    * Whether to validate StateStore commits for ForeachBatch sinks to ensure all partitions
@@ -181,7 +166,6 @@ class StateStoreConf(
 }
 
 object StateStoreConf {
-  val FORMAT_VALIDATION_ENABLED_CONFIG = "formatValidationEnabled"
   val FORMAT_VALIDATION_CHECK_VALUE_CONFIG = "formatValidationCheckValue"
 
   val empty = new StateStoreConf()