address comment

Author: zifeif2

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

@@ -1256,11 +1256,11 @@ object SymmetricHashJoinStateManager {
 
   private[streaming] sealed trait StateStoreType
 
-  private[streaming] case object KeyToNumValuesType extends StateStoreType {
+  private[sql] case object KeyToNumValuesType extends StateStoreType {
     override def toString(): String = "keyToNumValues"
   }
 
-  private[streaming] case object KeyWithIndexToValueType extends StateStoreType {
+  private[sql] case object KeyWithIndexToValueType extends StateStoreType {
     override def toString(): String = "keyWithIndexToValue"
   }
 

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

@@ -57,18 +57,19 @@ class StatePartitionAllColumnFamiliesWriter(
     colFamilyToWriterInfoMap: Map[String, StatePartitionWriterColumnFamilyInfo],
     operatorName: String,
     schemaProviderOpt: Option[StateSchemaProvider],
-    sqlConf: Map[String, String]) {
+    sqlConf: SQLConf) {
 
   private def isJoinV3Operator(
-      operatorName: String, sqlConf: Map[String, String]): Boolean = {
+      operatorName: String, sqlConf: SQLConf): Boolean = {
     operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
-      sqlConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION.key) == "3"
+      sqlConf.getConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION) == 3
   }
 
   private val defaultSchema = {
     colFamilyToWriterInfoMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) match {
       case Some(info) => info.schema
       case None =>
+        // joinV3 operator doesn't have default column family schema
         assert(isJoinV3Operator(operatorName, sqlConf),
           s"Please provide the schema of 'default' column family in StateStoreColFamilySchema" +
             s"for operator $operatorName")
@@ -173,14 +174,11 @@ class StatePartitionAllColumnFamiliesWriter(
     val valueRow = new UnsafeRow(columnFamilyToValueSchemaLenMap(colFamilyName))
     valueRow.pointTo(valueBytes, valueBytes.length)
 
-    if (colFamilyToWriterInfoMap(colFamilyName).useMultipleValuesPerKey) {
-      // if a column family useMultipleValuesPerKey (e.g. ListType), we will
-      // write with 1 put followed by merge
-      if (stateStore.keyExists(keyRow, colFamilyName)) {
-        stateStore.merge(keyRow, valueRow, colFamilyName)
-      } else {
-        stateStore.put(keyRow, valueRow, colFamilyName)
-      }
+    // if a column family useMultipleValuesPerKey (e.g. ListType), we will
+    // write with 1 put followed by merge
+    if (colFamilyToWriterInfoMap(colFamilyName).useMultipleValuesPerKey &&
+        stateStore.keyExists(keyRow, colFamilyName)) {
+      stateStore.merge(keyRow, valueRow, colFamilyName)
     } else {
       stateStore.put(keyRow, valueRow, colFamilyName)
     }

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

@@ -19,12 +19,14 @@ package org.apache.spark.sql.execution.datasources.v2.state
 import java.sql.Timestamp
 
 import org.apache.spark.sql.{DataFrame, Dataset}
+import org.apache.spark.sql.execution.streaming.operators.stateful.join.StreamingSymmetricHashJoinHelper.{LeftSide, RightSide}
+import org.apache.spark.sql.execution.streaming.operators.stateful.join.SymmetricHashJoinStateManager
 import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
-import org.apache.spark.sql.execution.streaming.state.{KeyStateEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
-import org.apache.spark.sql.streaming.util.StreamManualClock
+import org.apache.spark.sql.streaming.util.{ColumnFamilyMetadata, StreamManualClock}
 import org.apache.spark.sql.types.{BooleanType, IntegerType, LongType, NullType, StringType, StructField, StructType, TimestampType}
 
 trait StateDataSourceTestBase extends StreamTest with StateStoreMetricsTest {
@@ -459,12 +461,6 @@ case class SessionUpdate(
     numEvents: Int,
     expired: Boolean)
 
-case class ColumnFamilyMetadata(
-    keySchema: StructType,
-    valueSchema: StructType,
-    encoderSpec: KeyStateEncoderSpec,
-    useMultipleValuePerKey: Boolean = false)
-
 // Utility for runCompositeKeyStreamingAggregationQuery
 // todo: Move runCompositeKeyStreamingAggregationQuery to this class
 object CompositeKeyAggregationTestUtils {
@@ -568,10 +564,6 @@ object SimpleAggregationTestUtils {
     (metadata.keySchema, metadata.valueSchema)
   }
 
-  /**
-   * @param stateVersion The state format version:
-   * @return ColumnFamilyMetadata including schema and KeyEncoderSpec
-   */
   def getSchemasWithMetadata(stateVersion: Int): ColumnFamilyMetadata = {
     val keySchema = StructType(Array(
       StructField("groupKey", IntegerType, nullable = false)
@@ -606,19 +598,11 @@ object SimpleAggregationTestUtils {
  */
 object FlatMapGroupsWithStateTestUtils {
 
-  /**
-   * @param stateVersion The state format version:
-   * @return A tuple of (keySchema, valueSchema)
-   */
   def getSchemas(stateVersion: Int): (StructType, StructType) = {
     val metadata = getSchemasWithMetadata(stateVersion)
     (metadata.keySchema, metadata.valueSchema)
   }
 
-  /**
-   * @param stateVersion The state format version
-   * @return ColumnFamilyMetadata with schema and KeyEncoderSpec
-   */
   def getSchemasWithMetadata(stateVersion: Int): ColumnFamilyMetadata = {
     val keySchema = StructType(Array(
       StructField("value", StringType, nullable = true)
@@ -654,17 +638,11 @@ object FlatMapGroupsWithStateTestUtils {
  */
 object SessionWindowTestUtils {
 
-  /**
-   * @return A tuple of (keySchema, valueSchema)
-   */
   def getSchemas(): (StructType, StructType) = {
     val metadata = getSchemasWithMetadata()
     (metadata.keySchema, metadata.valueSchema)
   }
 
-  /**
-   * @return ColumnFamilyMetadata with schema and KeyEncoderSpec
-   */
   def getSchemasWithMetadata(): ColumnFamilyMetadata = {
     val keySchema = StructType(Array(
       StructField("sessionId", StringType, nullable = false),
@@ -687,33 +665,23 @@ object SessionWindowTestUtils {
  * Test utility object providing schema definitions and constants for runStreamStreamJoinQuery
  */
 object StreamStreamJoinTestUtils {
-  // Column family names for keyToNumValues stores
-  val KEY_TO_NUM_VALUES_LEFT = "left-keyToNumValues"
-  val KEY_TO_NUM_VALUES_RIGHT = "right-keyToNumValues"
-  val KEY_TO_NUM_VALUES_ALL: Seq[String] = Seq(
-    KEY_TO_NUM_VALUES_LEFT,
-    KEY_TO_NUM_VALUES_RIGHT
-  )
-
-  // Column family names for keyWithIndexToValue stores
-  val KEY_WITH_INDEX_LEFT = "left-keyWithIndexToValue"
-  val KEY_WITH_INDEX_RIGHT = "right-keyWithIndexToValue"
-  val KEY_WITH_INDEX_ALL: Seq[String] = Seq(
-    KEY_WITH_INDEX_LEFT,
-    KEY_WITH_INDEX_RIGHT
-  )
-
-  /**
-   * @return A tuple of (keySchema, valueSchema)
-   */
+  // All state store names from SymmetricHashJoinStateManager
+  private val allStoreNames: Seq[String] =
+    SymmetricHashJoinStateManager.allStateStoreNames(LeftSide, RightSide)
+
+  // Column family names for keyToNumValues stores (derived from allStateStoreNames)
+  val KEY_TO_NUM_VALUES_ALL: Seq[String] =
+    allStoreNames.filter(_.endsWith(SymmetricHashJoinStateManager.KeyToNumValuesType.toString))
+
+  // Column family names for keyWithIndexToValue stores (derived from allStateStoreNames)
+  val KEY_WITH_INDEX_ALL: Seq[String] =
+    allStoreNames.filter(_.endsWith(SymmetricHashJoinStateManager.KeyWithIndexToValueType.toString))
+
   def getKeyToNumValuesSchemas(): (StructType, StructType) = {
     val metadata = getKeyToNumValuesSchemasWithMetadata()
     (metadata.keySchema, metadata.valueSchema)
   }
 
-  /**
-   * @return ColumnFamilyMetadata with schema and KeyEncoderSpec
-   */
   def getKeyToNumValuesSchemasWithMetadata(): ColumnFamilyMetadata = {
     val keySchema = StructType(Array(
       StructField("key", IntegerType)
@@ -725,19 +693,11 @@ object StreamStreamJoinTestUtils {
     ColumnFamilyMetadata(keySchema, valueSchema, encoderSpec)
   }
 
-  /**
-   * @param stateVersion The join state format version:
-   * @return A tuple of (keySchema, valueSchema)
-   */
   def getKeyWithIndexToValueSchemas(stateVersion: Int): (StructType, StructType) = {
     val metadata = getKeyWithIndexToValueSchemasWithMetadata(stateVersion)
     (metadata.keySchema, metadata.valueSchema)
   }
 
-  /**
-   * @param stateVersion The state format version
-   * @return ColumnFamilyMetadata with schema and KeyEncoderSpec
-   */
   def getKeyWithIndexToValueSchemasWithMetadata(stateVersion: Int): ColumnFamilyMetadata = {
     val keySchema = StructType(Array(
       StructField("key", IntegerType, nullable = false),
@@ -761,37 +721,23 @@ object StreamStreamJoinTestUtils {
     ColumnFamilyMetadata(keySchema, valueSchema, encoderSpec)
   }
 
-  /**
-   * Returns all schemas for stream-stream join V3 (multi-column family) in legacy 2-tuple format.
-   * V3 uses a single state store with multiple column families instead of separate stores.
-   *
-   * @return Map of column family name to (keySchema, valueSchema)
-   */
   def getJoinV3ColumnSchemaMap(): Map[String, (StructType, StructType)] = {
     getJoinV3ColumnSchemaMapWithMetadata().view.mapValues { metadata =>
       (metadata.keySchema, metadata.valueSchema)
     }.toMap
   }
 
-  /**
-   * @return Map of column family name to ColumnFamilyMetadata
-   */
   def getJoinV3ColumnSchemaMapWithMetadata(): Map[String, ColumnFamilyMetadata] = {
     val (keyToNumKeySchema, keyToNumValueSchema) = getKeyToNumValuesSchemas()
     val (keyWithIndexKeySchema, keyWithIndexValueSchema) = getKeyWithIndexToValueSchemas(3)
 
     val keyToNumEncoderSpec = NoPrefixKeyStateEncoderSpec(keyToNumKeySchema)
     val keyWithIndexEncoderSpec = NoPrefixKeyStateEncoderSpec(keyWithIndexKeySchema)
 
-    Map(
-      KEY_TO_NUM_VALUES_LEFT -> ColumnFamilyMetadata(
-        keyToNumKeySchema, keyToNumValueSchema, keyToNumEncoderSpec),
-      KEY_TO_NUM_VALUES_RIGHT -> ColumnFamilyMetadata(
-        keyToNumKeySchema, keyToNumValueSchema, keyToNumEncoderSpec),
-      KEY_WITH_INDEX_LEFT -> ColumnFamilyMetadata(
-        keyWithIndexKeySchema, keyWithIndexValueSchema, keyWithIndexEncoderSpec),
-      KEY_WITH_INDEX_RIGHT -> ColumnFamilyMetadata(
-        keyWithIndexKeySchema, keyWithIndexValueSchema, keyWithIndexEncoderSpec)
-    )
+    KEY_TO_NUM_VALUES_ALL.map(name => name -> ColumnFamilyMetadata(
+      keyToNumKeySchema, keyToNumValueSchema, keyToNumEncoderSpec)).toMap ++
+    KEY_WITH_INDEX_ALL.map(name => name -> ColumnFamilyMetadata(
+      keyWithIndexKeySchema, keyWithIndexValueSchema, keyWithIndexEncoderSpec)).toMap
+
   }
 }

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

@@ -25,13 +25,13 @@ import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
-import org.apache.spark.sql.execution.datasources.v2.state.utils.{EventTimeTimerProcessor, MultiStateVarProcessor, MultiStateVarProcessorTestUtils, TimerTestUtils}
 import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
 import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider, StateRepartitionUnsupportedProviderError, StateStore}
 import org.apache.spark.sql.functions.{col, count, sum, timestamp_seconds}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.{InputEvent, ListStateTTLProcessor, MapInputEvent, MapStateTTLProcessor, OutputMode, RunningCountStatefulProcessorWithProcTimeTimer, TimeMode, Trigger, TTLConfig, ValueStateTTLProcessor}
 import org.apache.spark.sql.streaming.util.{StreamManualClock, TTLProcessorUtils}
+import org.apache.spark.sql.streaming.util.{EventTimeTimerProcessor, MultiStateVarProcessor, MultiStateVarProcessorTestUtils, TimerTestUtils}
 import org.apache.spark.sql.types.{DataType, NullType, StructField, StructType}
 
 /**