reduce duplicate code

Author: zifeif2

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

@@ -79,14 +79,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
     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
-    }
+    val stateFormatVersion = getStateFormatVersion(storeMetadata)
 
     // The key state encoder spec should be available for all operators except stream-stream joins
     val keyStateEncoderSpec = if (stateStoreReaderInfo.keyStateEncoderSpecOpt.isDefined) {
@@ -102,7 +95,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
       stateStoreReaderInfo.stateSchemaProviderOpt,
       stateStoreReaderInfo.joinColFamilyOpt,
       Option(stateStoreReaderInfo.allColumnFamiliesReaderInfo),
-      Option(stateFormatVersion))
+      stateFormatVersion)
   }
 
   override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
@@ -112,14 +105,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
     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 stateFormatVersion = getStateFormatVersion(storeMetadata)
 
     val stateCheckpointLocation = sourceOptions.stateCheckpointLocation
     try {
@@ -159,6 +145,20 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
 
   override def supportsExternalMetadata(): Boolean = false
 
+  /**
+   * Returns the state format version for SYMMETRIC_HASH_JOIN operators.
+   * For join operators, returns the configured version; for other operators returns None.
+   */
+  private def getStateFormatVersion(
+      storeMetadata: Array[StateMetadataTableEntry]): Option[Int] = {
+    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
+    }
+  }
+
   /**
    * 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).
@@ -389,13 +389,14 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
       }
     }
 
+    val operatorName = if (storeMetadata.nonEmpty) storeMetadata.head.operatorName else ""
     (StateStoreReaderInfo(
       keyStateEncoderSpecOpt,
       stateStoreColFamilySchemaOpt,
       transformWithStateVariableInfoOpt,
       stateSchemaProvider,
       joinColFamilyOpt,
-      AllColumnFamiliesReaderInfo(stateStoreColFamilySchemas, stateVariableInfos)
+      AllColumnFamiliesReaderInfo(stateStoreColFamilySchemas, stateVariableInfos, operatorName)
     ), storeMetadata)
   }
 

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

@@ -31,7 +31,8 @@ import org.apache.spark.util.{NextIterator, SerializableConfiguration}
 
 case class AllColumnFamiliesReaderInfo(
     colFamilySchemas: List[StateStoreColFamilySchema] = List.empty,
-    stateVariableInfos: List[TransformWithStateVariableInfo] = List.empty)
+    stateVariableInfos: List[TransformWithStateVariableInfo] = List.empty,
+    operatorName: String = "")
 
 /**
  * An implementation of [[PartitionReaderFactory]] for State data source. This is used to support
@@ -280,6 +281,15 @@ class StatePartitionAllColumnFamiliesReader(
 
   private val stateStoreColFamilySchemas = allColumnFamiliesReaderInfo.colFamilySchemas
   private val stateVariableInfos = allColumnFamiliesReaderInfo.stateVariableInfos
+  private val operatorName = allColumnFamiliesReaderInfo.operatorName
+
+  // Create the extractor for partition key extraction
+  private lazy val partitionKeyExtractor = SchemaUtil.getExtractor(
+    operatorName,
+    keySchema,
+    partition.sourceOptions.storeName,
+    stateVariableInfos.headOption,
+    stateFormatVersion)
 
   private def isListType(colFamilyName: String): Boolean = {
     SchemaUtil.checkVariableType(
@@ -364,13 +374,14 @@ class StatePartitionAllColumnFamiliesReader(
           pair =>
             store.valuesIterator(pair.key, cfSchema.colFamilyName).map {
               value =>
-                SchemaUtil.unifyStateRowPairAsRawBytes((pair.key, value), cfSchema.colFamilyName)
+                SchemaUtil.unifyStateRowPairAsRawBytes(
+                  (pair.key, value), cfSchema.colFamilyName, partitionKeyExtractor)
             }
         )
       } else {
         store.iterator(cfSchema.colFamilyName).map { pair =>
           SchemaUtil.unifyStateRowPairAsRawBytes(
-            (pair.key, pair.value), cfSchema.colFamilyName)
+            (pair.key, pair.value), cfSchema.colFamilyName, partitionKeyExtractor)
         }
       }
     }

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

@@ -123,18 +123,18 @@ object SchemaUtil {
 
   /**
    * Returns an InternalRow representing
-   * 1. partitionKey
+   * 1. partitionKey (extracted using the StatePartitionKeyExtractor)
    * 2. key in bytes
    * 3. value in bytes
    * 4. column family name
    */
   def unifyStateRowPairAsRawBytes(
       pair: (UnsafeRow, UnsafeRow),
-      colFamilyName: String): InternalRow = {
+      colFamilyName: String,
+      extractor: StatePartitionKeyExtractor): 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)
+    val partitionKey = extractor.partitionKey(pair._1)
+    row.update(0, partitionKey)
     row.update(1, pair._1.getBytes)
     row.update(2, pair._2.getBytes)
     row.update(3, UTF8String.fromString(colFamilyName))