address comment

addres comment 2

Author: zifeif2

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

@@ -19,20 +19,22 @@ package org.apache.spark.sql.execution.streaming.state
 import java.util.UUID
 
 import scala.collection.MapView
-import scala.collection.immutable.HashMap
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.execution.streaming.operators.stateful.StatefulOperatorsUtils
 import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.StateStoreColumnFamilySchemaUtils
 import org.apache.spark.sql.execution.streaming.runtime.StreamingCheckpointConstants.DIR_NAME_STATE
+import org.apache.spark.sql.internal.SQLConf
 
 case class StatePartitionWriterColumnFamilyInfo(
   schema: StateStoreColFamilySchema,
   // set this to true if state variable is ListType in TransformWithState
   useMultipleValuesPerKey: Boolean = false)
+
 /**
  * A writer that can directly write binary data to the streaming state store.
  *
@@ -52,11 +54,24 @@ class StatePartitionAllColumnFamiliesWriter(
     operatorId: Int,
     storeName: String,
     currentBatchId: Long,
-    columnFamilyToSchemaMap: HashMap[String, StatePartitionWriterColumnFamilyInfo]) {
+    columnFamilyToSchemaMap: Map[String, StatePartitionWriterColumnFamilyInfo],
+    operatorName: String,
+    schemaProviderOpt: Option[StateSchemaProvider],
+    sqlConf: Map[String, String]) {
+
+  private def isJoinV3Operator(
+      operatorName: String, sqlConf: Map[String, String]): Boolean = {
+    operatorName ==
+        StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
+        sqlConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION.key) == "3"
+  }
   private val defaultSchema = {
     columnFamilyToSchemaMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) match {
       case Some(info) => info.schema
       case None =>
+        assert(isJoinV3Operator(operatorName, sqlConf),
+          s"Please provide the schema of 'default' column family in StateStoreColFamilySchema" +
+            s"for operator $operatorName")
         // Return a dummy StateStoreColFamilySchema if not found
         val placeholderSchema = columnFamilyToSchemaMap.head._2.schema
         StateStoreColFamilySchema(
@@ -68,7 +83,7 @@ class StatePartitionAllColumnFamiliesWriter(
           keyStateEncoderSpec = Option(NoPrefixKeyStateEncoderSpec(placeholderSchema.keySchema)))
     }
   }
-
+  private val useColumnFamilies = columnFamilyToSchemaMap.size > 1
   private val columnFamilyToKeySchemaLenMap: MapView[String, Int] =
     columnFamilyToSchemaMap.view.mapValues(_.schema.keySchema.length)
   private val columnFamilyToValueSchemaLenMap: MapView[String, Int] =
@@ -80,12 +95,12 @@ class StatePartitionAllColumnFamiliesWriter(
       operatorId, partitionId, storeName)
     val stateStoreProviderId = StateStoreProviderId(stateStoreId, UUID.randomUUID())
 
-    val useColumnFamilies = columnFamilyToSchemaMap.size > 1
+
     val provider = StateStoreProvider.createAndInit(
       stateStoreProviderId, defaultSchema.keySchema, defaultSchema.valueSchema,
       defaultSchema.keyStateEncoderSpec.get,
       useColumnFamilies = useColumnFamilies, storeConf, hadoopConf,
-      useMultipleValuesPerKey = false, stateSchemaProvider = None)
+      useMultipleValuesPerKey = false, stateSchemaProvider = schemaProviderOpt)
     provider
   }
 
@@ -101,7 +116,7 @@ class StatePartitionAllColumnFamiliesWriter(
       stateStoreCkptId = None,
       loadEmpty = true
     )
-    if (columnFamilyToSchemaMap.size > 1) {
+    if (useColumnFamilies) {
       columnFamilyToSchemaMap.foreach { pair =>
         val colFamilyName = pair._1
         val cfSchema = pair._2.schema
@@ -117,7 +132,7 @@ class StatePartitionAllColumnFamiliesWriter(
               cfSchema.keySchema,
               cfSchema.valueSchema,
               cfSchema.keyStateEncoderSpec.get,
-              columnFamilyToSchemaMap(colFamilyName).useMultipleValuesPerKey,
+              pair._2.useMultipleValuesPerKey,
               isInternal)
         }
       }