initial commit

Author: zifeif2

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

@@ -44,6 +44,7 @@ import org.apache.spark.sql.execution.streaming.runtime.StreamingQueryCheckpoint
 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.internal.SQLConf
 import org.apache.spark.sql.sources.DataSourceRegister
 import org.apache.spark.sql.streaming.TimeMode
 import org.apache.spark.sql.types.StructType
@@ -75,9 +76,18 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
         sourceOptions.resolvedCpLocation,
         stateConf.providerClass)
     }
-    val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
+    val (stateStoreReaderInfo, storeMetadata) = getStoreMetadataAndRunChecks(
       sourceOptions)
 
+    // Extract stateFormatVersion from StateStoreConf for SYMMETRIC_HASH_JOIN operator
+    val isJoin = (
+      storeMetadata.head.operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME)
+    val stateFormatVersion: Int = if (storeMetadata.nonEmpty && isJoin) {
+      session.conf.get(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION)
+    } else {
+      1
+    }
+
     // The key state encoder spec should be available for all operators except stream-stream joins
     val keyStateEncoderSpec = if (stateStoreReaderInfo.keyStateEncoderSpecOpt.isDefined) {
       stateStoreReaderInfo.keyStateEncoderSpecOpt.get
@@ -91,19 +101,28 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
       stateStoreReaderInfo.stateStoreColFamilySchemaOpt,
       stateStoreReaderInfo.stateSchemaProviderOpt,
       stateStoreReaderInfo.joinColFamilyOpt,
-      Option(stateStoreReaderInfo.allColumnFamiliesReaderInfo))
+      Option(stateStoreReaderInfo.allColumnFamiliesReaderInfo),
+      Option(stateFormatVersion))
   }
 
   override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
     val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
       StateSourceOptions.apply(session, hadoopConf, options))
 
-    val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
-      sourceOptions)
+    val (stateStoreReaderInfo, storeMetadata) = getStoreMetadataAndRunChecks(sourceOptions)
     val oldSchemaFilePaths = StateDataSource.getOldSchemaFilePaths(sourceOptions, hadoopConf)
 
+    // Extract stateFormatVersion from StateStoreConf for SYMMETRIC_HASH_JOIN operator
+    val stateFormatVersion = if (storeMetadata.nonEmpty &&
+      (storeMetadata.head.operatorName ==
+        StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME)) {
+        Some(session.conf.get(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION))
+    } else {
+      None
+    }
+
     val stateCheckpointLocation = sourceOptions.stateCheckpointLocation
-    try {
+//    try {
       val (keySchema, valueSchema) = sourceOptions.joinSide match {
         case JoinSideValues.left =>
           StreamStreamJoinStateHelper.readKeyValueSchema(session, stateCheckpointLocation.toString,
@@ -120,14 +139,24 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
           (resultSchema.keySchema, resultSchema.valueSchema)
       }
 
+      println("transformWithStateVariableInfoOpt",
+        stateStoreReaderInfo.transformWithStateVariableInfoOpt)
+      val stateVarInfo: Option[TransformWithStateVariableInfo] = if (
+        sourceOptions.internalOnlyReadAllColumnFamilies) {
+        Option(stateStoreReaderInfo.allColumnFamiliesReaderInfo.stateVariableInfos.head)
+      } else {
+        stateStoreReaderInfo.transformWithStateVariableInfoOpt
+      }
       SchemaUtil.getSourceSchema(sourceOptions, keySchema,
         valueSchema,
-        stateStoreReaderInfo.transformWithStateVariableInfoOpt,
-        stateStoreReaderInfo.stateStoreColFamilySchemaOpt)
-    } catch {
-      case NonFatal(e) =>
-        throw StateDataSourceErrors.failedToReadStateSchema(sourceOptions, e)
-    }
+        stateVarInfo,
+        stateStoreReaderInfo.stateStoreColFamilySchemaOpt,
+        storeMetadata,
+        stateFormatVersion)
+//    } catch {
+//      case NonFatal(e) =>
+//        throw StateDataSourceErrors.failedToReadStateSchema(sourceOptions, e)
+//    }
   }
 
   override def supportsExternalMetadata(): Boolean = false
@@ -256,7 +285,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
   }
 
   private def getStoreMetadataAndRunChecks(sourceOptions: StateSourceOptions):
-    StateStoreReaderInfo = {
+    (StateStoreReaderInfo, Array[StateMetadataTableEntry]) = {
     val storeMetadata = StateDataSource.getStateStoreMetadata(sourceOptions, hadoopConf)
     if (!sourceOptions.internalOnlyReadAllColumnFamilies) {
       // skipping runStateVarChecks for StatePartitionAllColumnFamiliesReader because
@@ -354,14 +383,14 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
       }
     }
 
-    StateStoreReaderInfo(
+    (StateStoreReaderInfo(
       keyStateEncoderSpecOpt,
       stateStoreColFamilySchemaOpt,
       transformWithStateVariableInfoOpt,
       stateSchemaProvider,
       joinColFamilyOpt,
       AllColumnFamiliesReaderInfo(stateStoreColFamilySchemas, stateVariableInfos)
-    )
+    ), storeMetadata)
   }
 
   private def getKeyStateEncoderSpec(

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

@@ -49,7 +49,8 @@ class StatePartitionReaderFactory(
     stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
     stateSchemaProviderOpt: Option[StateSchemaProvider],
     joinColFamilyOpt: Option[String],
-    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo])
+    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo],
+    stateFormatVersion: Option[Int])
   extends PartitionReaderFactory {
 
   override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
@@ -58,7 +59,7 @@ class StatePartitionReaderFactory(
       require(allColumnFamiliesReaderInfo.isDefined)
       new StatePartitionAllColumnFamiliesReader(storeConf, hadoopConf,
         stateStoreInputPartition, schema, keyStateEncoderSpec, stateStoreColFamilySchemaOpt,
-        stateSchemaProviderOpt, allColumnFamiliesReaderInfo.get)
+        stateSchemaProviderOpt, allColumnFamiliesReaderInfo.get, stateFormatVersion)
     } else if (stateStoreInputPartition.sourceOptions.readChangeFeed) {
       new StateStoreChangeDataPartitionReader(storeConf, hadoopConf,
         stateStoreInputPartition, schema, keyStateEncoderSpec, stateVariableInfoOpt,
@@ -267,7 +268,8 @@ class StatePartitionAllColumnFamiliesReader(
     keyStateEncoderSpec: KeyStateEncoderSpec,
     defaultStateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
     stateSchemaProviderOpt: Option[StateSchemaProvider],
-    allColumnFamiliesReaderInfo: AllColumnFamiliesReaderInfo)
+    allColumnFamiliesReaderInfo: AllColumnFamiliesReaderInfo,
+    stateFormatVersion: Option[Int])
   extends StatePartitionReaderBase(
     storeConf,
     hadoopConf, partition, schema,

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

@@ -47,10 +47,11 @@ class StateScanBuilder(
     stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
     stateSchemaProviderOpt: Option[StateSchemaProvider],
     joinColFamilyOpt: Option[String],
-    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo]) extends ScanBuilder {
+    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo],
+    stateFormatVersion: Option[Int]) extends ScanBuilder {
   override def build(): Scan = new StateScan(session, schema, sourceOptions, stateStoreConf,
     keyStateEncoderSpec, stateVariableInfoOpt, stateStoreColFamilySchemaOpt, stateSchemaProviderOpt,
-    joinColFamilyOpt, allColumnFamiliesReaderInfo)
+    joinColFamilyOpt, allColumnFamiliesReaderInfo, stateFormatVersion)
 }
 
 /** An implementation of [[InputPartition]] for State Store data source. */
@@ -70,7 +71,8 @@ class StateScan(
     stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
     stateSchemaProviderOpt: Option[StateSchemaProvider],
     joinColFamilyOpt: Option[String],
-    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo])
+    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo],
+    stateFormatVersion: Option[Int])
   extends Scan with Batch {
 
   // A Hadoop Configuration can be about 10 KB, which is pretty big, so broadcast it
@@ -146,7 +148,7 @@ class StateScan(
     case JoinSideValues.none =>
       new StatePartitionReaderFactory(stateStoreConf, hadoopConfBroadcast.value, schema,
         keyStateEncoderSpec, stateVariableInfoOpt, stateStoreColFamilySchemaOpt,
-        stateSchemaProviderOpt, joinColFamilyOpt, allColumnFamiliesReaderInfo)
+        stateSchemaProviderOpt, joinColFamilyOpt, allColumnFamiliesReaderInfo, stateFormatVersion)
   }
 
   override def toBatch: Batch = this

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

@@ -46,7 +46,8 @@ class StateTable(
     stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
     stateSchemaProviderOpt: Option[StateSchemaProvider],
     joinColFamilyOpt: Option[String],
-    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo] = None)
+    allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo] = None,
+    stateFormatVersion: Option[Int] = None)
   extends Table with SupportsRead with SupportsMetadataColumns {
 
   import StateTable._
@@ -88,7 +89,7 @@ class StateTable(
   override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
     new StateScanBuilder(session, schema, sourceOptions, stateConf, keyStateEncoderSpec,
       stateVariableInfoOpt, stateStoreColFamilySchemaOpt, stateSchemaProviderOpt,
-      joinColFamilyOpt, allColumnFamiliesReaderInfo)
+      joinColFamilyOpt, allColumnFamiliesReaderInfo, stateFormatVersion)
 
   override def properties(): util.Map[String, String] = Map.empty[String, String].asJava
 

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

@@ -25,9 +25,13 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow}
 import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData}
 import org.apache.spark.sql.execution.datasources.v2.state.{StateDataSourceErrors, StateSourceOptions}
+import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataTableEntry
+import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorsUtils, StatePartitionKeyExtractorFactory}
+import org.apache.spark.sql.execution.streaming.operators.stateful.join.StreamingSymmetricHashJoinHelper.LeftSide
+import org.apache.spark.sql.execution.streaming.operators.stateful.join.SymmetricHashJoinStateManager
 import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.{StateVariableType, TransformWithStateVariableInfo}
 import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.StateVariableType._
-import org.apache.spark.sql.execution.streaming.state.{ReadStateStore, StateStoreColFamilySchema, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{ReadStateStore, StatePartitionKeyExtractor, StateStore, StateStoreColFamilySchema, UnsafeRowPair}
 import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, IntegerType, LongType, MapType, StringType, StructType}
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.ArrayImplicits._
@@ -49,7 +53,10 @@ object SchemaUtil {
       keySchema: StructType,
       valueSchema: StructType,
       transformWithStateVariableInfoOpt: Option[TransformWithStateVariableInfo],
-      stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema]): StructType = {
+      stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
+      storeMetadata: Array[StateMetadataTableEntry],
+      stateFormatVersion: Option[Int] = None): StructType = {
+    println("stateFormatVersion", stateFormatVersion)
     if (transformWithStateVariableInfoOpt.isDefined) {
       require(stateStoreColFamilySchemaOpt.isDefined)
       generateSchemaForStateVar(transformWithStateVariableInfoOpt.get,
@@ -62,10 +69,13 @@ object SchemaUtil {
         .add("value", valueSchema)
         .add("partition_id", IntegerType)
     } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+      // Extract partition key schema using StatePartitionKeyExtractor
+      require(storeMetadata.nonEmpty)
+      val extractor = getExtractor(
+        storeMetadata.head.operatorName, keySchema, sourceOptions.storeName,
+        transformWithStateVariableInfoOpt, stateFormatVersion)
       new StructType()
-        // TODO [SPARK-54443]: change keySchema to a more specific type after we
-        // can extract partition key from keySchema
-        .add("partition_key", keySchema)
+        .add("partition_key", extractor.partitionKeySchema)
         .add("key_bytes", BinaryType)
         .add("value_bytes", BinaryType)
         .add("column_family_name", StringType)
@@ -77,6 +87,33 @@ object SchemaUtil {
     }
   }
 
+  /**
+   * Creates a StatePartitionKeyExtractor for the given operator.
+   * This is used to extract partition keys from state store keys for state repartitioning.
+   */
+  def getExtractor(
+      operatorName: String,
+      keySchema: StructType,
+      storeName: String,
+      transformWithStateVariableInfoOpt: Option[TransformWithStateVariableInfo],
+      stateFormatVersion: Option[Int]): StatePartitionKeyExtractor = {
+    val colFamilyName: String =
+      if (operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME) {
+        SymmetricHashJoinStateManager.allStateStoreNames(LeftSide).head
+      } else {
+        transformWithStateVariableInfoOpt.map(_.stateName)
+          .getOrElse(StateStore.DEFAULT_COL_FAMILY_NAME)
+      }
+    StatePartitionKeyExtractorFactory.create(
+      operatorName,
+      keySchema,
+      storeName = storeName,
+      colFamilyName = colFamilyName,
+      stateFormatVersion = stateFormatVersion,
+      transformWithStateVariableInfoOpt
+    )
+  }
+
   def unifyStateRowPair(pair: (UnsafeRow, UnsafeRow), partition: Int): InternalRow = {
     val row = new GenericInternalRow(3)
     row.update(0, pair._1)