From ac41ba5946b08c9d4499223410221a2a6986962d Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 20 May 2024 10:36:22 -0700
Subject: [PATCH 01/22] can write metadata

---
 .../spark/sql/execution/SparkPlan.scala       |  8 +++
 .../streaming/ProgressReporter.scala          |  4 ++
 .../StatefulProcessorHandleImpl.scala         | 10 ++++
 .../streaming/TransformWithStateExec.scala    | 52 ++++++++++++++++++-
 4 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
index 7bc770a0c9e33..2cb7788f48f99 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
+++ b/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.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
index 3842ed5743558..748f54af89371 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
@@ -269,6 +269,10 @@ abstract class ProgressContext(
       currentTriggerStartOffsets != null && currentTriggerEndOffsets != null &&
         currentTriggerLatestOffsets != null
     )
+    if (lastExecution.isFirstBatch) {
+      lastExecution.executedPlan.writeOperatorStateMetadata()
+    }
+
     currentTriggerEndTimestamp = triggerClock.getTimeMillis()
     val processingTimeMills = currentTriggerEndTimestamp - currentTriggerStartTimestamp
     assert(lastExecution != null, "executed batch should provide the information for execution.")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index dcc77e94de280..48779e0db0353 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.streaming
 import java.util
 import java.util.UUID
 
+import scala.collection.mutable
+
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
@@ -94,6 +96,8 @@ class StatefulProcessorHandleImpl(
    */
   private[sql] val ttlStates: util.List[TTLState] = new util.ArrayList[TTLState]()
 
+  val operatorMetadata: mutable.Map[String, String] = mutable.Map.empty
+
   private val BATCH_QUERY_ID = "00000000-0000-0000-0000-000000000000"
 
   private def buildQueryInfo(): QueryInfo = {
@@ -129,6 +133,7 @@ class StatefulProcessorHandleImpl(
       stateName: String,
       valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
+    operatorMetadata.put(stateName, "VALUE, FALSE")
     incrementMetric("numValueStateVars")
     val resultState = new ValueStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -139,6 +144,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[T],
       ttlConfig: TTLConfig): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
+    operatorMetadata.put(stateName, "VALUE, TRUE")
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -240,6 +246,7 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state")
+    operatorMetadata.put(stateName, "LIST, FALSE")
     incrementMetric("numListStateVars")
     val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -266,6 +273,7 @@ class StatefulProcessorHandleImpl(
       ttlConfig: TTLConfig): ListState[T] = {
 
     verifyStateVarOperations("get_list_state")
+    operatorMetadata.put(stateName, "LIST, TRUE")
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -282,6 +290,7 @@ class StatefulProcessorHandleImpl(
       userKeyEnc: Encoder[K],
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
+    operatorMetadata.put(stateName, "MAP, FALSE")
     incrementMetric("numMapStateVars")
     val resultState = new MapStateImpl[K, V](store, stateName, keyEncoder, userKeyEnc, valEncoder)
     resultState
@@ -293,6 +302,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[V],
       ttlConfig: TTLConfig): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
+    operatorMetadata.put(stateName, "MAP, TRUE")
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index fbd062acfb5dd..bba5727e9c2f4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.streaming
 import java.util.UUID
 import java.util.concurrent.TimeUnit.NANOSECONDS
 
+import org.apache.spark.SparkContext
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
@@ -32,8 +33,42 @@ import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSch
 import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
-import org.apache.spark.util.{CompletionIterator, SerializableConfiguration, Utils}
+import org.apache.spark.util.{AccumulatorV2, CompletionIterator, SerializableConfiguration, Utils}
 
+class ColFamilyMetadata(initValue: Map[String, String] = Map.empty)
+  extends AccumulatorV2[Map[String, String], Map[String, String]] {
+
+  private var _value: Map[String, String] = initValue
+
+  override def isZero: Boolean = _value.isEmpty
+
+  override def copy(): AccumulatorV2[Map[String, String], Map[String, String]] = {
+    val newAcc = new ColFamilyMetadata
+    newAcc._value = _value
+    newAcc
+  }
+
+  override def reset(): Unit = _value = Map.empty[String, String]
+
+  override def add(v: Map[String, String]): Unit = _value ++= v
+
+  override def merge(other: AccumulatorV2[Map[String, String], Map[String, String]]): Unit = {
+    _value ++= other.value
+  }
+
+  override def value: Map[String, String] = _value
+}
+
+object ColFamilyMetadata {
+  def create(
+      sc: SparkContext,
+      name: String,
+      initValue: Map[String, String] = Map.empty): ColFamilyMetadata = {
+    val acc = new ColFamilyMetadata(initValue)
+    acc.register(sc, name = Some(name))
+    acc
+  }
+}
 /**
  * Physical operator for executing `TransformWithState`
  *
@@ -77,6 +112,20 @@ case class TransformWithStateExec(
 
   override def shortName: String = "transformWithStateExec"
 
+  val colFamilyMetadataAccumulator: ColFamilyMetadata =
+    ColFamilyMetadata.create(sparkContext, "colFamilyMetadata")
+
+  /**
+   * 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 = {
+    colFamilyMetadataAccumulator.value.foreach { case (key, value) =>
+      logError(s"### $key: $value")
+    }
+    children.foreach(_.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 +355,7 @@ case class TransformWithStateExec(
           store.abort()
         }
       }
+      colFamilyMetadataAccumulator.add(processorHandle.operatorMetadata.toMap)
       setStoreMetrics(store)
       setOperatorMetrics()
       statefulProcessor.close()

From 62b1aac265aa53671cc1ed8804e700908559fbf1 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 20 May 2024 12:14:15 -0700
Subject: [PATCH 02/22] operatorstatemetadata

---
 .../StatefulProcessorHandleImpl.scala         | 15 ++++---
 .../streaming/TransformWithStateExec.scala    | 41 ++++++++++++++----
 .../state/OperatorStateMetadata.scala         | 43 ++++++++++++++++---
 3 files changed, 77 insertions(+), 22 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 48779e0db0353..9827020a7eac8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -96,7 +96,8 @@ class StatefulProcessorHandleImpl(
    */
   private[sql] val ttlStates: util.List[TTLState] = new util.ArrayList[TTLState]()
 
-  val operatorMetadata: mutable.Map[String, String] = mutable.Map.empty
+  var stateVariables: mutable.ListBuffer[StateVariableInfo] =
+    mutable.ListBuffer.empty[StateVariableInfo]
 
   private val BATCH_QUERY_ID = "00000000-0000-0000-0000-000000000000"
 
@@ -133,7 +134,7 @@ class StatefulProcessorHandleImpl(
       stateName: String,
       valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
-    operatorMetadata.put(stateName, "VALUE, FALSE")
+    stateVariables.append(new StateVariableInfo(stateName, ValueState, false))
     incrementMetric("numValueStateVars")
     val resultState = new ValueStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -144,7 +145,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[T],
       ttlConfig: TTLConfig): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
-    operatorMetadata.put(stateName, "VALUE, TRUE")
+    stateVariables.append(new StateVariableInfo(stateName, ValueState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -246,7 +247,7 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state")
-    operatorMetadata.put(stateName, "LIST, FALSE")
+    stateVariables.append(new StateVariableInfo(stateName, ListState, false))
     incrementMetric("numListStateVars")
     val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -273,7 +274,7 @@ class StatefulProcessorHandleImpl(
       ttlConfig: TTLConfig): ListState[T] = {
 
     verifyStateVarOperations("get_list_state")
-    operatorMetadata.put(stateName, "LIST, TRUE")
+    stateVariables.append(new StateVariableInfo(stateName, ListState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -290,7 +291,7 @@ class StatefulProcessorHandleImpl(
       userKeyEnc: Encoder[K],
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
-    operatorMetadata.put(stateName, "MAP, FALSE")
+    stateVariables.append(new StateVariableInfo(stateName, MapState, false))
     incrementMetric("numMapStateVars")
     val resultState = new MapStateImpl[K, V](store, stateName, keyEncoder, userKeyEnc, valEncoder)
     resultState
@@ -302,7 +303,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[V],
       ttlConfig: TTLConfig): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
-    operatorMetadata.put(stateName, "MAP, TRUE")
+    stateVariables.append(new StateVariableInfo(stateName, MapState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index bba5727e9c2f4..6cb03f76cf541 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -35,7 +35,22 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
 import org.apache.spark.util.{AccumulatorV2, CompletionIterator, SerializableConfiguration, Utils}
 
-class ColFamilyMetadata(initValue: Map[String, String] = Map.empty)
+// define state variable type enum with types VALUE, LIST, MAP
+sealed trait StateVariableType
+case object ValueState extends StateVariableType
+case object ListState extends StateVariableType
+case object MapState extends StateVariableType
+class StateVariableInfo(
+    val stateName: String,
+    val stateType: StateVariableType,
+    val isTtlEnabled: Boolean
+) {
+    override def toString: String = {
+      s"StateVariableInfo(stateName=$stateName, stateType=$stateType, isTtlEnabled=$isTtlEnabled)"
+    }
+}
+
+class OperatorProperties(initValue: Map[String, String] = Map.empty)
   extends AccumulatorV2[Map[String, String], Map[String, String]] {
 
   private var _value: Map[String, String] = initValue
@@ -43,7 +58,7 @@ class ColFamilyMetadata(initValue: Map[String, String] = Map.empty)
   override def isZero: Boolean = _value.isEmpty
 
   override def copy(): AccumulatorV2[Map[String, String], Map[String, String]] = {
-    val newAcc = new ColFamilyMetadata
+    val newAcc = new OperatorProperties
     newAcc._value = _value
     newAcc
   }
@@ -59,12 +74,12 @@ class ColFamilyMetadata(initValue: Map[String, String] = Map.empty)
   override def value: Map[String, String] = _value
 }
 
-object ColFamilyMetadata {
+object OperatorProperties {
   def create(
       sc: SparkContext,
       name: String,
-      initValue: Map[String, String] = Map.empty): ColFamilyMetadata = {
-    val acc = new ColFamilyMetadata(initValue)
+      initValue: Map[String, String] = Map.empty): OperatorProperties = {
+    val acc = new OperatorProperties(initValue)
     acc.register(sc, name = Some(name))
     acc
   }
@@ -112,15 +127,22 @@ case class TransformWithStateExec(
 
   override def shortName: String = "transformWithStateExec"
 
-  val colFamilyMetadataAccumulator: ColFamilyMetadata =
-    ColFamilyMetadata.create(sparkContext, "colFamilyMetadata")
+  val operatorProperties: OperatorProperties =
+    OperatorProperties.create(
+      sparkContext,
+      "colFamilyMetadata",
+      Map(
+        "timeMode" -> timeMode.toString,
+        "outputMode" -> outputMode.toString
+      )
+    )
 
   /**
    * 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 = {
-    colFamilyMetadataAccumulator.value.foreach { case (key, value) =>
+    operatorProperties.value.foreach { case (key, value) =>
       logError(s"### $key: $value")
     }
     children.foreach(_.writeOperatorStateMetadata())
@@ -355,7 +377,8 @@ case class TransformWithStateExec(
           store.abort()
         }
       }
-      colFamilyMetadataAccumulator.add(processorHandle.operatorMetadata.toMap)
+      operatorProperties.add(Map
+        ("stateVariables" -> processorHandle.stateVariables.toString()))
       setStoreMetrics(store)
       setOperatorMetrics()
       statefulProcessor.close()
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index b58c805af9d60..1140e56160c48 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -56,12 +56,23 @@ trait OperatorStateMetadata {
   def version: Int
 }
 
+object OperatorStateMetadata {
+  def metadataFilePath(stateCheckpointPath: Path): Path =
+    new Path(new Path(stateCheckpointPath, "_metadata"), "metadata")
+}
 case class OperatorStateMetadataV1(
     operatorInfo: OperatorInfoV1,
     stateStoreInfo: Array[StateStoreMetadataV1]) extends OperatorStateMetadata {
   override def version: Int = 1
 }
 
+case class OperatorStateMetadataV2(
+    operatorInfo: OperatorInfoV1,
+    stateStoreInfo: Array[StateStoreMetadataV1],
+    operatorProperties: Map[String, String]) extends OperatorStateMetadata {
+  override def version: Int = 2
+}
+
 object OperatorStateMetadataV1 {
 
   private implicit val formats: Formats = Serialization.formats(NoTypeHints)
@@ -70,9 +81,6 @@ object OperatorStateMetadataV1 {
   private implicit val manifest = Manifest
     .classType[OperatorStateMetadataV1](implicitly[ClassTag[OperatorStateMetadataV1]].runtimeClass)
 
-  def metadataFilePath(stateCheckpointPath: Path): Path =
-    new Path(new Path(stateCheckpointPath, "_metadata"), "metadata")
-
   def deserialize(in: BufferedReader): OperatorStateMetadata = {
     Serialization.read[OperatorStateMetadataV1](in)
   }
@@ -84,13 +92,31 @@ object OperatorStateMetadataV1 {
   }
 }
 
+object OperatorStateMetadataV2 {
+  private implicit val formats: Formats = Serialization.formats(NoTypeHints)
+
+  @scala.annotation.nowarn
+  private implicit val manifest = Manifest
+    .classType[OperatorStateMetadataV1](implicitly[ClassTag[OperatorStateMetadataV2]].runtimeClass)
+
+  def deserialize(in: BufferedReader): OperatorStateMetadata = {
+    Serialization.read[OperatorStateMetadataV2](in)
+  }
+
+  def serialize(
+      out: FSDataOutputStream,
+      operatorStateMetadata: OperatorStateMetadata): Unit = {
+    Serialization.write(operatorStateMetadata.asInstanceOf[OperatorStateMetadataV2], out)
+  }
+}
+
 /**
  * Write OperatorStateMetadata into the state checkpoint directory.
  */
 class OperatorStateMetadataWriter(stateCheckpointPath: Path, hadoopConf: Configuration)
   extends Logging {
 
-  private val metadataFilePath = OperatorStateMetadataV1.metadataFilePath(stateCheckpointPath)
+  private val metadataFilePath = OperatorStateMetadata.metadataFilePath(stateCheckpointPath)
 
   private lazy val fm = CheckpointFileManager.create(stateCheckpointPath, hadoopConf)
 
@@ -101,7 +127,12 @@ class OperatorStateMetadataWriter(stateCheckpointPath: Path, hadoopConf: Configu
     val outputStream = fm.createAtomic(metadataFilePath, overwriteIfPossible = false)
     try {
       outputStream.write(s"v${operatorMetadata.version}\n".getBytes(StandardCharsets.UTF_8))
-      OperatorStateMetadataV1.serialize(outputStream, operatorMetadata)
+      operatorMetadata.version match {
+        case 1 =>
+          OperatorStateMetadataV1.serialize(outputStream, operatorMetadata)
+        case 2 =>
+          OperatorStateMetadataV2.serialize(outputStream, operatorMetadata)
+      }
       outputStream.close()
     } catch {
       case e: Throwable =>
@@ -117,7 +148,7 @@ class OperatorStateMetadataWriter(stateCheckpointPath: Path, hadoopConf: Configu
  */
 class OperatorStateMetadataReader(stateCheckpointPath: Path, hadoopConf: Configuration) {
 
-  private val metadataFilePath = OperatorStateMetadataV1.metadataFilePath(stateCheckpointPath)
+  private val metadataFilePath = OperatorStateMetadata.metadataFilePath(stateCheckpointPath)
 
   private lazy val fm = CheckpointFileManager.create(stateCheckpointPath, hadoopConf)
 

From 696377b72b2fe9f9b41135fce547e721ec2e32fa Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 20 May 2024 14:28:49 -0700
Subject: [PATCH 03/22] writing operator state metadata

---
 .../streaming/IncrementalExecution.scala      | 22 ++++++++++---------
 .../streaming/TransformWithStateExec.scala    | 20 +++++++++++++++++
 .../state/OperatorStateMetadata.scala         |  9 +++++---
 .../streaming/statefulOperators.scala         |  2 ++
 4 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
index f72e2eb407f84..d27b2546a89c3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
+++ b/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,13 @@ 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
+        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 +435,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 {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 6cb03f76cf541..5dddff92b4e3c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.streaming
 import java.util.UUID
 import java.util.concurrent.TimeUnit.NANOSECONDS
 
+import org.apache.hadoop.fs.Path
+
 import org.apache.spark.SparkContext
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
@@ -137,6 +139,17 @@ case class TransformWithStateExec(
       )
     )
 
+  /** 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))
+    OperatorStateMetadataV2(operatorInfo, stateStoreInfo, operatorProperties.value)
+  }
+
+  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.
@@ -145,6 +158,13 @@ case class TransformWithStateExec(
     operatorProperties.value.foreach { case (key, value) =>
       logError(s"### $key: $value")
     }
+    val metadata = operatorStateMetadata()
+    val metadataWriter = new OperatorStateMetadataWriter(
+      new Path(stateInfo.get.checkpointLocation,
+        getStateInfo.operatorId.toString),
+      session.sqlContext.sessionState.newHadoopConf()
+    )
+    metadataWriter.write(metadata)
     children.foreach(_.writeOperatorStateMetadata())
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index 1140e56160c48..003c8985cdebe 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -158,9 +158,12 @@ class OperatorStateMetadataReader(stateCheckpointPath: Path, hadoopConf: Configu
       new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))
     try {
       val versionStr = inputReader.readLine()
-      val version = MetadataVersionUtil.validateVersion(versionStr, 1)
-      assert(version == 1)
-      OperatorStateMetadataV1.deserialize(inputReader)
+      val version = MetadataVersionUtil.validateVersion(versionStr, 2)
+      assert(version == 1 || version == 2)
+      version match {
+        case 1 => OperatorStateMetadataV1.deserialize(inputReader)
+        case 2 => OperatorStateMetadataV2.deserialize(inputReader)
+      }
     } finally {
       inputStream.close()
     }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 9add574e01fc5..d7e8428318b85 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -176,6 +176,8 @@ trait StateStoreWriter extends StatefulOperator with PythonSQLMetrics { self: Sp
     )
   }
 
+  def operatorStateMetadataVersion: Int = 1
+
   /** Records the duration of running `body` for the next query progress update. */
   protected def timeTakenMs(body: => Unit): Long = Utils.timeTakenMs(body)._2
 

From 1ade6bef49372cf12b2a1c6fff0f35c7a1f130f0 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 20 May 2024 14:39:55 -0700
Subject: [PATCH 04/22] moving stuff around

---
 .../StatefulProcessorHandleImpl.scala         | 16 ++++++
 .../streaming/TransformWithStateExec.scala    | 49 -------------------
 .../state/OperatorStateMetadata.scala         | 45 ++++++++++++++++-
 3 files changed, 59 insertions(+), 51 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 9827020a7eac8..44df402db8c26 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -69,6 +69,22 @@ class QueryInfoImpl(
   }
 }
 
+sealed trait StateVariableType
+
+case object ValueState extends StateVariableType
+case object ListState extends StateVariableType
+case object MapState extends StateVariableType
+
+class StateVariableInfo(
+    val stateName: String,
+    val stateType: StateVariableType,
+    val isTtlEnabled: Boolean
+) {
+  override def toString: String = {
+    s"StateVariableInfo(stateName=$stateName, stateType=$stateType, isTtlEnabled=$isTtlEnabled)"
+  }
+}
+
 /**
  * Class that provides a concrete implementation of a StatefulProcessorHandle. Note that we keep
  * track of valid transitions as various functions are invoked to track object lifecycle.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 5dddff92b4e3c..34eb0d5f27588 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -37,55 +37,6 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
 import org.apache.spark.util.{AccumulatorV2, CompletionIterator, SerializableConfiguration, Utils}
 
-// define state variable type enum with types VALUE, LIST, MAP
-sealed trait StateVariableType
-case object ValueState extends StateVariableType
-case object ListState extends StateVariableType
-case object MapState extends StateVariableType
-class StateVariableInfo(
-    val stateName: String,
-    val stateType: StateVariableType,
-    val isTtlEnabled: Boolean
-) {
-    override def toString: String = {
-      s"StateVariableInfo(stateName=$stateName, stateType=$stateType, isTtlEnabled=$isTtlEnabled)"
-    }
-}
-
-class OperatorProperties(initValue: Map[String, String] = Map.empty)
-  extends AccumulatorV2[Map[String, String], Map[String, String]] {
-
-  private var _value: Map[String, String] = initValue
-
-  override def isZero: Boolean = _value.isEmpty
-
-  override def copy(): AccumulatorV2[Map[String, String], Map[String, String]] = {
-    val newAcc = new OperatorProperties
-    newAcc._value = _value
-    newAcc
-  }
-
-  override def reset(): Unit = _value = Map.empty[String, String]
-
-  override def add(v: Map[String, String]): Unit = _value ++= v
-
-  override def merge(other: AccumulatorV2[Map[String, String], Map[String, String]]): Unit = {
-    _value ++= other.value
-  }
-
-  override def value: Map[String, String] = _value
-}
-
-object OperatorProperties {
-  def create(
-      sc: SparkContext,
-      name: String,
-      initValue: Map[String, String] = Map.empty): OperatorProperties = {
-    val acc = new OperatorProperties(initValue)
-    acc.register(sc, name = Some(name))
-    acc
-  }
-}
 /**
  * Physical operator for executing `TransformWithState`
  *
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index 003c8985cdebe..fe806cb15a8ef 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -19,16 +19,16 @@ package org.apache.spark.sql.execution.streaming.state
 
 import java.io.{BufferedReader, InputStreamReader}
 import java.nio.charset.StandardCharsets
-
 import scala.reflect.ClassTag
-
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FSDataOutputStream, Path}
 import org.json4s.{Formats, NoTypeHints}
 import org.json4s.jackson.Serialization
 
+import org.apache.spark.SparkContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, MetadataVersionUtil}
+import org.apache.spark.util.AccumulatorV2
 
 /**
  * Metadata for a state store instance.
@@ -66,6 +66,47 @@ case class OperatorStateMetadataV1(
   override def version: Int = 1
 }
 
+/**
+ * Accumulator to store arbitrary Operator properties.
+ * This accumulator is used to store the properties of an operator that are not
+ * available on the driver at the time of planning, and will only be known from
+ * the executor side.
+ */
+class OperatorProperties(initValue: Map[String, String] = Map.empty)
+  extends AccumulatorV2[Map[String, String], Map[String, String]] {
+
+  private var _value: Map[String, String] = initValue
+
+  override def isZero: Boolean = _value.isEmpty
+
+  override def copy(): AccumulatorV2[Map[String, String], Map[String, String]] = {
+    val newAcc = new OperatorProperties
+    newAcc._value = _value
+    newAcc
+  }
+
+  override def reset(): Unit = _value = Map.empty[String, String]
+
+  override def add(v: Map[String, String]): Unit = _value ++= v
+
+  override def merge(other: AccumulatorV2[Map[String, String], Map[String, String]]): Unit = {
+    _value ++= other.value
+  }
+
+  override def value: Map[String, String] = _value
+}
+
+object OperatorProperties {
+  def create(
+      sc: SparkContext,
+      name: String,
+      initValue: Map[String, String] = Map.empty): OperatorProperties = {
+    val acc = new OperatorProperties(initValue)
+    acc.register(sc, name = Some(name))
+    acc
+  }
+}
+
 case class OperatorStateMetadataV2(
     operatorInfo: OperatorInfoV1,
     stateStoreInfo: Array[StateStoreMetadataV1],

From 8f5c6b42ad9b7394197deac33ec75d74718acbef Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 20 May 2024 15:24:24 -0700
Subject: [PATCH 05/22] comments

---
 .../spark/sql/execution/streaming/IncrementalExecution.scala   | 2 ++
 .../spark/sql/execution/streaming/TransformWithStateExec.scala | 3 ---
 .../sql/execution/streaming/state/OperatorStateMetadata.scala  | 2 ++
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
index d27b2546a89c3..eb04ba0abb715 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
@@ -190,6 +190,8 @@ class IncrementalExecution(
   object WriteStatefulOperatorMetadataRule extends SparkPlanPartialRule {
     override val rule: PartialFunction[SparkPlan, SparkPlan] = {
       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(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 34eb0d5f27588..9ab46d53acc82 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -106,9 +106,6 @@ case class TransformWithStateExec(
    * to write the metadata of the operator to the checkpoint file.
    */
   override def writeOperatorStateMetadata(): Unit = {
-    operatorProperties.value.foreach { case (key, value) =>
-      logError(s"### $key: $value")
-    }
     val metadata = operatorStateMetadata()
     val metadataWriter = new OperatorStateMetadataWriter(
       new Path(stateInfo.get.checkpointLocation,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index fe806cb15a8ef..de637bf8841b3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -19,7 +19,9 @@ package org.apache.spark.sql.execution.streaming.state
 
 import java.io.{BufferedReader, InputStreamReader}
 import java.nio.charset.StandardCharsets
+
 import scala.reflect.ClassTag
+
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FSDataOutputStream, Path}
 import org.json4s.{Formats, NoTypeHints}

From 83ffc451294f318f1b26e3db6f0a3be4ce2587a7 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 20 May 2024 15:39:15 -0700
Subject: [PATCH 06/22] import errors

---
 .../spark/sql/execution/streaming/TransformWithStateExec.scala | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 9ab46d53acc82..50e26712ca738 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -21,7 +21,6 @@ import java.util.concurrent.TimeUnit.NANOSECONDS
 
 import org.apache.hadoop.fs.Path
 
-import org.apache.spark.SparkContext
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
@@ -35,7 +34,7 @@ import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSch
 import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
-import org.apache.spark.util.{AccumulatorV2, CompletionIterator, SerializableConfiguration, Utils}
+import org.apache.spark.util.{CompletionIterator, SerializableConfiguration, Utils}
 
 /**
  * Physical operator for executing `TransformWithState`

From e2fa500fd0aea4935c39cea8b7953f3de95c0853 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 09:03:59 -0700
Subject: [PATCH 07/22] adding test case for state reader integration

---
 .../state/metadata/StateMetadataSource.scala  | 37 ++++++++----
 .../streaming/TransformWithStateSuite.scala   | 58 +++++++++++++++++++
 2 files changed, 83 insertions(+), 12 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index 28de21aaf9389..9696a990a4cc6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/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
@@ -199,17 +199,30 @@ 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,
-          stateStoreMetadata.storeName,
-          stateStoreMetadata.numPartitions,
-          if (batchIds.nonEmpty) batchIds.head else -1,
-          if (batchIds.nonEmpty) batchIds.last else -1,
-          stateStoreMetadata.numColsPrefixKey
-        )
+      require(operatorStateMetadata.version == 1 || operatorStateMetadata.version == 2)
+      operatorStateMetadata match {
+        case v1: OperatorStateMetadataV1 =>
+          v1.stateStoreInfo.map { stateStoreMetadata =>
+            StateMetadataTableEntry(v1.operatorInfo.operatorId,
+              v1.operatorInfo.operatorName,
+              stateStoreMetadata.storeName,
+              stateStoreMetadata.numPartitions,
+              if (batchIds.nonEmpty) batchIds.head else -1,
+              if (batchIds.nonEmpty) batchIds.last else -1,
+              stateStoreMetadata.numColsPrefixKey
+            )
+          }
+        case v2: OperatorStateMetadataV2 =>
+          v2.stateStoreInfo.map { stateStoreMetadata =>
+            StateMetadataTableEntry(v2.operatorInfo.operatorId,
+              v2.operatorInfo.operatorName,
+              stateStoreMetadata.storeName,
+              stateStoreMetadata.numPartitions,
+              if (batchIds.nonEmpty) batchIds.head else -1,
+              if (batchIds.nonEmpty) batchIds.last else -1,
+              stateStoreMetadata.numColsPrefixKey
+            )
+          }
       }
     }
   }.iterator
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 0057af44d3e37..cf857d2df869b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -24,6 +24,7 @@ import org.apache.spark.SparkRuntimeException
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{Dataset, Encoders}
 import org.apache.spark.sql.catalyst.util.stringToFile
+import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions
 import org.apache.spark.sql.execution.streaming._
 import org.apache.spark.sql.execution.streaming.state.{AlsoTestWithChangelogCheckpointingEnabled, RocksDBStateStoreProvider, StatefulProcessorCannotPerformOperationWithInvalidHandleState, StateStoreMultipleColumnFamiliesNotSupportedException}
 import org.apache.spark.sql.functions.timestamp_seconds
@@ -367,6 +368,63 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
+  test("state reader") {
+    withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+      classOf[RocksDBStateStoreProvider].getName) {
+      withTempDir { chkptDir =>
+        val clock = new StreamManualClock
+
+        val inputData = MemoryStream[String]
+        val result = inputData.toDS()
+          .groupByKey(x => x)
+          .transformWithState(new RunningCountStatefulProcessorWithProcTimeTimer(),
+            TimeMode.ProcessingTime(),
+            OutputMode.Update())
+
+        testStream(result, OutputMode.Update())(
+          StartStream(
+            Trigger.ProcessingTime("1 second"),
+            triggerClock = clock,
+            checkpointLocation = chkptDir.getCanonicalPath
+          ),
+          AddData(inputData, "a"),
+          AdvanceManualClock(1 * 1000),
+          CheckNewAnswer(("a", "1")),
+          Execute { q =>
+            assert(q.lastProgress.stateOperators(0).customMetrics.get("numValueStateVars") > 0)
+            assert(q.lastProgress.stateOperators(0).customMetrics.get("numRegisteredTimers") === 1)
+          },
+          AddData(inputData, "b"),
+          AdvanceManualClock(1 * 1000),
+          CheckNewAnswer(("b", "1")),
+
+          AddData(inputData, "b"),
+          AdvanceManualClock(10 * 1000),
+          CheckNewAnswer(("a", "-1"), ("b", "2")),
+
+          AddData(inputData, "b"),
+          AddData(inputData, "c"),
+          AdvanceManualClock(1 * 1000),
+          CheckNewAnswer(("c", "1")), // should remove 'b' as count reaches 3
+
+          AddData(inputData, "d"),
+          AdvanceManualClock(10 * 1000),
+          CheckNewAnswer(("c", "-1"), ("d", "1")),
+          StopStream
+        )
+
+        val stateReadDf = spark.read
+          .format("statestore")
+          .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath)
+          // skip version and operator ID to test out functionalities
+          .load()
+
+//        // print out all the contents of the stateReadDf
+        stateReadDf.show()
+      }
+    }
+  }
+
   test("transformWithState - streaming with rocksdb and processing time timer " +
    "should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->

From d672d36000509a9628c38c682624a9abeccf7621 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 09:04:18 -0700
Subject: [PATCH 08/22] comment spacing

---
 .../apache/spark/sql/streaming/TransformWithStateSuite.scala    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index cf857d2df869b..06d6aaf8641a0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -419,7 +419,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           // skip version and operator ID to test out functionalities
           .load()
 
-//        // print out all the contents of the stateReadDf
+        // print out all the contents of the stateReadDf
         stateReadDf.show()
       }
     }

From 7f3a9a8910f5e6f4f41aabab772fa27eb5ec20b5 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 10:54:00 -0700
Subject: [PATCH 09/22] added field

---
 .../state/metadata/StateMetadataSource.scala  | 20 +++++++++++++++----
 .../streaming/TransformWithStateSuite.scala   |  2 +-
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index 9696a990a4cc6..d341539fe4a7f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -26,6 +26,7 @@ import org.apache.hadoop.fs.{Path, PathFilter}
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, MapData}
 import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
 import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder}
@@ -34,7 +35,7 @@ import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.PA
 import org.apache.spark.sql.execution.streaming.CheckpointFileManager
 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.types.{DataType, IntegerType, LongType, MapType, StringType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.SerializableConfiguration
@@ -46,8 +47,15 @@ case class StateMetadataTableEntry(
     numPartitions: Int,
     minBatchId: Long,
     maxBatchId: Long,
-    numColsPrefixKey: Int) {
+    numColsPrefixKey: Int,
+    operatorProperties: Map[String, String] = Map.empty[String, String]) {
   def toRow(): InternalRow = {
+    // create MapData from Map
+    val mapData: MapData = ArrayBasedMapData.apply(
+      operatorProperties.keys.map(UTF8String.fromString).toArray,
+      operatorProperties.values.map(UTF8String.fromString).toArray
+    )
+
     new GenericInternalRow(
       Array[Any](operatorId,
         UTF8String.fromString(operatorName),
@@ -55,7 +63,8 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
-        numColsPrefixKey))
+        numColsPrefixKey,
+        mapData))
   }
 }
 
@@ -68,6 +77,8 @@ object StateMetadataTableEntry {
       .add("numPartitions", IntegerType)
       .add("minBatchId", LongType)
       .add("maxBatchId", LongType)
+      .add("numColsPrefixKey", IntegerType)
+      .add("operatorProperties", MapType(StringType, StringType, valueContainsNull = false))
   }
 }
 
@@ -220,7 +231,8 @@ class StateMetadataPartitionReader(
               stateStoreMetadata.numPartitions,
               if (batchIds.nonEmpty) batchIds.head else -1,
               if (batchIds.nonEmpty) batchIds.last else -1,
-              stateStoreMetadata.numColsPrefixKey
+              stateStoreMetadata.numColsPrefixKey,
+              v2.operatorProperties
             )
           }
       }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 06d6aaf8641a0..001f408d3f293 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -414,7 +414,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
         )
 
         val stateReadDf = spark.read
-          .format("statestore")
+          .format("state-metadata")
           .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath)
           // skip version and operator ID to test out functionalities
           .load()

From ec346188f167db08431f5092ce3a31dc43eba15c Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 11:57:13 -0700
Subject: [PATCH 10/22] moving numColsPrefixKey back

---
 .../v2/state/metadata/StateMetadataSource.scala    | 14 ++++++--------
 .../streaming/TransformWithStateExec.scala         |  2 +-
 .../sql/streaming/TransformWithStateSuite.scala    |  7 +++----
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index d341539fe4a7f..3ea790e30d05b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -47,14 +47,13 @@ case class StateMetadataTableEntry(
     numPartitions: Int,
     minBatchId: Long,
     maxBatchId: Long,
-    numColsPrefixKey: Int,
-    operatorProperties: Map[String, String] = Map.empty[String, String]) {
+    operatorProperties: Map[String, String],
+    numColsPrefixKey: Int) {
   def toRow(): InternalRow = {
     // create MapData from Map
     val mapData: MapData = ArrayBasedMapData.apply(
       operatorProperties.keys.map(UTF8String.fromString).toArray,
-      operatorProperties.values.map(UTF8String.fromString).toArray
-    )
+      operatorProperties.values.map(UTF8String.fromString).toArray)
 
     new GenericInternalRow(
       Array[Any](operatorId,
@@ -63,7 +62,6 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
-        numColsPrefixKey,
         mapData))
   }
 }
@@ -77,7 +75,6 @@ object StateMetadataTableEntry {
       .add("numPartitions", IntegerType)
       .add("minBatchId", LongType)
       .add("maxBatchId", LongType)
-      .add("numColsPrefixKey", IntegerType)
       .add("operatorProperties", MapType(StringType, StringType, valueContainsNull = false))
   }
 }
@@ -220,6 +217,7 @@ class StateMetadataPartitionReader(
               stateStoreMetadata.numPartitions,
               if (batchIds.nonEmpty) batchIds.head else -1,
               if (batchIds.nonEmpty) batchIds.last else -1,
+              Map.empty[String, String],
               stateStoreMetadata.numColsPrefixKey
             )
           }
@@ -231,8 +229,8 @@ class StateMetadataPartitionReader(
               stateStoreMetadata.numPartitions,
               if (batchIds.nonEmpty) batchIds.head else -1,
               if (batchIds.nonEmpty) batchIds.last else -1,
-              stateStoreMetadata.numColsPrefixKey,
-              v2.operatorProperties
+              v2.operatorProperties,
+              stateStoreMetadata.numColsPrefixKey
             )
           }
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 50e26712ca738..8fea0aef57a25 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -112,7 +112,7 @@ case class TransformWithStateExec(
       session.sqlContext.sessionState.newHadoopConf()
     )
     metadataWriter.write(metadata)
-    children.foreach(_.writeOperatorStateMetadata())
+    super.writeOperatorStateMetadata()
   }
 
   override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 001f408d3f293..0309e114ecd75 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -368,7 +368,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("state reader") {
+  test("verify that operatorProperties contain all stateVariables") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       withTempDir { chkptDir =>
@@ -413,14 +413,13 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           StopStream
         )
 
-        val stateReadDf = spark.read
+        val df = spark.read
           .format("state-metadata")
           .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath)
           // skip version and operator ID to test out functionalities
           .load()
 
-        // print out all the contents of the stateReadDf
-        stateReadDf.show()
+        df.show()
       }
     }
   }

From 6b05daf1ba93317bf4cf55bbca73eb116f1ed09c Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 19:09:07 -0700
Subject: [PATCH 11/22] test changes, need to find a way to check the state
 variables list

---
 .../state/metadata/StateMetadataSource.scala  | 15 ++++------
 .../StatefulProcessorHandleImpl.scala         | 10 +++++++
 .../streaming/TransformWithStateExec.scala    | 29 ++++++++++++++-----
 .../state/OperatorStateMetadata.scala         | 21 +++++++-------
 .../streaming/TransformWithStateSuite.scala   | 19 +++++++++++-
 5 files changed, 66 insertions(+), 28 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index 3ea790e30d05b..3a3d685a6287f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -26,7 +26,6 @@ import org.apache.hadoop.fs.{Path, PathFilter}
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
-import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, MapData}
 import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
 import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder}
@@ -35,7 +34,7 @@ import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.PA
 import org.apache.spark.sql.execution.streaming.CheckpointFileManager
 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, MapType, StringType, StructType}
+import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StringType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.SerializableConfiguration
@@ -47,13 +46,9 @@ case class StateMetadataTableEntry(
     numPartitions: Int,
     minBatchId: Long,
     maxBatchId: Long,
-    operatorProperties: Map[String, String],
+    operatorProperties: String,
     numColsPrefixKey: Int) {
   def toRow(): InternalRow = {
-    // create MapData from Map
-    val mapData: MapData = ArrayBasedMapData.apply(
-      operatorProperties.keys.map(UTF8String.fromString).toArray,
-      operatorProperties.values.map(UTF8String.fromString).toArray)
 
     new GenericInternalRow(
       Array[Any](operatorId,
@@ -62,7 +57,7 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
-        mapData))
+        UTF8String.fromString(operatorProperties)))
   }
 }
 
@@ -75,7 +70,7 @@ object StateMetadataTableEntry {
       .add("numPartitions", IntegerType)
       .add("minBatchId", LongType)
       .add("maxBatchId", LongType)
-      .add("operatorProperties", MapType(StringType, StringType, valueContainsNull = false))
+      .add("operatorProperties", StringType)
   }
 }
 
@@ -217,7 +212,7 @@ class StateMetadataPartitionReader(
               stateStoreMetadata.numPartitions,
               if (batchIds.nonEmpty) batchIds.head else -1,
               if (batchIds.nonEmpty) batchIds.last else -1,
-              Map.empty[String, String],
+              "",
               stateStoreMetadata.numColsPrefixKey
             )
           }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 44df402db8c26..54913b6af0912 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -21,6 +21,9 @@ import java.util.UUID
 
 import scala.collection.mutable
 
+import org.json4s.JsonAST.{JBool, JString, JValue}
+import org.json4s.JsonDSL._
+
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
@@ -83,6 +86,13 @@ class StateVariableInfo(
   override def toString: String = {
     s"StateVariableInfo(stateName=$stateName, stateType=$stateType, isTtlEnabled=$isTtlEnabled)"
   }
+
+  def jsonValue: JValue = {
+    ("stateName" -> JString(stateName)) ~
+      ("stateType" -> JString(stateType.toString)) ~
+      ("isTtlEnabled" -> JBool(isTtlEnabled))
+
+  }
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 8fea0aef57a25..444c67ccf4cd7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -20,6 +20,11 @@ import java.util.UUID
 import java.util.concurrent.TimeUnit.NANOSECONDS
 
 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
@@ -82,11 +87,7 @@ case class TransformWithStateExec(
   val operatorProperties: OperatorProperties =
     OperatorProperties.create(
       sparkContext,
-      "colFamilyMetadata",
-      Map(
-        "timeMode" -> timeMode.toString,
-        "outputMode" -> outputMode.toString
-      )
+      "colFamilyMetadata"
     )
 
   /** Metadata of this stateful operator and its states stores. */
@@ -95,7 +96,13 @@ case class TransformWithStateExec(
     val operatorInfo = OperatorInfoV1(info.operatorId, shortName)
     val stateStoreInfo =
       Array(StateStoreMetadataV1(StateStoreId.DEFAULT_STORE_NAME, 0, info.numPartitions))
-    OperatorStateMetadataV2(operatorInfo, stateStoreInfo, operatorProperties.value)
+
+    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
@@ -345,7 +352,7 @@ case class TransformWithStateExec(
         }
       }
       operatorProperties.add(Map
-        ("stateVariables" -> processorHandle.stateVariables.toString()))
+        ("stateVariables" -> JArray(processorHandle.stateVariables.map(_.jsonValue).toList)))
       setStoreMetrics(store)
       setOperatorMetrics()
       statefulProcessor.close()
@@ -601,6 +608,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
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index de637bf8841b3..8d92c04d4cceb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -25,6 +25,7 @@ import scala.reflect.ClassTag
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FSDataOutputStream, Path}
 import org.json4s.{Formats, NoTypeHints}
+import org.json4s.JsonAST.JValue
 import org.json4s.jackson.Serialization
 
 import org.apache.spark.SparkContext
@@ -74,35 +75,35 @@ case class OperatorStateMetadataV1(
  * available on the driver at the time of planning, and will only be known from
  * the executor side.
  */
-class OperatorProperties(initValue: Map[String, String] = Map.empty)
-  extends AccumulatorV2[Map[String, String], Map[String, String]] {
+class OperatorProperties(initValue: Map[String, JValue] = Map.empty)
+  extends AccumulatorV2[Map[String, JValue], Map[String, JValue]] {
 
-  private var _value: Map[String, String] = initValue
+  private var _value: Map[String, JValue] = initValue
 
   override def isZero: Boolean = _value.isEmpty
 
-  override def copy(): AccumulatorV2[Map[String, String], Map[String, String]] = {
+  override def copy(): AccumulatorV2[Map[String, JValue], Map[String, JValue]] = {
     val newAcc = new OperatorProperties
     newAcc._value = _value
     newAcc
   }
 
-  override def reset(): Unit = _value = Map.empty[String, String]
+  override def reset(): Unit = _value = Map.empty[String, JValue]
 
-  override def add(v: Map[String, String]): Unit = _value ++= v
+  override def add(v: Map[String, JValue]): Unit = _value ++= v
 
-  override def merge(other: AccumulatorV2[Map[String, String], Map[String, String]]): Unit = {
+  override def merge(other: AccumulatorV2[Map[String, JValue], Map[String, JValue]]): Unit = {
     _value ++= other.value
   }
 
-  override def value: Map[String, String] = _value
+  override def value: Map[String, JValue] = _value
 }
 
 object OperatorProperties {
   def create(
       sc: SparkContext,
       name: String,
-      initValue: Map[String, String] = Map.empty): OperatorProperties = {
+      initValue: Map[String, JValue] = Map.empty): OperatorProperties = {
     val acc = new OperatorProperties(initValue)
     acc.register(sc, name = Some(name))
     acc
@@ -112,7 +113,7 @@ object OperatorProperties {
 case class OperatorStateMetadataV2(
     operatorInfo: OperatorInfoV1,
     stateStoreInfo: Array[StateStoreMetadataV1],
-    operatorProperties: Map[String, String]) extends OperatorStateMetadata {
+    operatorProperties: String) extends OperatorStateMetadata {
   override def version: Int = 2
 }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 0309e114ecd75..bcc6865211e63 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -416,10 +416,27 @@ class TransformWithStateSuite extends StateStoreMetricsTest
         val df = spark.read
           .format("state-metadata")
           .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath)
-          // skip version and operator ID to test out functionalities
           .load()
 
         df.show()
+
+        val propsString = df.select("operatorProperties").
+          collect().head.getString(0)
+
+        val map = TransformWithStateExec.
+          deserializeOperatorProperties(propsString)
+        assert(map("timeMode") === "ProcessingTime")
+        assert(map("outputMode") === "Update")
+
+        // TODO: Find a way to verify the stateVariables, need to deserialize from JArray
+//       // get "stateVariables" as list and cast this list as a collection of StateVariableInfo
+//        val stateVariables = map("stateVariables").asInstanceOf[JArray].arr
+//        val stateVariableInfos = stateVariables.map { sv =>
+//
+//        }
+//
+//        // check that all stateVariables are present in operatorProperties
+//        assert(stateVariableInfos.size === 1)
       }
     }
   }

From 169403db33e0fc8ead1046dd03358b2bb3ebdb1f Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 19:13:21 -0700
Subject: [PATCH 12/22] adding numColsPrefixKey

---
 .../datasources/v2/state/metadata/StateMetadataSource.scala   | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index 3a3d685a6287f..2ca13581ba7d5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -57,7 +57,9 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
-        UTF8String.fromString(operatorProperties)))
+        UTF8String.fromString(operatorProperties),
+        numColsPrefixKey
+      ))
   }
 }
 

From 4286502b9233060d7a2b2034f194eb6f15faec10 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 21 May 2024 22:42:47 -0700
Subject: [PATCH 13/22] maybe the dumbest serialization methods known to man

---
 .../StatefulProcessorHandleImpl.scala         | 26 ++++++++++++++++++-
 .../streaming/TransformWithStateSuite.scala   | 13 +++-------
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 54913b6af0912..06e07e12cf64d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -21,7 +21,8 @@ import java.util.UUID
 
 import scala.collection.mutable
 
-import org.json4s.JsonAST.{JBool, JString, JValue}
+import org.json4s.{DefaultFormats, Formats, JValue}
+import org.json4s.JsonAST.{JBool, JString}
 import org.json4s.JsonDSL._
 
 import org.apache.spark.TaskContext
@@ -78,6 +79,14 @@ case object ValueState extends StateVariableType
 case object ListState extends StateVariableType
 case object MapState extends StateVariableType
 
+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")
+  }
+}
 class StateVariableInfo(
     val stateName: String,
     val stateType: StateVariableType,
@@ -95,6 +104,21 @@ class StateVariableInfo(
   }
 }
 
+object StateVariableInfo {
+  implicit val formats: Formats = DefaultFormats
+  def fromJson(json: Any): List[StateVariableInfo] = {
+    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]
+      )
+    }
+  }
+}
+
 /**
  * Class that provides a concrete implementation of a StatefulProcessorHandle. Note that we keep
  * track of valid transitions as various functions are invoked to track object lifecycle.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index bcc6865211e63..14508bfd221e1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -428,15 +428,10 @@ class TransformWithStateSuite extends StateStoreMetricsTest
         assert(map("timeMode") === "ProcessingTime")
         assert(map("outputMode") === "Update")
 
-        // TODO: Find a way to verify the stateVariables, need to deserialize from JArray
-//       // get "stateVariables" as list and cast this list as a collection of StateVariableInfo
-//        val stateVariables = map("stateVariables").asInstanceOf[JArray].arr
-//        val stateVariableInfos = stateVariables.map { sv =>
-//
-//        }
-//
-//        // check that all stateVariables are present in operatorProperties
-//        assert(stateVariableInfos.size === 1)
+        val stateVariableInfos = StateVariableInfo.fromJson(
+          map("stateVariables"))
+        assert(stateVariableInfos.size === 1)
+        assert(stateVariableInfos.exists(_.stateName === "countState"))
       }
     }
   }

From 2c1aa946bea812a900940797c6c4b694f2640a06 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Wed, 22 May 2024 09:17:46 -0700
Subject: [PATCH 14/22] tests work

---
 .../spark/sql/streaming/TransformWithStateSuite.scala      | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 14508bfd221e1..862b02af354e6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -418,8 +418,6 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath)
           .load()
 
-        df.show()
-
         val propsString = df.select("operatorProperties").
           collect().head.getString(0)
 
@@ -431,7 +429,10 @@ class TransformWithStateSuite extends StateStoreMetricsTest
         val stateVariableInfos = StateVariableInfo.fromJson(
           map("stateVariables"))
         assert(stateVariableInfos.size === 1)
-        assert(stateVariableInfos.exists(_.stateName === "countState"))
+        val stateVariableInfo = stateVariableInfos.head
+        assert(stateVariableInfo.stateName === "countState")
+        assert(stateVariableInfo.isTtlEnabled === false)
+        assert(stateVariableInfo.stateType === ValueState)
       }
     }
   }

From 94b7a67b7bf014ba2c3aa33496e645ea597ebe6f Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Wed, 22 May 2024 09:26:25 -0700
Subject: [PATCH 15/22] comment

---
 .../sql/execution/streaming/state/OperatorStateMetadata.scala | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index 8d92c04d4cceb..d5d7a5f55d6f2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -110,6 +110,8 @@ object OperatorProperties {
   }
 }
 
+// operatorProperties is an arbitrary JSON formatted string that contains
+// any properties that we would want to store for a particular operator.
 case class OperatorStateMetadataV2(
     operatorInfo: OperatorInfoV1,
     stateStoreInfo: Array[StateStoreMetadataV1],
@@ -141,7 +143,7 @@ object OperatorStateMetadataV2 {
 
   @scala.annotation.nowarn
   private implicit val manifest = Manifest
-    .classType[OperatorStateMetadataV1](implicitly[ClassTag[OperatorStateMetadataV2]].runtimeClass)
+    .classType[OperatorStateMetadataV2](implicitly[ClassTag[OperatorStateMetadataV2]].runtimeClass)
 
   def deserialize(in: BufferedReader): OperatorStateMetadata = {
     Serialization.read[OperatorStateMetadataV2](in)

From 76aa1d0fae6d6fef9a20a12b2e2299fbd6e66d1b Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Wed, 22 May 2024 09:34:23 -0700
Subject: [PATCH 16/22] creating new file

---
 .../streaming/StateVariableInfo.scala         | 71 +++++++++++++++++++
 .../StatefulProcessorHandleImpl.scala         | 50 -------------
 2 files changed, 71 insertions(+), 50 deletions(-)
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateVariableInfo.scala

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateVariableInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateVariableInfo.scala
new file mode 100644
index 0000000000000..e381984b4672d
--- /dev/null
+++ b/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]
+      )
+    }
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 06e07e12cf64d..9827020a7eac8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -21,10 +21,6 @@ import java.util.UUID
 
 import scala.collection.mutable
 
-import org.json4s.{DefaultFormats, Formats, JValue}
-import org.json4s.JsonAST.{JBool, JString}
-import org.json4s.JsonDSL._
-
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
@@ -73,52 +69,6 @@ class QueryInfoImpl(
   }
 }
 
-sealed trait StateVariableType
-
-case object ValueState extends StateVariableType
-case object ListState extends StateVariableType
-case object MapState extends StateVariableType
-
-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")
-  }
-}
-class StateVariableInfo(
-    val stateName: String,
-    val stateType: StateVariableType,
-    val isTtlEnabled: Boolean
-) {
-  override def toString: String = {
-    s"StateVariableInfo(stateName=$stateName, stateType=$stateType, isTtlEnabled=$isTtlEnabled)"
-  }
-
-  def jsonValue: JValue = {
-    ("stateName" -> JString(stateName)) ~
-      ("stateType" -> JString(stateType.toString)) ~
-      ("isTtlEnabled" -> JBool(isTtlEnabled))
-
-  }
-}
-
-object StateVariableInfo {
-  implicit val formats: Formats = DefaultFormats
-  def fromJson(json: Any): List[StateVariableInfo] = {
-    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]
-      )
-    }
-  }
-}
-
 /**
  * Class that provides a concrete implementation of a StatefulProcessorHandle. Note that we keep
  * track of valid transitions as various functions are invoked to track object lifecycle.

From 005dbeab4a8ed07c91674d1920b426710a79fac9 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Wed, 22 May 2024 09:38:05 -0700
Subject: [PATCH 17/22] renaming variable

---
 .../v2/state/metadata/StateMetadataSource.scala            | 7 +++----
 .../execution/streaming/state/OperatorStateMetadata.scala  | 2 +-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index 2ca13581ba7d5..9eedda0ecac36 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -46,10 +46,9 @@ case class StateMetadataTableEntry(
     numPartitions: Int,
     minBatchId: Long,
     maxBatchId: Long,
-    operatorProperties: String,
+    operatorPropertiesJson: String,
     numColsPrefixKey: Int) {
   def toRow(): InternalRow = {
-
     new GenericInternalRow(
       Array[Any](operatorId,
         UTF8String.fromString(operatorName),
@@ -57,7 +56,7 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
-        UTF8String.fromString(operatorProperties),
+        UTF8String.fromString(operatorPropertiesJson),
         numColsPrefixKey
       ))
   }
@@ -226,7 +225,7 @@ class StateMetadataPartitionReader(
               stateStoreMetadata.numPartitions,
               if (batchIds.nonEmpty) batchIds.head else -1,
               if (batchIds.nonEmpty) batchIds.last else -1,
-              v2.operatorProperties,
+              v2.operatorPropertiesJson,
               stateStoreMetadata.numColsPrefixKey
             )
           }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index d5d7a5f55d6f2..beb2f46d9dd51 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -115,7 +115,7 @@ object OperatorProperties {
 case class OperatorStateMetadataV2(
     operatorInfo: OperatorInfoV1,
     stateStoreInfo: Array[StateStoreMetadataV1],
-    operatorProperties: String) extends OperatorStateMetadata {
+    operatorPropertiesJson: String) extends OperatorStateMetadata {
   override def version: Int = 2
 }
 

From 109191a100a66e57728d126d776dde01906abb12 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 3 Jun 2024 09:41:47 -0700
Subject: [PATCH 18/22] using collection converter

---
 .../streaming/StatefulProcessorHandleImpl.scala  | 16 ++++++++--------
 .../streaming/TransformWithStateExec.scala       |  5 ++++-
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 9827020a7eac8..d60a830fe4825 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -96,8 +96,8 @@ class StatefulProcessorHandleImpl(
    */
   private[sql] val ttlStates: util.List[TTLState] = new util.ArrayList[TTLState]()
 
-  var stateVariables: mutable.ListBuffer[StateVariableInfo] =
-    mutable.ListBuffer.empty[StateVariableInfo]
+  private[sql] val stateVariables: util.List[StateVariableInfo] =
+    new util.ArrayList[StateVariableInfo]()
 
   private val BATCH_QUERY_ID = "00000000-0000-0000-0000-000000000000"
 
@@ -134,7 +134,7 @@ class StatefulProcessorHandleImpl(
       stateName: String,
       valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
-    stateVariables.append(new StateVariableInfo(stateName, ValueState, false))
+    stateVariables.add(new StateVariableInfo(stateName, ValueState, false))
     incrementMetric("numValueStateVars")
     val resultState = new ValueStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -145,7 +145,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[T],
       ttlConfig: TTLConfig): ValueState[T] = {
     verifyStateVarOperations("get_value_state")
-    stateVariables.append(new StateVariableInfo(stateName, ValueState, true))
+    stateVariables.add(new StateVariableInfo(stateName, ValueState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -247,7 +247,7 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state")
-    stateVariables.append(new StateVariableInfo(stateName, ListState, false))
+    stateVariables.add(new StateVariableInfo(stateName, ListState, false))
     incrementMetric("numListStateVars")
     val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder)
     resultState
@@ -274,7 +274,7 @@ class StatefulProcessorHandleImpl(
       ttlConfig: TTLConfig): ListState[T] = {
 
     verifyStateVarOperations("get_list_state")
-    stateVariables.append(new StateVariableInfo(stateName, ListState, true))
+    stateVariables.add(new StateVariableInfo(stateName, ListState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
@@ -291,7 +291,7 @@ class StatefulProcessorHandleImpl(
       userKeyEnc: Encoder[K],
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
-    stateVariables.append(new StateVariableInfo(stateName, MapState, false))
+    stateVariables.add(new StateVariableInfo(stateName, MapState, false))
     incrementMetric("numMapStateVars")
     val resultState = new MapStateImpl[K, V](store, stateName, keyEncoder, userKeyEnc, valEncoder)
     resultState
@@ -303,7 +303,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[V],
       ttlConfig: TTLConfig): MapState[K, V] = {
     verifyStateVarOperations("get_map_state")
-    stateVariables.append(new StateVariableInfo(stateName, MapState, true))
+    stateVariables.add(new StateVariableInfo(stateName, MapState, true))
     validateTTLConfig(ttlConfig, stateName)
 
     assert(batchTimestampMs.isDefined)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 444c67ccf4cd7..6e6f56e3169c8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -19,6 +19,8 @@ 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
@@ -352,7 +354,8 @@ case class TransformWithStateExec(
         }
       }
       operatorProperties.add(Map
-        ("stateVariables" -> JArray(processorHandle.stateVariables.map(_.jsonValue).toList)))
+        ("stateVariables" -> JArray(processorHandle.stateVariables.
+          asScala.map(_.jsonValue).toList)))
       setStoreMetrics(store)
       setOperatorMetrics()
       statefulProcessor.close()

From 2e89fb42688e90ea34f3c56560fa99a83f3e4e88 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 3 Jun 2024 10:28:31 -0700
Subject: [PATCH 19/22] checking previous execution runId

---
 .../streaming/MicroBatchExecution.scala          |  7 ++++++-
 .../execution/streaming/ProgressReporter.scala   | 16 +++++++++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala
index ef49de5e0857d..6d763390f2a3d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala
+++ b/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,
+          Some(execCtx.previousContext)
+        )
       } else {
         execCtx.finishNoExecutionTrigger(execCtx.batchId)
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
index 748f54af89371..37a1629704ebb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
@@ -264,13 +264,18 @@ abstract class ProgressContext(
       hasNewData: Boolean,
       sourceToNumInputRowsMap: Map[SparkDataStream, Long],
       lastExecution: IncrementalExecution,
-      lastEpochId: Long): Unit = {
+      lastEpochId: Long,
+      previousExecution: Option[StreamExecutionContext]): Unit = {
     assert(
       currentTriggerStartOffsets != null && currentTriggerEndOffsets != null &&
         currentTriggerLatestOffsets != null
     )
-    if (lastExecution.isFirstBatch) {
-      lastExecution.executedPlan.writeOperatorStateMetadata()
+
+    previousExecution match {
+      case Some(prevContext) => if (!prevContext.runId.equals(lastExecution.runId)) {
+        lastExecution.executedPlan.writeOperatorStateMetadata()
+      }
+      case None => lastExecution.executedPlan.writeOperatorStateMetadata()
     }
 
     currentTriggerEndTimestamp = triggerClock.getTimeMillis()
@@ -383,10 +388,11 @@ abstract class ProgressContext(
   def finishTrigger(
       hasNewData: Boolean,
       lastExecution: IncrementalExecution,
-      lastEpoch: Long): Unit = {
+      lastEpoch: Long,
+      previousContext: Option[StreamExecutionContext] = 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(

From 54e63ff2f33e484af6a9e4e7eee881f2fd790826 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 3 Jun 2024 11:00:33 -0700
Subject: [PATCH 20/22] removing unused import

---
 .../sql/execution/streaming/StatefulProcessorHandleImpl.scala   | 2 --
 1 file changed, 2 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index d60a830fe4825..f8948b07457e5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -19,8 +19,6 @@ package org.apache.spark.sql.execution.streaming
 import java.util
 import java.util.UUID
 
-import scala.collection.mutable
-
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder

From fec850dbc19f0c5daacb94b90a84618c6f920732 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 3 Jun 2024 11:04:16 -0700
Subject: [PATCH 21/22] no compilation errors

---
 .../spark/sql/execution/streaming/MicroBatchExecution.scala   | 2 +-
 .../spark/sql/execution/streaming/ProgressReporter.scala      | 4 ++--
 .../sql/execution/streaming/StreamExecutionContext.scala      | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala
index 6d763390f2a3d..a7ea02d957830 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala
@@ -390,7 +390,7 @@ class MicroBatchExecution(
           currentBatchHasNewData,
           execCtx.executionPlan,
           execCtx.batchId,
-          Some(execCtx.previousContext)
+          execCtx.previousContext
         )
       } else {
         execCtx.finishNoExecutionTrigger(execCtx.batchId)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
index 37a1629704ebb..301154c515901 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ProgressReporter.scala
@@ -265,7 +265,7 @@ abstract class ProgressContext(
       sourceToNumInputRowsMap: Map[SparkDataStream, Long],
       lastExecution: IncrementalExecution,
       lastEpochId: Long,
-      previousExecution: Option[StreamExecutionContext]): Unit = {
+      previousExecution: Option[MicroBatchExecutionContext]): Unit = {
     assert(
       currentTriggerStartOffsets != null && currentTriggerEndOffsets != null &&
         currentTriggerLatestOffsets != null
@@ -389,7 +389,7 @@ abstract class ProgressContext(
       hasNewData: Boolean,
       lastExecution: IncrementalExecution,
       lastEpoch: Long,
-      previousContext: Option[StreamExecutionContext] = None): Unit = {
+      previousContext: Option[MicroBatchExecutionContext] = None): Unit = {
     val map: Map[SparkDataStream, Long] =
       if (hasNewData) extractSourceToNumInputRows(lastExecution) else Map.empty
     finishTrigger(hasNewData, map, lastExecution, lastEpoch, previousContext)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecutionContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecutionContext.scala
index c5e14df3e20e1..6ca1945c7599e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecutionContext.scala
+++ b/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],

From 0ad6080a6e74f8ac83bc450cc8c149b8fe8899b0 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 3 Jun 2024 11:56:49 -0700
Subject: [PATCH 22/22] refactor statemetadatasource method

---
 .../state/metadata/StateMetadataSource.scala  | 39 +++++++------------
 .../state/OperatorStateMetadata.scala         |  4 ++
 2 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index 9eedda0ecac36..fe1c6f248cb97 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -204,31 +204,20 @@ class StateMetadataPartitionReader(
   private[state] lazy val stateMetadata: Iterator[StateMetadataTableEntry] = {
     allOperatorStateMetadata.flatMap { operatorStateMetadata =>
       require(operatorStateMetadata.version == 1 || operatorStateMetadata.version == 2)
-      operatorStateMetadata match {
-        case v1: OperatorStateMetadataV1 =>
-          v1.stateStoreInfo.map { stateStoreMetadata =>
-            StateMetadataTableEntry(v1.operatorInfo.operatorId,
-              v1.operatorInfo.operatorName,
-              stateStoreMetadata.storeName,
-              stateStoreMetadata.numPartitions,
-              if (batchIds.nonEmpty) batchIds.head else -1,
-              if (batchIds.nonEmpty) batchIds.last else -1,
-              "",
-              stateStoreMetadata.numColsPrefixKey
-            )
-          }
-        case v2: OperatorStateMetadataV2 =>
-          v2.stateStoreInfo.map { stateStoreMetadata =>
-            StateMetadataTableEntry(v2.operatorInfo.operatorId,
-              v2.operatorInfo.operatorName,
-              stateStoreMetadata.storeName,
-              stateStoreMetadata.numPartitions,
-              if (batchIds.nonEmpty) batchIds.head else -1,
-              if (batchIds.nonEmpty) batchIds.last else -1,
-              v2.operatorPropertiesJson,
-              stateStoreMetadata.numColsPrefixKey
-            )
-          }
+      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
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
index beb2f46d9dd51..e210e669b6c10 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OperatorStateMetadata.scala
@@ -57,6 +57,10 @@ case class OperatorInfoV1(operatorId: Long, operatorName: String) extends Operat
 
 trait OperatorStateMetadata {
   def version: Int
+
+  def operatorInfo: OperatorInfo
+
+  def stateStoreInfo: Array[StateStoreMetadataV1]
 }
 
 object OperatorStateMetadata {