address comment

Author: zifeif2

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

@@ -5507,6 +5507,19 @@
     },
     "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

@@ -41,7 +41,8 @@ import org.apache.spark.sql.execution.streaming.operators.stateful.transformwith
 import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.timers.TimerStateUtils
 import org.apache.spark.sql.execution.streaming.runtime.StreamingCheckpointConstants.DIR_NAME_STATE
 import org.apache.spark.sql.execution.streaming.runtime.StreamingQueryCheckpointMetadata
-import org.apache.spark.sql.execution.streaming.state.{InMemoryStateSchemaProvider, KeyStateEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, StateSchemaCompatibilityChecker, StateSchemaMetadata, StateSchemaProvider, StateStore, StateStoreColFamilySchema, StateStoreConf, StateStoreId, StateStoreProviderId}
+import org.apache.spark.sql.execution.streaming.state.{InMemoryStateSchemaProvider, KeyStateEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RocksDBStateStoreProvider, StateSchemaCompatibilityChecker, StateSchemaMetadata, StateSchemaProvider, StateStore, StateStoreColFamilySchema, StateStoreConf, StateStoreId, StateStoreProviderId}
+import org.apache.spark.sql.execution.streaming.state.OfflineStateRepartitionErrors
 import org.apache.spark.sql.execution.streaming.utils.StreamingUtils
 import org.apache.spark.sql.sources.DataSourceRegister
 import org.apache.spark.sql.streaming.TimeMode
@@ -67,11 +68,9 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
       StateSourceOptions.apply(session, hadoopConf, properties))
     val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
     if (sourceOptions.internalOnlyReadAllColumnFamilies
-      && !stateConf.providerClass.contains("RocksDB")) {
-      throw StateDataSourceErrors.invalidOptionValue(
-        StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
-        "internalOnlyReadAllColumnFamilies=true is only supported with " +
-          s"RocksDBStateStoreProvider. Current provider: ${stateConf.providerClass}")
+      && stateConf.providerClass != classOf[RocksDBStateStoreProvider].getName) {
+      throw OfflineStateRepartitionErrors.unsupportedStateStoreProviderError(
+        stateConf.providerClass)
     }
     val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
       sourceOptions)
@@ -379,7 +378,7 @@ case class StateSourceOptions(
     stateVarName: Option[String],
     readRegisteredTimers: Boolean,
     flattenCollectionTypes: Boolean,
-    internalOnlyReadAllColumnFamilies: Boolean,
+    internalOnlyReadAllColumnFamilies: Boolean = false,
     startOperatorStateUniqueIds: Option[Array[Array[String]]] = None,
     endOperatorStateUniqueIds: Option[Array[Array[String]]] = None) {
   def stateCheckpointLocation: Path = new Path(resolvedCpLocation, DIR_NAME_STATE)

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

@@ -50,7 +50,11 @@ class StatePartitionReaderFactory(
   override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
     val stateStoreInputPartition = partition.asInstanceOf[StateStoreInputPartition]
     if (stateStoreInputPartition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-      new StatePartitionReaderAllColumnFamilies(storeConf, hadoopConf,
+      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)
     } else if (stateStoreInputPartition.sourceOptions.readChangeFeed) {
       new StateStoreChangeDataPartitionReader(storeConf, hadoopConf,
@@ -87,13 +91,17 @@ abstract class StatePartitionReaderBase(
   protected val keySchema = {
     if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
       SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
+    } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+      schemaForValueRow
     } else {
       SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
     }
   }
 
   protected val valueSchema = if (stateVariableInfoOpt.isDefined) {
     schemaForValueRow
+  } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+    schemaForValueRow
   } else {
     SchemaUtil.getSchemaAsDataType(
       schema, "value").asInstanceOf[StructType]
@@ -243,14 +251,17 @@ class StatePartitionReader(
 /**
  * An implementation of [[StatePartitionReaderBase]] for reading all column families
  * in binary format. This reader returns raw key and value bytes along with column family names.
+ * We are returning key/value bytes because each column family can have different schema
  */
-class StatePartitionReaderAllColumnFamilies(
+class StatePartitionAllColumnFamiliesReader(
     storeConf: StateStoreConf,
     hadoopConf: SerializableConfiguration,
     partition: StateStoreInputPartition,
     schema: StructType,
     keyStateEncoderSpec: KeyStateEncoderSpec)
-  extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+  extends StatePartitionReaderBase(
+    storeConf,
+    hadoopConf, partition, schema,
     keyStateEncoderSpec, None, None, None, None) {
 
   private lazy val store: ReadStateStore = {
@@ -263,7 +274,6 @@ class StatePartitionReaderAllColumnFamilies(
   }
 
   override lazy val iter: Iterator[InternalRow] = {
-    // Single store with column families (join v3, transformWithState, or simple operators)
     store
       .iterator()
       .map { pair =>

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

@@ -63,15 +63,12 @@ object SchemaUtil {
         .add("partition_id", IntegerType)
     } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
       new StructType()
-        // todo: change this to some more specific type after we
+        // todo [SPARK-54443]: change keySchema to more specific type after we
         //  can extract partition key from keySchema
         .add("partition_key", keySchema)
         .add("key_bytes", BinaryType)
         .add("value_bytes", BinaryType)
         .add("column_family_name", StringType)
-        // need key and value schema so that state store can encode data
-        .add("value", valueSchema)
-        .add("key", keySchema)
     } else {
       new StructType()
         .add("key", keySchema)
@@ -89,14 +86,18 @@ object SchemaUtil {
   }
 
   /**
-   * Creates a unified row from raw key and value bytes.
-   * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
-   * instead of a tuple for better readability.
+   * Returns an InternalRow representing
+   * 1. partitionKey
+   * 2. key in bytes
+   * 3. value in bytes
+   * 4. column family name
    */
   def unifyStateRowPairAsRawBytes(
-     pair: (UnsafeRow, UnsafeRow),
-     colFamilyName: String): InternalRow = {
-    val row = new GenericInternalRow(6)
+      pair: (UnsafeRow, UnsafeRow),
+      colFamilyName: String): InternalRow = {
+    val row = new GenericInternalRow(4)
+    // todo [SPARK-54443]: change keySchema to more specific type after we
+    //  can extract partition key from keySchema
     row.update(0, pair._1)
     row.update(1, pair._1.getBytes)
     row.update(2, pair._2.getBytes)
@@ -261,9 +262,9 @@ object SchemaUtil {
       "expiration_timestamp_ms" -> classOf[LongType],
       "partition_id" -> classOf[IntegerType],
       "partition_key" -> classOf[StructType],
-      "key_bytes"->classOf[BinaryType],
-      "value_bytes"->classOf[BinaryType],
-      "column_family_name"->classOf[StringType])
+      "key_bytes" -> classOf[BinaryType],
+      "value_bytes" -> classOf[BinaryType],
+      "column_family_name" -> classOf[StringType])
 
     val expectedFieldNames = if (transformWithStateVariableInfoOpt.isDefined) {
       val stateVarInfo = transformWithStateVariableInfoOpt.get
@@ -305,7 +306,7 @@ object SchemaUtil {
     } else if (sourceOptions.readChangeFeed) {
       Seq("batch_id", "change_type", "key", "value", "partition_id")
     } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
-      Seq("partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key")
+      Seq("partition_key", "key_bytes", "value_bytes", "column_family_name")
     } else {
       Seq("key", "value", "partition_id")
     }

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

@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
-import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException}
+import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, SparkRuntimeException}
 
 /**
  * Errors thrown by Offline state repartitioning.
@@ -85,6 +85,12 @@ object OfflineStateRepartitionErrors {
       version: Int): StateRepartitionInvalidCheckpointError = {
     new StateRepartitionUnsupportedOffsetSeqVersionError(checkpointLocation, version)
   }
+
+  def unsupportedStateStoreProviderError(
+      providerClass: String
+    ): StateRepartitionInvalidStateStoreConfigUnsupportedProviderError = {
+    new StateRepartitionInvalidStateStoreConfigUnsupportedProviderError(providerClass)
+  }
 }
 
 /**
@@ -201,3 +207,19 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
     checkpointLocation,
     subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
     messageParameters = Map("version" -> version.toString))
+
+abstract class StateRepartitionInvalidStateStoreConfigError(
+    configName: String,
+    subClass: String,
+    messageParameters: Map[String, String] = Map.empty,
+    cause: Throwable = null)
+  extends SparkRuntimeException(
+    errorClass = s"STATE_REPARTITION_INVALID_STATE_STORE_CONFIG.$subClass",
+    messageParameters = Map("configName" -> configName) ++ messageParameters,
+    cause = cause)
+
+class StateRepartitionInvalidStateStoreConfigUnsupportedProviderError(
+    provider: String) extends StateRepartitionInvalidStateStoreConfigError(
+    "SQLConf.STATE_STORE_PROVIDER_CLASS.key",
+    subClass = "UNSUPPORTED_PROVIDER",
+    messageParameters = Map("provider" -> provider))

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

@@ -30,6 +30,18 @@ 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
    */
@@ -83,7 +95,9 @@ class StateStoreConf(
   val providerClass: String = sqlConf.stateStoreProviderClass
 
   /** Whether validate the underlying format or not. */
-  val formatValidationEnabled: Boolean = sqlConf.stateStoreFormatValidationEnabled
+  val formatValidationEnabled: Boolean = extraOptions.getOrElse(
+    StateStoreConf.FORMAT_VALIDATION_ENABLED_CONFIG,
+    sqlConf.stateStoreFormatValidationEnabled) == "true"
 
   /**
    * Whether to validate StateStore commits for ForeachBatch sinks to ensure all partitions
@@ -166,6 +180,7 @@ class StateStoreConf(
 }
 
 object StateStoreConf {
+  val FORMAT_VALIDATION_ENABLED_CONFIG = "formatValidationEnabled"
   val FORMAT_VALIDATION_CHECK_VALUE_CONFIG = "formatValidationCheckValue"
 
   val empty = new StateStoreConf()