Introducing OperatorStateMetadataV2 for the TransformWithState operator

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala

@@ -182,6 +182,14 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
   /** Specifies sort order for each partition requirements on the input data for this operator. */
   def requiredChildOrdering: Seq[Seq[SortOrder]] = Seq.fill(children.size)(Nil)
 
+  /**
+   * This method will be called at the end of a MicrobatchExecution
+   * to write the metadata of the operator to the checkpoint file.
+   */
+  def writeOperatorStateMetadata(): Unit = {
+    children.foreach(_.writeOperatorStateMetadata())
+  }
+
   /**
    * Returns the result of this query as an RDD[InternalRow] by delegating to `doExecute` after
    * preparations.

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

@@ -32,7 +32,7 @@ import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionRead
 import org.apache.spark.sql.execution.datasources.v2.state.StateDataSourceErrors
 import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.PATH
 import org.apache.spark.sql.execution.streaming.CheckpointFileManager
-import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadata, OperatorStateMetadataReader, OperatorStateMetadataV1}
+import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadata, OperatorStateMetadataReader, OperatorStateMetadataV1, OperatorStateMetadataV2}
 import org.apache.spark.sql.sources.DataSourceRegister
 import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StringType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -46,6 +46,7 @@ case class StateMetadataTableEntry(
     numPartitions: Int,
     minBatchId: Long,
     maxBatchId: Long,
+    operatorPropertiesJson: String,
     numColsPrefixKey: Int) {
   def toRow(): InternalRow = {
     new GenericInternalRow(
@@ -55,7 +56,9 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
-        numColsPrefixKey))
+        UTF8String.fromString(operatorPropertiesJson),
+        numColsPrefixKey
+      ))
   }
 }
 
@@ -68,6 +71,7 @@ object StateMetadataTableEntry {
       .add("numPartitions", IntegerType)
       .add("minBatchId", LongType)
       .add("maxBatchId", LongType)
+      .add("operatorProperties", StringType)
   }
 }
 
@@ -199,17 +203,21 @@ class StateMetadataPartitionReader(
 
   private[state] lazy val stateMetadata: Iterator[StateMetadataTableEntry] = {
     allOperatorStateMetadata.flatMap { operatorStateMetadata =>
-      require(operatorStateMetadata.version == 1)
-      val operatorStateMetadataV1 = operatorStateMetadata.asInstanceOf[OperatorStateMetadataV1]
-      operatorStateMetadataV1.stateStoreInfo.map { stateStoreMetadata =>
-        StateMetadataTableEntry(operatorStateMetadataV1.operatorInfo.operatorId,
-          operatorStateMetadataV1.operatorInfo.operatorName,
+      require(operatorStateMetadata.version == 1 || operatorStateMetadata.version == 2)
+      val operatorProperties = operatorStateMetadata match {
+        case _: OperatorStateMetadataV1 => ""
+        case v2: OperatorStateMetadataV2 => v2.operatorPropertiesJson
+      }
+      operatorStateMetadata.stateStoreInfo.map { stateStoreMetadata =>
+          StateMetadataTableEntry(operatorStateMetadata.operatorInfo.operatorId,
+          operatorStateMetadata.operatorInfo.operatorName,
           stateStoreMetadata.storeName,
           stateStoreMetadata.numPartitions,
           if (batchIds.nonEmpty) batchIds.head else -1,
           if (batchIds.nonEmpty) batchIds.last else -1,
+          operatorProperties,
           stateStoreMetadata.numColsPrefixKey
-        )
+          )
       }
     }
   }.iterator

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

@@ -37,7 +37,7 @@ import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadat
 import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
 import org.apache.spark.sql.execution.python.FlatMapGroupsInPandasWithStateExec
 import org.apache.spark.sql.execution.streaming.sources.WriteToMicroBatchDataSourceV1
-import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadataV1, OperatorStateMetadataWriter}
+import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadataV1, OperatorStateMetadataV2, OperatorStateMetadataWriter}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.OutputMode
 import org.apache.spark.util.{SerializableConfiguration, Utils}
@@ -189,12 +189,15 @@ class IncrementalExecution(
 
   object WriteStatefulOperatorMetadataRule extends SparkPlanPartialRule {
     override val rule: PartialFunction[SparkPlan, SparkPlan] = {
-      case stateStoreWriter: StateStoreWriter if isFirstBatch =>
-        val metadata = stateStoreWriter.operatorStateMetadata()
-        val metadataWriter = new OperatorStateMetadataWriter(new Path(
-          checkpointLocation, stateStoreWriter.getStateInfo.operatorId.toString), hadoopConf)
-        metadataWriter.write(metadata)
-        stateStoreWriter
+      case stateStoreWriter: StateStoreWriter
+        // We only want to write the OperatorStateMetadata via this rule if
+        // we are in the first batch and the state format version is 1.
+        if isFirstBatch && stateStoreWriter.operatorStateMetadataVersion == 1 =>
+          val metadata = stateStoreWriter.operatorStateMetadata()
+          val metadataWriter = new OperatorStateMetadataWriter(new Path(
+            checkpointLocation, stateStoreWriter.getStateInfo.operatorId.toString), hadoopConf)
+          metadataWriter.write(metadata)
+          stateStoreWriter
     }
   }
 
@@ -434,9 +437,10 @@ class IncrementalExecution(
               new SerializableConfiguration(hadoopConf))
             val opMetadataList = reader.allOperatorStateMetadata
             ret = opMetadataList.map { operatorMetadata =>
-              val metadataInfoV1 = operatorMetadata
-                .asInstanceOf[OperatorStateMetadataV1]
-                .operatorInfo
+              val metadataInfoV1 = operatorMetadata match {
+                case v1: OperatorStateMetadataV1 => v1.operatorInfo
+                case v2: OperatorStateMetadataV2 => v2.operatorInfo
+              }
               metadataInfoV1.operatorId -> metadataInfoV1.operatorName
             }.toMap
           } catch {

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

@@ -386,7 +386,12 @@ class MicroBatchExecution(
       execCtx.carryOverExecStatsOnLatestExecutedBatch()
       // Must be outside reportTimeTaken so it is recorded
       if (execCtx.isCurrentBatchConstructed) {
-        execCtx.finishTrigger(currentBatchHasNewData, execCtx.executionPlan, execCtx.batchId)
+        execCtx.finishTrigger(
+          currentBatchHasNewData,
+          execCtx.executionPlan,
+          execCtx.batchId,
+          execCtx.previousContext
+        )
       } else {
         execCtx.finishNoExecutionTrigger(execCtx.batchId)
       }

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

@@ -264,11 +264,20 @@ abstract class ProgressContext(
       hasNewData: Boolean,
       sourceToNumInputRowsMap: Map[SparkDataStream, Long],
       lastExecution: IncrementalExecution,
-      lastEpochId: Long): Unit = {
+      lastEpochId: Long,
+      previousExecution: Option[MicroBatchExecutionContext]): Unit = {
     assert(
       currentTriggerStartOffsets != null && currentTriggerEndOffsets != null &&
         currentTriggerLatestOffsets != null
     )
+
+    previousExecution match {
+      case Some(prevContext) => if (!prevContext.runId.equals(lastExecution.runId)) {
+        lastExecution.executedPlan.writeOperatorStateMetadata()
+      }
+      case None => lastExecution.executedPlan.writeOperatorStateMetadata()
+    }
+
     currentTriggerEndTimestamp = triggerClock.getTimeMillis()
     val processingTimeMills = currentTriggerEndTimestamp - currentTriggerStartTimestamp
     assert(lastExecution != null, "executed batch should provide the information for execution.")
@@ -379,10 +388,11 @@ abstract class ProgressContext(
   def finishTrigger(
       hasNewData: Boolean,
       lastExecution: IncrementalExecution,
-      lastEpoch: Long): Unit = {
+      lastEpoch: Long,
+      previousContext: Option[MicroBatchExecutionContext] = None): Unit = {
     val map: Map[SparkDataStream, Long] =
       if (hasNewData) extractSourceToNumInputRows(lastExecution) else Map.empty
-    finishTrigger(hasNewData, map, lastExecution, lastEpoch)
+    finishTrigger(hasNewData, map, lastExecution, lastEpoch, previousContext)
   }
 
   private def warnIfFinishTriggerTakesTooLong(

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

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.streaming
+
+import org.json4s.{DefaultFormats, Formats, JValue}
+import org.json4s.JsonAST.{JBool, JString}
+import org.json4s.JsonDSL._
+
+// Enum to store the types of state variables we support
+sealed trait StateVariableType
+
+case object ValueState extends StateVariableType
+case object ListState extends StateVariableType
+case object MapState extends StateVariableType
+
+// This object is used to convert the state type from string to the corresponding enum
+object StateVariableType {
+  def withName(name: String): StateVariableType = name match {
+    case "ValueState" => ValueState
+    case "ListState" => ListState
+    case "MapState" => MapState
+    case _ => throw new IllegalArgumentException(s"Unknown state type: $name")
+  }
+}
+
+// This class is used to store the information about a state variable.
+// It is stored in operatorProperties for the TransformWithStateExec operator
+// to be able to validate that the State Variables are the same across restarts.
+class StateVariableInfo(
+   val stateName: String,
+   val stateType: StateVariableType,
+   val isTtlEnabled: Boolean
+ ) {
+  def jsonValue: JValue = {
+    ("stateName" -> JString(stateName)) ~
+      ("stateType" -> JString(stateType.toString)) ~
+      ("isTtlEnabled" -> JBool(isTtlEnabled))
+  }
+}
+
+// This object is used to convert the state variable information
+// from JSON to a list of StateVariableInfo
+object StateVariableInfo {
+  implicit val formats: Formats = DefaultFormats
+  def fromJson(json: Any): List[StateVariableInfo] = {
+    assert(json.isInstanceOf[List[_]], s"Expected List but got ${json.getClass}")
+    val stateVariables = json.asInstanceOf[List[Map[String, Any]]]
+    // Extract each JValue to StateVariableInfo
+    stateVariables.map { stateVariable =>
+      new StateVariableInfo(
+        stateVariable("stateName").asInstanceOf[String],
+        StateVariableType.withName(stateVariable("stateType").asInstanceOf[String]),
+        stateVariable("isTtlEnabled").asInstanceOf[Boolean]
+      )
+    }
+  }
+}

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

@@ -94,6 +94,9 @@ class StatefulProcessorHandleImpl(
    */
   private[sql] val ttlStates: util.List[TTLState] = new util.ArrayList[TTLState]()
 
+  private[sql] val stateVariables: util.List[StateVariableInfo] =
+    new util.ArrayList[StateVariableInfo]()
+
   private val BATCH_QUERY_ID = "00000000-0000-0000-0000-000000000000"
 
   private def buildQueryInfo(): QueryInfo = {
@@ -129,6 +132,7 @@ class StatefulProcessorHandleImpl(
       stateName: String,
       valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
+    stateVariables.add(new StateVariableInfo(stateName, ValueState, false))
     incrementMetric("numValueStateVars")
     val resultState = new ValueStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -139,6 +143,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[T],
       ttlConfig: TTLConfig): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
+    stateVariables.add(new StateVariableInfo(stateName, ValueState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -240,6 +245,7 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state")
+    stateVariables.add(new StateVariableInfo(stateName, ListState, false))
     incrementMetric("numListStateVars")
     val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -266,6 +272,7 @@ class StatefulProcessorHandleImpl(
       ttlConfig: TTLConfig): ListState[T] = {
 
     verifyStateVarOperations("get_list_state")
+    stateVariables.add(new StateVariableInfo(stateName, ListState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -282,6 +289,7 @@ class StatefulProcessorHandleImpl(
       userKeyEnc: Encoder[K],
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
+    stateVariables.add(new StateVariableInfo(stateName, MapState, false))
     incrementMetric("numMapStateVars")
     val resultState = new MapStateImpl[K, V](store, stateName, keyEncoder, userKeyEnc, valEncoder)
     resultState
@@ -293,6 +301,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[V],
       ttlConfig: TTLConfig): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
+    stateVariables.add(new StateVariableInfo(stateName, MapState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)

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

@@ -118,7 +118,7 @@ class ContinuousExecutionContext(
  */
 class MicroBatchExecutionContext(
     id: UUID,
-    runId: UUID,
+    val runId: UUID,
     name: String,
     triggerClock: Clock,
     sources: Seq[SparkDataStream],

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

@@ -19,6 +19,15 @@ package org.apache.spark.sql.execution.streaming
 import java.util.UUID
 import java.util.concurrent.TimeUnit.NANOSECONDS
 
+import scala.jdk.CollectionConverters.CollectionHasAsScala
+
+import org.apache.hadoop.fs.Path
+import org.json4s.{DefaultFormats, JArray, JString}
+import org.json4s.JsonAST.JValue
+import org.json4s.JsonDSL._
+import org.json4s.jackson.JsonMethods
+import org.json4s.jackson.JsonMethods.{compact, render}
+
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
@@ -77,6 +86,44 @@ case class TransformWithStateExec(
 
   override def shortName: String = "transformWithStateExec"
 
+  val operatorProperties: OperatorProperties =
+    OperatorProperties.create(
+      sparkContext,
+      "colFamilyMetadata"
+    )
+
+  /** Metadata of this stateful operator and its states stores. */
+  override def operatorStateMetadata(): OperatorStateMetadata = {
+    val info = getStateInfo
+    val operatorInfo = OperatorInfoV1(info.operatorId, shortName)
+    val stateStoreInfo =
+      Array(StateStoreMetadataV1(StateStoreId.DEFAULT_STORE_NAME, 0, info.numPartitions))
+
+    val operatorPropertiesJson: JValue = ("timeMode" -> JString(timeMode.toString)) ~
+      ("outputMode" -> JString(outputMode.toString)) ~
+      ("stateVariables" -> operatorProperties.value.get("stateVariables"))
+
+    val json = compact(render(operatorPropertiesJson))
+    OperatorStateMetadataV2(operatorInfo, stateStoreInfo, json)
+  }
+
+  override def operatorStateMetadataVersion: Int = 2
+
+  /**
+   * This method will be called at the end of a MicrobatchExecution
+   * to write the metadata of the operator to the checkpoint file.
+   */
+  override def writeOperatorStateMetadata(): Unit = {
+    val metadata = operatorStateMetadata()
+    val metadataWriter = new OperatorStateMetadataWriter(
+      new Path(stateInfo.get.checkpointLocation,
+        getStateInfo.operatorId.toString),
+      session.sqlContext.sessionState.newHadoopConf()
+    )
+    metadataWriter.write(metadata)
+    super.writeOperatorStateMetadata()
+  }
+
   override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = {
     if (timeMode == ProcessingTime) {
       // TODO: check if we can return true only if actual timers are registered, or there is
@@ -306,6 +353,9 @@ case class TransformWithStateExec(
           store.abort()
         }
       }
+      operatorProperties.add(Map
+        ("stateVariables" -> JArray(processorHandle.stateVariables.
+          asScala.map(_.jsonValue).toList)))
       setStoreMetrics(store)
       setOperatorMetrics()
       statefulProcessor.close()
@@ -561,6 +611,14 @@ object TransformWithStateExec {
       initialStateDeserializer,
       initialState)
   }
+
+  def deserializeOperatorProperties(json: String): Map[String, Any] = {
+    val parsedJson = JsonMethods.parse(json)
+
+    implicit val formats = DefaultFormats
+    val deserializedMap: Map[String, Any] = parsedJson.extract[Map[String, Any]]
+    deserializedMap
+  }
 }
 // scalastyle:on argcount