creating operatorstatemetadata log

Author: ericm-db

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

@@ -31,8 +31,8 @@ import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder}
 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.{CheckpointFileManager, OperatorStateMetadataLog}
+import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadata, OperatorStateMetadataReader, OperatorStateMetadataV1, OperatorStateMetadataV2}
 import org.apache.spark.sql.sources.DataSourceRegister
 import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StringType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -46,6 +46,7 @@ case class StateMetadataTableEntry(
     numPartitions: Int,
     minBatchId: Long,
     maxBatchId: Long,
+    operatorPropertiesJson: String,
     numColsPrefixKey: Int) {
   def toRow(): InternalRow = {
     new GenericInternalRow(
@@ -55,6 +56,7 @@ case class StateMetadataTableEntry(
         numPartitions,
         minBatchId,
         maxBatchId,
+        UTF8String.fromString(operatorPropertiesJson),
         numColsPrefixKey))
   }
 }
@@ -68,6 +70,7 @@ object StateMetadataTableEntry {
       .add("numPartitions", IntegerType)
       .add("minBatchId", LongType)
       .add("maxBatchId", LongType)
+      .add("operatorProperties", StringType)
   }
 }
 
@@ -192,22 +195,35 @@ class StateMetadataPartitionReader(
     val stateDir = new Path(checkpointLocation, "state")
     val opIds = fileManager
       .list(stateDir, pathNameCanBeParsedAsLongFilter).map(f => pathToLong(f.getPath)).sorted
-    opIds.map { opId =>
-      new OperatorStateMetadataReader(new Path(stateDir, opId.toString), hadoopConf).read()
+    opIds.flatMap { opId =>
+      val operatorIdPath = new Path(stateDir, opId.toString)
+      // check all OperatorStateMetadataV2
+      val operatorStateMetadataV2Path = OperatorStateMetadataV2.metadataFilePath(operatorIdPath)
+      if (fileManager.exists(operatorStateMetadataV2Path)) {
+        val operatorStateMetadataLog = new OperatorStateMetadataLog(
+          hadoopConf, operatorStateMetadataV2Path.toString)
+        operatorStateMetadataLog.listBatchesOnDisk.flatMap(operatorStateMetadataLog.get)
+      } else {
+        Array(new OperatorStateMetadataReader(operatorIdPath, hadoopConf).read())
+      }
     }
   }
 
   private[state] lazy val stateMetadata: Iterator[StateMetadataTableEntry] = {
     allOperatorStateMetadata.flatMap { operatorStateMetadata =>
-      require(operatorStateMetadata.version == 1)
-      val operatorStateMetadataV1 = operatorStateMetadata.asInstanceOf[OperatorStateMetadataV1]
-      operatorStateMetadataV1.stateStoreInfo.map { stateStoreMetadata =>
-        StateMetadataTableEntry(operatorStateMetadataV1.operatorInfo.operatorId,
-          operatorStateMetadataV1.operatorInfo.operatorName,
+      require(operatorStateMetadata.version == 1 || operatorStateMetadata.version == 2)
+      val operatorProperties = operatorStateMetadata match {
+        case _: OperatorStateMetadataV1 => ""
+        case v2: OperatorStateMetadataV2 => v2.operatorPropertiesJson
+      }
+      operatorStateMetadata.stateStoreInfo.map { stateStoreMetadata =>
+        StateMetadataTableEntry(operatorStateMetadata.operatorInfo.operatorId,
+          operatorStateMetadata.operatorInfo.operatorName,
           stateStoreMetadata.storeName,
           stateStoreMetadata.numPartitions,
           if (batchIds.nonEmpty) batchIds.head else -1,
           if (batchIds.nonEmpty) batchIds.last else -1,
+          operatorProperties,
           stateStoreMetadata.numColsPrefixKey
         )
       }

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,6 +189,15 @@ class IncrementalExecution(
 
   object WriteStatefulOperatorMetadataRule extends SparkPlanPartialRule {
     override val rule: PartialFunction[SparkPlan, SparkPlan] = {
+      case tws: TransformWithStateExec if isFirstBatch =>
+        val metadata = tws.operatorStateMetadata()
+        // use a subdirectory for v2
+        val metadataPath = OperatorStateMetadataV2.metadataFilePath(new Path(
+          checkpointLocation, tws.getStateInfo.operatorId.toString))
+        val operatorStateMetadataLog = new OperatorStateMetadataLog(sparkSession,
+          metadataPath.toString)
+        operatorStateMetadataLog.add(currentBatchId, metadata)
+        tws
       case stateStoreWriter: StateStoreWriter if isFirstBatch =>
         val metadata = stateStoreWriter.operatorStateMetadata()
         val metadataWriter = new OperatorStateMetadataWriter(new Path(
@@ -454,11 +463,11 @@ class IncrementalExecution(
               new Path(checkpointLocation).getParent.toString,
               new SerializableConfiguration(hadoopConf))
             val opMetadataList = reader.allOperatorStateMetadata
-            ret = opMetadataList.map { operatorMetadata =>
-              val metadataInfoV1 = operatorMetadata
-                .asInstanceOf[OperatorStateMetadataV1]
-                .operatorInfo
-              metadataInfoV1.operatorId -> metadataInfoV1.operatorName
+            ret = opMetadataList.map {
+              case OperatorStateMetadataV1(operatorInfo, _) =>
+                operatorInfo.operatorId -> operatorInfo.operatorName
+              case OperatorStateMetadataV2(operatorInfo, _, _) =>
+                operatorInfo.operatorId -> operatorInfo.operatorName
             }.toMap
           } catch {
             case e: Exception =>
@@ -495,8 +504,8 @@ class IncrementalExecution(
 
       // The two rules below don't change the plan but can cause the side effect that
       // metadata/schema is written in the checkpoint directory of stateful operator.
-      planWithStateOpId transform StateSchemaValidationRule.rule
-      planWithStateOpId transform WriteStatefulOperatorMetadataRule.rule
+      val schemaValidatedPlan = planWithStateOpId transform StateSchemaValidationRule.rule
+      schemaValidatedPlan transform WriteStatefulOperatorMetadataRule.rule
 
       simulateWatermarkPropagation(planWithStateOpId)
       planWithStateOpId transform WatermarkPropagationRule.rule

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

@@ -0,0 +1,69 @@
+/*
+ * 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 java.io.{BufferedReader, InputStream, InputStreamReader, OutputStream}
+import java.nio.charset.StandardCharsets
+import java.nio.charset.StandardCharsets._
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.FSDataOutputStream
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadata, OperatorStateMetadataV1, OperatorStateMetadataV2}
+import org.apache.spark.sql.internal.SQLConf
+
+
+class OperatorStateMetadataLog(
+    hadoopConf: Configuration,
+    path: String,
+    metadataCacheEnabled: Boolean = false)
+  extends HDFSMetadataLog[OperatorStateMetadata](hadoopConf, path, metadataCacheEnabled) {
+
+  def this(sparkSession: SparkSession, path: String) = {
+    this(
+      sparkSession.sessionState.newHadoopConf(),
+      path,
+      metadataCacheEnabled = sparkSession.sessionState.conf.getConf(
+        SQLConf.STREAMING_METADATA_CACHE_ENABLED)
+    )
+  }
+
+  override protected def serialize(metadata: OperatorStateMetadata, out: OutputStream): Unit = {
+    val fsDataOutputStream = out.asInstanceOf[FSDataOutputStream]
+    fsDataOutputStream.write(s"v${metadata.version}\n".getBytes(StandardCharsets.UTF_8))
+    metadata.version match {
+      case 1 =>
+        OperatorStateMetadataV1.serialize(fsDataOutputStream, metadata)
+      case 2 =>
+        OperatorStateMetadataV2.serialize(fsDataOutputStream, metadata)
+    }
+  }
+
+  override protected def deserialize(in: InputStream): OperatorStateMetadata = {
+    // called inside a try-finally where the underlying stream is closed in the caller
+    // create buffered reader from input stream
+    val bufferedReader = new BufferedReader(new InputStreamReader(in, UTF_8))
+    // read first line for version number, in the format "v{version}"
+    val version = bufferedReader.readLine()
+    version match {
+      case "v1" => OperatorStateMetadataV1.deserialize(bufferedReader)
+      case "v2" => OperatorStateMetadataV2.deserialize(bufferedReader)
+    }
+  }
+}

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

@@ -230,7 +230,8 @@ case class StreamingSymmetricHashJoinExec(
   override def operatorStateMetadata(): OperatorStateMetadata = {
     val info = getStateInfo
     val operatorInfo = OperatorInfoV1(info.operatorId, shortName)
-    val stateStoreInfo = stateStoreNames.map(StateStoreMetadataV1(_, 0, info.numPartitions)).toArray
+    val stateStoreInfo: Array[StateStoreMetadata] =
+      stateStoreNames.map(StateStoreMetadataV1(_, 0, info.numPartitions)).toArray
     OperatorStateMetadataV1(operatorInfo, stateStoreInfo)
   }
 

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

@@ -23,6 +23,10 @@ import scala.jdk.CollectionConverters.CollectionHasAsScala
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
+import org.json4s.JsonAST.JValue
+import org.json4s.JsonDSL._
+import org.json4s.JString
+import org.json4s.jackson.JsonMethods.{compact, render}
 
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
@@ -95,6 +99,8 @@ case class TransformWithStateExec(
     }
   }
 
+  override def operatorStateMetadataVersion: Int = 2
+
   private def getDriverProcessorHandle: DriverStatefulProcessorHandleImpl = {
     val driverProcessorHandle = new DriverStatefulProcessorHandleImpl
     statefulProcessor.setHandle(driverProcessorHandle)
@@ -382,7 +388,25 @@ case class TransformWithStateExec(
   private def validateSchemas(
       oldSchema: List[ColumnFamilySchema],
       newSchema: List[ColumnFamilySchema]): Unit = {
-    // TODO: Implement logic that allows for schema evolution
+    // TODO: Implement logic that allows for schema validation and evolution
+  }
+
+  /** Metadata of this stateful operator and its states stores. */
+  override def operatorStateMetadata(): OperatorStateMetadata = {
+    val info = getStateInfo
+    val operatorInfo = OperatorInfoV1(info.operatorId, shortName)
+    // stateSchemaFilePath should be populated at this point
+    assert(info.stateSchemaPath.isDefined)
+    val stateStoreInfo: Array[StateStoreMetadata] =
+      Array(StateStoreMetadataV2(
+        StateStoreId.DEFAULT_STORE_NAME, 0, info.numPartitions, info.stateSchemaPath.get))
+
+    val operatorPropertiesJson: JValue =
+      ("timeMode" -> JString(timeMode.toString)) ~
+      ("outputMode" -> JString(outputMode.toString))
+
+    val json = compact(render(operatorPropertiesJson))
+    OperatorStateMetadataV2(operatorInfo, stateStoreInfo, json)
   }
 
   private def stateSchemaFilePath(storeName: Option[String] = None): Path = {

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

@@ -33,7 +33,7 @@ import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, Metadata
 /**
  * Metadata for a state store instance.
  */
-trait StateStoreMetadata {
+trait StateStoreMetadata extends Serializable {
   def storeName: String
   def numColsPrefixKey: Int
   def numPartitions: Int
@@ -42,6 +42,21 @@ trait StateStoreMetadata {
 case class StateStoreMetadataV1(storeName: String, numColsPrefixKey: Int, numPartitions: Int)
   extends StateStoreMetadata
 
+case class StateStoreMetadataV2(
+    storeName: String,
+    numColsPrefixKey: Int,
+    numPartitions: Int,
+    stateSchemaFilePath: String)
+  extends StateStoreMetadata
+
+object StateStoreMetadataV2 {
+  private implicit val formats: Formats = Serialization.formats(NoTypeHints)
+
+  @scala.annotation.nowarn
+  private implicit val manifest = Manifest
+    .classType[StateStoreMetadataV2](implicitly[ClassTag[StateStoreMetadataV2]].runtimeClass)
+}
+
 /**
  * Information about a stateful operator.
  */
@@ -54,14 +69,25 @@ case class OperatorInfoV1(operatorId: Long, operatorName: String) extends Operat
 
 trait OperatorStateMetadata {
   def version: Int
+
+  def operatorInfo: OperatorInfo
+
+  def stateStoreInfo: Array[StateStoreMetadata]
 }
 
 case class OperatorStateMetadataV1(
     operatorInfo: OperatorInfoV1,
-    stateStoreInfo: Array[StateStoreMetadataV1]) extends OperatorStateMetadata {
+    stateStoreInfo: Array[StateStoreMetadata]) extends OperatorStateMetadata {
   override def version: Int = 1
 }
 
+case class OperatorStateMetadataV2(
+    operatorInfo: OperatorInfoV1,
+    stateStoreInfo: Array[StateStoreMetadata],
+    operatorPropertiesJson: String) extends OperatorStateMetadata {
+  override def version: Int = 2
+}
+
 object OperatorStateMetadataV1 {
 
   private implicit val formats: Formats = Serialization.formats(NoTypeHints)
@@ -84,6 +110,27 @@ object OperatorStateMetadataV1 {
   }
 }
 
+object OperatorStateMetadataV2 {
+  private implicit val formats: Formats = Serialization.formats(NoTypeHints)
+
+  @scala.annotation.nowarn
+  private implicit val manifest = Manifest
+    .classType[OperatorStateMetadataV2](implicitly[ClassTag[OperatorStateMetadataV2]].runtimeClass)
+
+  def metadataFilePath(stateCheckpointPath: Path): Path =
+    new Path(new Path(new Path(stateCheckpointPath, "v2"), "_metadata"), "metadata")
+
+  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.
  */
@@ -114,7 +161,9 @@ class OperatorStateMetadataWriter(stateCheckpointPath: Path, hadoopConf: Configu
 }
 
 /**
- * Read OperatorStateMetadata from the state checkpoint directory.
+ * Read OperatorStateMetadata from the state checkpoint directory. This class will only be
+ * used to read OperatorStateMetadataV1.
+ * OperatorStateMetadataV2 will be read by the OperatorStateMetadataLog.
  */
 class OperatorStateMetadataReader(stateCheckpointPath: Path, hadoopConf: Configuration) {