fix rebase error

Author: zifeif2

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

@@ -44,6 +44,16 @@ 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,
@@ -62,13 +72,7 @@ object SchemaUtil {
         .add("value", valueSchema)
         .add("partition_id", IntegerType)
     } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
-      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)
+      getScanAllColumnFamiliesSchema(keySchema)
     } else {
       new StructType()
         .add("key", keySchema)

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

@@ -53,20 +53,19 @@ class StatePartitionAllColumnFamiliesWriter(
     storeName: String,
     currentBatchId: Long,
     columnFamilyToSchemaMap: HashMap[String, StatePartitionWriterColumnFamilyInfo]) {
-  private val dummySchema: StructType =
-    StructType(Array(StructField("__dummy__", NullType)))
   private val defaultSchema = {
     columnFamilyToSchemaMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) match {
       case Some(info) => info.schema
       case None =>
         // Return a dummy StateStoreColFamilySchema if not found
+        val placeholderSchema = columnFamilyToSchemaMap.head._2.schema
         StateStoreColFamilySchema(
           colFamilyName = "__dummy__",
           keySchemaId = 0,
-          keySchema = dummySchema,
+          keySchema = placeholderSchema.keySchema,
           valueSchemaId = 0,
-          valueSchema = dummySchema,
-          keyStateEncoderSpec = Option(NoPrefixKeyStateEncoderSpec(dummySchema)))
+          valueSchema = placeholderSchema.valueSchema,
+          keyStateEncoderSpec = Option(NoPrefixKeyStateEncoderSpec(placeholderSchema.keySchema)))
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionAllColumnFamiliesWriterSuite.scala

@@ -24,7 +24,9 @@ import scala.collection.immutable.HashMap
 import org.apache.spark.TaskContext
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
-import org.apache.spark.sql.execution.datasources.v2.state.{StateDataSourceTestBase, StateSourceOptions}
+import org.apache.spark.sql.execution.datasources.v2.state.{EventTimeTimerProcessor, ListStateTTLProcessor, MultiStateVarProcessor, RunningCountStatefulProcessorWithProcTimeTimer, StateDataSourceTestBase, StateSourceOptions}
+import org.apache.spark.sql.execution.datasources.v2.state.utils.SchemaUtil
+import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.timers.TimerStateUtils
 import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryCheckpointMetadata}
 import org.apache.spark.sql.execution.streaming.utils.StreamingUtils
 import org.apache.spark.sql.functions.{col, timestamp_seconds}