address comment

Author: zifeif2

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

@@ -81,11 +81,6 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
     // The key state encoder spec should be available for all operators except stream-stream joins
     val keyStateEncoderSpec = if (stateStoreReaderInfo.keyStateEncoderSpecOpt.isDefined) {
       stateStoreReaderInfo.keyStateEncoderSpecOpt.get
-    } else if (isReadAllColFamiliesOnJoinV3(sourceOptions)) {
-      // Create keyStateEncoderSpec here because getStoreMetadataAndRunChecks
-      // doesn't assign it in stateStoreReaderInfo.keyStateEncoderSpecOpt
-      NoPrefixKeyStateEncoderSpec(
-        stateStoreReaderInfo.allColumnFamiliesReaderInfo.colFamilySchemas.head.keySchema)
     } else {
       val keySchema = SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
       NoPrefixKeyStateEncoderSpec(keySchema)
@@ -119,17 +114,10 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
             sourceOptions.operatorId, RightSide, oldSchemaFilePaths)
 
         case JoinSideValues.none =>
-          if (isReadAllColFamiliesOnJoinV3(sourceOptions)) {
-            // readAllColumnFamiliesReader on joinV3 reads schema with StreamStreamJoinStateHelper
-            StreamStreamJoinStateHelper.readKeyValueSchema(session,
-              stateCheckpointLocation.toString,
-              sourceOptions.operatorId, LeftSide, oldSchemaFilePaths)
-          } else {
-            // we should have the schema for the state store if joinSide is none
-            require(stateStoreReaderInfo.stateStoreColFamilySchemaOpt.isDefined)
-            val resultSchema = stateStoreReaderInfo.stateStoreColFamilySchemaOpt.get
-            (resultSchema.keySchema, resultSchema.valueSchema)
-          }
+          // we should have the schema for the state store if joinSide is none
+          require(stateStoreReaderInfo.stateStoreColFamilySchemaOpt.isDefined)
+          val resultSchema = stateStoreReaderInfo.stateStoreColFamilySchemaOpt.get
+          (resultSchema.keySchema, resultSchema.valueSchema)
       }
 
       SchemaUtil.getSourceSchema(sourceOptions, keySchema,
@@ -148,8 +136,9 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
    * Returns true if this is a read-all-column-families request for a stream-stream join
    * that uses virtual column families (state format version 3).
    */
-  private def isReadAllColFamiliesOnJoinV3(sourceOptions: StateSourceOptions): Boolean = {
-    val storeMetadata = StateDataSource.getStateStoreMetadata(sourceOptions, hadoopConf)
+  private def isReadAllColFamiliesOnJoinV3(
+      sourceOptions: StateSourceOptions,
+      storeMetadata: Array[StateMetadataTableEntry]): Boolean = {
     sourceOptions.internalOnlyReadAllColumnFamilies &&
       storeMetadata.head.operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
       StreamStreamJoinStateHelper.usesVirtualColumnFamilies(
@@ -270,6 +259,8 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
     StateStoreReaderInfo = {
     val storeMetadata = StateDataSource.getStateStoreMetadata(sourceOptions, hadoopConf)
     if (!sourceOptions.internalOnlyReadAllColumnFamilies) {
+      // skipping runStateVarChecks for StatePartitionAllColumnFamiliesReader because
+      // we won't specify any stateVars when querying a TWS operator
       runStateVarChecks(sourceOptions, storeMetadata)
     }
 
@@ -299,15 +290,15 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
           if (sourceOptions.readRegisteredTimers) {
             stateVarName = TimerStateUtils.getTimerStateVarNames(timeMode)._1
           }
-          if (!sourceOptions.internalOnlyReadAllColumnFamilies) {
+          if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+            stateVariableInfos = operatorProperties.stateVariables
+          } else {
             val stateVarInfoList = operatorProperties.stateVariables
               .filter(stateVar => stateVar.stateName == stateVarName)
             require(stateVarInfoList.size == 1, s"Failed to find unique state variable info " +
               s"for state variable $stateVarName in operator ${sourceOptions.operatorId}")
             val stateVarInfo = stateVarInfoList.head
             transformWithStateVariableInfoOpt = Some(stateVarInfo)
-          } else {
-            stateVariableInfos = operatorProperties.stateVariables
           }
           val schemaFilePaths = storeMetadataEntry.stateSchemaFilePaths
           val stateSchemaMetadata = StateSchemaMetadata.createStateSchemaMetadata(
@@ -335,8 +326,8 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
         val storeId = new StateStoreId(stateCheckpointLocation.toString, sourceOptions.operatorId,
           partitionId, sourceOptions.storeName)
         val providerId = new StateStoreProviderId(storeId, UUID.randomUUID())
-        val manager = new StateSchemaCompatibilityChecker(
-          providerId, hadoopConf, oldSchemaFilePaths)
+        val manager = new StateSchemaCompatibilityChecker(providerId, hadoopConf,
+          oldSchemaFilePaths = oldSchemaFilePaths)
         val stateSchema = manager.readSchemaFile()
 
         if (sourceOptions.internalOnlyReadAllColumnFamilies) {
@@ -345,16 +336,18 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
         }
         // When reading all column families for Join V3, no specific state variable is targeted,
         // so stateVarName defaults to DEFAULT_COL_FAMILY_NAME.
-        // However, Join V3 does not have a "default" column family. Therefore, we skip populating
-        // keyStateEncoderSpec and stateStoreColFamilySchemaOpt in this case, as there is no
-        // matching schema for the default column family name.
-        if (!isReadAllColFamiliesOnJoinV3(sourceOptions)) {
-          // Based on the version and read schema, populate the keyStateEncoderSpec used for
-          // reading the column families
-          val resultSchema = stateSchema.filter(_.colFamilyName == stateVarName).head
-          keyStateEncoderSpecOpt = Some(getKeyStateEncoderSpec(resultSchema, storeMetadata))
-          stateStoreColFamilySchemaOpt = Some(resultSchema)
+        // However, Join V3 does not have a "default" column family. Therefore, we pick the first
+        // schema as resultSchema which will be used as placeholder schema for default schema
+        // in StatePartitionAllColumnFamiliesReader
+        val resultSchema = if (isReadAllColFamiliesOnJoinV3(sourceOptions, storeMetadata)) {
+          stateSchema.head
+        } else {
+          stateSchema.filter(_.colFamilyName == stateVarName).head
         }
+        // Based on the version and read schema, populate the keyStateEncoderSpec used for
+        // reading the column families
+        keyStateEncoderSpecOpt = Some(getKeyStateEncoderSpec(resultSchema, storeMetadata))
+        stateStoreColFamilySchemaOpt = Some(resultSchema)
       } catch {
         case NonFatal(ex) =>
           throw StateDataSourceErrors.failedToReadStateSchema(sourceOptions, ex)

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

@@ -30,8 +30,8 @@ import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.{NextIterator, SerializableConfiguration}
 
 case class AllColumnFamiliesReaderInfo(
-  colFamilySchemas: List[StateStoreColFamilySchema] = List.empty,
-  stateVariableInfos: List[TransformWithStateVariableInfo] = List.empty)
+    colFamilySchemas: List[StateStoreColFamilySchema] = List.empty,
+    stateVariableInfos: List[TransformWithStateVariableInfo] = List.empty)
 
 /**
  * An implementation of [[PartitionReaderFactory]] for State data source. This is used to support
@@ -55,9 +55,10 @@ class StatePartitionReaderFactory(
   override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
     val stateStoreInputPartition = partition.asInstanceOf[StateStoreInputPartition]
     if (stateStoreInputPartition.sourceOptions.internalOnlyReadAllColumnFamilies) {
+      require(allColumnFamiliesReaderInfo.isDefined)
       new StatePartitionAllColumnFamiliesReader(storeConf, hadoopConf,
-        stateStoreInputPartition, schema, keyStateEncoderSpec,
-        allColumnFamiliesReaderInfo.getOrElse(AllColumnFamiliesReaderInfo()))
+        stateStoreInputPartition, schema, keyStateEncoderSpec, stateStoreColFamilySchemaOpt,
+        stateSchemaProviderOpt, allColumnFamiliesReaderInfo.get)
     } else if (stateStoreInputPartition.sourceOptions.readChangeFeed) {
       new StateStoreChangeDataPartitionReader(storeConf, hadoopConf,
         stateStoreInputPartition, schema, keyStateEncoderSpec, stateVariableInfoOpt,
@@ -87,23 +88,25 @@ 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 dummySchema: StructType =
+  private val schemaForValueRow: StructType =
     StructType(Array(StructField("__dummy__", NullType)))
 
   protected val keySchema : StructType = {
     if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
       SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
     } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-      stateStoreColFamilySchemaOpt.map(_.keySchema).getOrElse(dummySchema)
+      require(stateStoreColFamilySchemaOpt.isDefined)
+      stateStoreColFamilySchemaOpt.map(_.keySchema).get
     } else {
       SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
     }
   }
 
   protected val valueSchema : StructType = if (stateVariableInfoOpt.isDefined) {
-    dummySchema
+    schemaForValueRow
   } else if (partition.sourceOptions.internalOnlyReadAllColumnFamilies) {
-    stateStoreColFamilySchemaOpt.map(_.valueSchema).getOrElse(dummySchema)
+    require(stateStoreColFamilySchemaOpt.isDefined)
+    stateStoreColFamilySchemaOpt.map(_.valueSchema).get
   } else {
     SchemaUtil.getSchemaAsDataType(
       schema, "value").asInstanceOf[StructType]
@@ -262,14 +265,15 @@ class StatePartitionAllColumnFamiliesReader(
     partition: StateStoreInputPartition,
     schema: StructType,
     keyStateEncoderSpec: KeyStateEncoderSpec,
+    defaultStateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
+    stateSchemaProviderOpt: Option[StateSchemaProvider],
     allColumnFamiliesReaderInfo: AllColumnFamiliesReaderInfo)
   extends StatePartitionReaderBase(
     storeConf,
     hadoopConf, partition, schema,
     keyStateEncoderSpec, None,
-    allColumnFamiliesReaderInfo.colFamilySchemas.find(
-      _.colFamilyName == StateStore.DEFAULT_COL_FAMILY_NAME),
-    None, None) {
+    defaultStateStoreColFamilySchemaOpt,
+    stateSchemaProviderOpt, None) {
 
   private val stateStoreColFamilySchemas = allColumnFamiliesReaderInfo.colFamilySchemas
   private val stateVariableInfos = allColumnFamiliesReaderInfo.stateVariableInfos
@@ -280,7 +284,6 @@ class StatePartitionAllColumnFamiliesReader(
       StateVariableType.ListState)
   }
 
-  // Override provider to register ALL column families
   override protected lazy val provider: StateStoreProvider = {
     val stateStoreId = StateStoreId(partition.sourceOptions.stateCheckpointLocation.toString,
       partition.sourceOptions.operatorId, partition.partition, partition.sourceOptions.storeName)
@@ -289,7 +292,30 @@ class StatePartitionAllColumnFamiliesReader(
     StateStoreProvider.createAndInit(
       stateStoreProviderId, keySchema, valueSchema, keyStateEncoderSpec,
       useColumnFamilies, storeConf, hadoopConf.value,
-      useMultipleValuesPerKey = false, stateSchemaProvider = None)
+      useMultipleValuesPerKey = false, stateSchemaProviderOpt)
+  }
+
+
+  private def checkAllColFamiliesExist(
+      colFamilyNames: List[String], stateStore: StateStore
+    ): Unit = {
+    // Filter out DEFAULT column family from validation for two reasons:
+    // 1. Some operators (e.g., stream-stream join v3) don't include DEFAULT in their schema
+    //    because the underlying RocksDB creates "default" column family automatically
+    // 2. The default column family schema is handled separately via
+    //    defaultStateStoreColFamilySchemaOpt, so no need to verify it here
+    val actualCFs = colFamilyNames.toSet.filter(_ != StateStore.DEFAULT_COL_FAMILY_NAME)
+    val expectedCFs = stateStore.allColumnFamilyNames
+      .filter(_ != StateStore.DEFAULT_COL_FAMILY_NAME)
+
+    // Validation: All column families found in the checkpoint must be declared in the schema.
+    // It's acceptable if some schema CFs are not in expectedCFs - this just means those
+    // column families have no data yet in the checkpoint
+    // (they'll be created during registration).
+    // However, if the checkpoint contains CFs not in the schema, it indicates a mismatch.
+    require(expectedCFs.subsetOf(actualCFs),
+      s"Checkpoint contains unexpected column families. " +
+        s"Column families in checkpoint but not in schema: ${expectedCFs.diff(actualCFs)}")
   }
 
   // Use a single store instance for both registering column families and iteration.
@@ -305,6 +331,7 @@ class StatePartitionAllColumnFamiliesReader(
 
     // Register all column families from the schema
     if (stateStoreColFamilySchemas.length > 1) {
+      checkAllColFamiliesExist(stateStoreColFamilySchemas.map(_.colFamilyName), stateStore)
       stateStoreColFamilySchemas.foreach { cfSchema =>
         cfSchema.colFamilyName match {
           case StateStore.DEFAULT_COL_FAMILY_NAME => // createAndInit has registered default

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

@@ -146,6 +146,9 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
       throw StateStoreErrors.multipleColumnFamiliesNotSupported(providerName)
     }
 
+    override def allColumnFamilyNames: collection.Set[String] =
+      collection.Set[String](StateStore.DEFAULT_COL_FAMILY_NAME)
+
     // Multiple col families are not supported with HDFSBackedStateStoreProvider. Throw an exception
     // if the user tries to use a non-default col family.
     private def assertUseOfDefaultColFamily(colFamilyName: String): Unit = {

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

@@ -336,6 +336,14 @@ class RocksDB(
     colFamilyNameToInfoMap.asScala.values.toSeq.count(_.isInternal == isInternal)
   }
 
+  /**
+   * Returns all column family names currently registered in RocksDB.
+   * This includes column families loaded from checkpoint metadata.
+   */
+  def allColumnFamilyNames: collection.Set[String] = {
+    colFamilyNameToInfoMap.asScala.keySet.toSet
+  }
+
   private val rocksDBFileMapping: RocksDBFileMapping = new RocksDBFileMapping()
 
   // We send snapshots that needs to be uploaded by the maintenance thread to this queue

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

@@ -610,6 +610,10 @@ private[sql] class RocksDBStateStoreProvider
 
     override def hasCommitted: Boolean = state == COMMITTED
 
+    override def allColumnFamilyNames: collection.Set[String] = {
+      rocksDB.allColumnFamilyNames
+    }
+
     override def toString: String = {
       s"RocksDBStateStore[stateStoreId=$stateStoreId_, version=$version]"
     }

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

@@ -310,6 +310,13 @@ trait StateStore extends ReadStateStore {
    * Whether all updates have been committed
    */
   def hasCommitted: Boolean
+
+  /**
+   * Returns all column family names in this state store.
+   *
+   * @return Set of all column family names
+   */
+  def allColumnFamilyNames: collection.Set[String]
 }
 
 /** Wraps the instance of StateStore to make the instance read-only. */

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala

@@ -42,6 +42,9 @@ class MemoryStateStore extends StateStore() {
     throw StateStoreErrors.multipleColumnFamiliesNotSupported("MemoryStateStoreProvider")
   }
 
+  override def allColumnFamilyNames: collection.Set[String] =
+    collection.Set[String](StateStore.DEFAULT_COL_FAMILY_NAME)
+
   override def removeColFamilyIfExists(colFamilyName: String): Boolean = {
     throw StateStoreErrors.removingColumnFamiliesNotSupported("MemoryStateStoreProvider")
   }

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala

@@ -122,6 +122,8 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
     )
   }
 
+  override def allColumnFamilyNames: collection.Set[String] = innerStore.allColumnFamilyNames
+
   override def put(
       key: UnsafeRow,
       value: UnsafeRow,