address commenet

Author: zifeif2

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

@@ -5493,6 +5493,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"
@@ -5520,19 +5525,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,16 +52,7 @@ 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,
+      new StatePartitionAllColumnFamiliesReader(storeConf, hadoopConf,
         stateStoreInputPartition, schema, keyStateEncoderSpec)
     } else if (stateStoreInputPartition.sourceOptions.readChangeFeed) {
       new StateStoreChangeDataPartitionReader(storeConf, hadoopConf,
@@ -95,19 +88,34 @@ abstract class StatePartitionReaderBase(
   private val placeholderSchema: StructType =
     StructType(Array(StructField("__dummy__", NullType)))
 
+  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
+  }
+
   protected val keySchema = {
     if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
       SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
     } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-      placeholderSchema
+      colFamilyToSchema(StateStore.DEFAULT_COL_FAMILY_NAME).keySchema
     } else {
       SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
     }
   }
 
-  protected val valueSchema = if (stateVariableInfoOpt.isDefined ||
-      partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+  protected val valueSchema = if (stateVariableInfoOpt.isDefined) {
     placeholderSchema
+  } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+    colFamilyToSchema(StateStore.DEFAULT_COL_FAMILY_NAME).valueSchema
   } else {
     SchemaUtil.getSchemaAsDataType(
       schema, "value").asInstanceOf[StructType]

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()