small changes

Author: zifeif2

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

@@ -67,6 +67,8 @@ 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
+    // is built for only supports RocksDB
     if (sourceOptions.internalOnlyReadAllColumnFamilies
       && stateConf.providerClass != classOf[RocksDBStateStoreProvider].getName) {
       throw OfflineStateRepartitionErrors.unsupportedStateStoreProviderError(
@@ -386,8 +388,8 @@ case class StateSourceOptions(
   override def toString: String = {
     var desc = s"StateSourceOptions(checkpointLocation=$resolvedCpLocation, batchId=$batchId, " +
       s"operatorId=$operatorId, storeName=$storeName, joinSide=$joinSide, " +
-      s"stateVarName=${stateVarName.getOrElse("None")}, +" +
-      s"flattenCollectionTypes=$flattenCollectionTypes" +
+      s"stateVarName=${stateVarName.getOrElse("None")}, " +
+      s"flattenCollectionTypes=$flattenCollectionTypes, " +
       s"internalOnlyReadAllColumnFamilies=$internalOnlyReadAllColumnFamilies"
     if (fromSnapshotOptions.isDefined) {
       desc += s", snapshotStartBatchId=${fromSnapshotOptions.get.snapshotStartBatchId}"
@@ -502,8 +504,7 @@ object StateSourceOptions extends DataSourceOptions {
     val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
 
     val internalOnlyReadAllColumnFamilies = try {
-      Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES))
-        .map(_.toBoolean).getOrElse(false)
+      Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
     } catch {
       case _: IllegalArgumentException =>
         throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
@@ -648,8 +649,8 @@ object StateSourceOptions extends DataSourceOptions {
     StateSourceOptions(
       resolvedCpLocation, batchId.get, operatorId, storeName, joinSide,
       readChangeFeed, fromSnapshotOptions, readChangeFeedOptions,
-      stateVarName, readRegisteredTimers, flattenCollectionTypes,
-      internalOnlyReadAllColumnFamilies, startOperatorStateUniqueIds, endOperatorStateUniqueIds)
+      stateVarName, readRegisteredTimers, flattenCollectionTypes, internalOnlyReadAllColumnFamilies,
+      startOperatorStateUniqueIds, endOperatorStateUniqueIds)
   }
 
   private def getLastCommittedBatch(session: SparkSession, checkpointLocation: String): Long = {

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

@@ -50,6 +50,11 @@ 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"
@@ -85,23 +90,24 @@ abstract class StatePartitionReaderBase(
   extends PartitionReader[InternalRow] with Logging {
   // Used primarily as a placeholder for the value schema in the context of
   // state variables used within the transformWithState operator.
-  private val schemaForValueRow: StructType =
+  // Also used as a placeholder for both key and value schema for
+  // StatePartitionAllColumnFamiliesReader
+  private val placeholderSchema: StructType =
     StructType(Array(StructField("__dummy__", NullType)))
 
   protected val keySchema = {
     if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
       SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
     } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-      schemaForValueRow
+      placeholderSchema
     } else {
       SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
     }
   }
 
-  protected val valueSchema = if (stateVariableInfoOpt.isDefined) {
-    schemaForValueRow
-  } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-    schemaForValueRow
+  protected val valueSchema = if (stateVariableInfoOpt.isDefined ||
+      partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+    placeholderSchema
   } else {
     SchemaUtil.getSchemaAsDataType(
       schema, "value").asInstanceOf[StructType]
@@ -252,6 +258,7 @@ 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
+ * It will also return the partition key
  */
 class StatePartitionAllColumnFamiliesReader(
     storeConf: StateStoreConf,

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

@@ -63,7 +63,7 @@ object SchemaUtil {
         .add("partition_id", IntegerType)
     } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
       new StructType()
-        // todo [SPARK-54443]: change keySchema to more specific type after we
+        // todo [SPARK-54443]: change keySchema to a more specific type after we
         //  can extract partition key from keySchema
         .add("partition_key", keySchema)
         .add("key_bytes", BinaryType)

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

@@ -17,7 +17,8 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
-import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, SparkRuntimeException}
+import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.internal.SQLConf
 
 /**
  * Errors thrown by Offline state repartitioning.
@@ -211,15 +212,13 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
 abstract class StateRepartitionInvalidStateStoreConfigError(
     configName: String,
     subClass: String,
-    messageParameters: Map[String, String] = Map.empty,
-    cause: Throwable = null)
-  extends SparkRuntimeException(
+    messageParameters: Map[String, String] = Map.empty)
+  extends SparkUnsupportedOperationException(
     errorClass = s"STATE_REPARTITION_INVALID_STATE_STORE_CONFIG.$subClass",
-    messageParameters = Map("configName" -> configName) ++ messageParameters,
-    cause = cause)
+    messageParameters = Map("configName" -> configName) ++ messageParameters)
 
 class StateRepartitionInvalidStateStoreConfigUnsupportedProviderError(
     provider: String) extends StateRepartitionInvalidStateStoreConfigError(
-    "SQLConf.STATE_STORE_PROVIDER_CLASS.key",
+    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

@@ -95,9 +95,10 @@ class StateStoreConf(
   val providerClass: String = sqlConf.stateStoreProviderClass
 
   /** Whether validate the underlying format or not. */
-  val formatValidationEnabled: Boolean = extraOptions.getOrElse(
-    StateStoreConf.FORMAT_VALIDATION_ENABLED_CONFIG,
-    sqlConf.stateStoreFormatValidationEnabled) == "true"
+  val formatValidationEnabled: Boolean = extraOptions.get(
+      StateStoreConf.FORMAT_VALIDATION_ENABLED_CONFIG)
+    .map(_ == "true")
+    .getOrElse(sqlConf.stateStoreFormatValidationEnabled)
 
   /**
    * Whether to validate StateStore commits for ForeachBatch sinks to ensure all partitions

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.state
 import java.nio.ByteOrder
 import java.util.Arrays
 
-import org.apache.spark.SparkRuntimeException
+import org.apache.spark.SparkUnsupportedOperationException
 import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
 import org.apache.spark.sql.catalyst.InternalRow
@@ -100,20 +100,16 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
 
     // Create projections to convert Row to UnsafeRow bytes
     val keyProjection = UnsafeProjection.create(keySchema)
-    val valueProjection = if (valueSchema.isEmpty) null else UnsafeProjection.create(valueSchema)
+    val valueProjection = UnsafeProjection.create(valueSchema)
 
     // Create converters to convert external Row types to internal Catalyst types
     val keyConverter = CatalystTypeConverters.createToCatalystConverter(keySchema)
-    val valueConverter = if (valueSchema.isEmpty) {
-      null
-    } else {
-      CatalystTypeConverters.createToCatalystConverter(valueSchema)
-    }
+    val valueConverter = CatalystTypeConverters.createToCatalystConverter(valueSchema)
 
     // Convert normal data to bytes
     val normalAsBytes = normalDf.map { row =>
       val key = row.getStruct(1)
-      val value = if (row.isNullAt(2) || valueSchema.isEmpty) null else row.getStruct(2)
+      val value = if (row.isNullAt(2)) null else row.getStruct(2)
 
       // Convert key to InternalRow, then to UnsafeRow, then get bytes
       val keyInternalRow = keyConverter(key).asInstanceOf[InternalRow]
@@ -123,7 +119,7 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
       val keyBytes = keyUnsafeRow.getBytes.clone()
 
       // Convert value to bytes
-      val valueBytes = if (value == null || valueSchema.isEmpty) {
+      val valueBytes = if (value == null) {
         Array.empty[Byte]
       } else {
         val valueInternalRow = valueConverter(value).asInstanceOf[InternalRow]
@@ -433,7 +429,7 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
         )
 
         checkError(
-          exception = intercept[SparkRuntimeException] {
+          exception = intercept[SparkUnsupportedOperationException] {
             spark.read
               .format("statestore")
               .option(StateSourceOptions.PATH, tempDir.getAbsolutePath)
@@ -443,7 +439,7 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
           },
           condition = "STATE_REPARTITION_INVALID_STATE_STORE_CONFIG.UNSUPPORTED_PROVIDER",
           parameters = Map(
-            "configName" -> "SQLConf.STATE_STORE_PROVIDER_CLASS.key",
+            "configName" -> SQLConf.STATE_STORE_PROVIDER_CLASS.key,
             "provider" -> classOf[HDFSBackedStateStoreProvider].getName
           )
         )