address comment

addres comment 2

Author: zifeif2

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

@@ -44,16 +44,6 @@ object SchemaUtil {
     }
   }
 
-  def getScanAllColumnFamiliesSchema(keySchema: StructType): StructType = {
-    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("key_bytes", BinaryType)
-      .add("value_bytes", BinaryType)
-      .add("column_family_name", StringType)
-  }
-
   def getSourceSchema(
       sourceOptions: StateSourceOptions,
       keySchema: StructType,
@@ -72,7 +62,13 @@ object SchemaUtil {
         .add("value", valueSchema)
         .add("partition_id", IntegerType)
     } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
-      getScanAllColumnFamiliesSchema(keySchema)
+      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("key_bytes", BinaryType)
+        .add("value_bytes", BinaryType)
+        .add("column_family_name", StringType)
     } else {
       new StructType()
         .add("key", keySchema)

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateStoreColumnFamilySchemaUtils.scala

@@ -264,8 +264,4 @@ object StateStoreColumnFamilySchemaUtils {
       valSchema,
       Some(RangeKeyScanStateEncoderSpec(keySchema, Seq(0))))
   }
-
-  def isInternalColFamily(name: String): Boolean = {
-    name.startsWith("$")
-  }
 }

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,13 +54,26 @@ class StatePartitionAllColumnFamiliesWriter(
     operatorId: Int,
     storeName: String,
     currentBatchId: Long,
-    columnFamilyToSchemaMap: HashMap[String, StatePartitionWriterColumnFamilyInfo]) {
+    colFamilyToWriterInfoMap: 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 {
+    colFamilyToWriterInfoMap.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
+        val placeholderSchema = colFamilyToWriterInfoMap.head._2.schema
         StateStoreColFamilySchema(
           colFamilyName = "__dummy__",
           keySchemaId = 0,
@@ -69,23 +84,23 @@ class StatePartitionAllColumnFamiliesWriter(
     }
   }
 
+  private val useColumnFamilies = colFamilyToWriterInfoMap.size > 1
   private val columnFamilyToKeySchemaLenMap: MapView[String, Int] =
-    columnFamilyToSchemaMap.view.mapValues(_.schema.keySchema.length)
+    colFamilyToWriterInfoMap.view.mapValues(_.schema.keySchema.length)
   private val columnFamilyToValueSchemaLenMap: MapView[String, Int] =
-    columnFamilyToSchemaMap.view.mapValues(_.schema.valueSchema.length)
+    colFamilyToWriterInfoMap.view.mapValues(_.schema.valueSchema.length)
 
   protected lazy val provider: StateStoreProvider = {
     val stateCheckpointLocation = new Path(targetCpLocation, DIR_NAME_STATE).toString
     val stateStoreId = StateStoreId(stateCheckpointLocation,
       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,23 +116,22 @@ class StatePartitionAllColumnFamiliesWriter(
       stateStoreCkptId = None,
       loadEmpty = true
     )
-    if (columnFamilyToSchemaMap.size > 1) {
-      columnFamilyToSchemaMap.foreach { pair =>
+    if (useColumnFamilies) {
+      colFamilyToWriterInfoMap.foreach { pair =>
         val colFamilyName = pair._1
         val cfSchema = pair._2.schema
         colFamilyName match {
           case StateStore.DEFAULT_COL_FAMILY_NAME => // createAndInit has registered default
           case _ =>
-            val isInternal = StateStoreColumnFamilySchemaUtils.isInternalColFamily(colFamilyName)
-
             require(cfSchema.keyStateEncoderSpec.isDefined,
               s"keyStateEncoderSpec must be defined for column family ${cfSchema.colFamilyName}")
+            val isInternal = StateStoreColumnFamilySchemaUtils.isInternalColFamily(colFamilyName)
             store.createColFamilyIfAbsent(
               colFamilyName,
               cfSchema.keySchema,
               cfSchema.valueSchema,
               cfSchema.keyStateEncoderSpec.get,
-              columnFamilyToSchemaMap(colFamilyName).useMultipleValuesPerKey,
+              pair._2.useMultipleValuesPerKey,
               isInternal)
         }
       }
@@ -159,10 +173,14 @@ class StatePartitionAllColumnFamiliesWriter(
     val valueRow = new UnsafeRow(columnFamilyToValueSchemaLenMap(colFamilyName))
     valueRow.pointTo(valueBytes, valueBytes.length)
 
-    if (columnFamilyToSchemaMap(colFamilyName).useMultipleValuesPerKey) {
+    if (colFamilyToWriterInfoMap(colFamilyName).useMultipleValuesPerKey) {
       // if a column family useMultipleValuesPerKey (e.g. ListType), we will
       // write with 1 put followed by merge
-      stateStore.merge(keyRow, valueRow, colFamilyName)
+      if (stateStore.keyExists(keyRow, colFamilyName)) {
+        stateStore.merge(keyRow, valueRow, colFamilyName)
+      } else {
+        stateStore.put(keyRow, valueRow, colFamilyName)
+      }
     } else {
       stateStore.put(keyRow, valueRow, colFamilyName)
     }