[SPARK-57228][SS] Support transformWithState in Real-Time Mode

common/utils/src/main/resources/error/error-conditions.json

@@ -7787,6 +7787,11 @@
         "message" : [
           "The following session configuration(s) are incompatible with Real-Time Mode: <invalidReasons>. Update or remove them and restart the query."
         ]
+      },
+      "STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED" : {
+        "message" : [
+          "Streaming queries in real-time mode cannot include stateful operators (e.g. aggregate, deduplicate, transformWithState) before a union. Please restructure your query to apply the union before any stateful operations."
+        ]
       }
     },
     "sqlState" : "0A000"

connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeSuite.scala

@@ -18,23 +18,50 @@
 package org.apache.spark.sql.kafka010
 
 import java.util.UUID
+import java.util.regex.Pattern
 
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.time.SpanSugar._
 
 import org.apache.spark.{SparkConf, SparkContext, SparkIllegalStateException}
-import org.apache.spark.sql.execution.datasources.v2.LowLatencyClock
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.datasources.v2.{LowLatencyClock, RealTimeStreamScanExec}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
 import org.apache.spark.sql.execution.streaming._
+import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateExec
 import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream}
 import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.kafka010.consumer.KafkaDataConsumer
-import org.apache.spark.sql.streaming.{StreamingQuery, Trigger}
+import org.apache.spark.sql.streaming.{OutputMode, StatefulProcessor, StreamingQuery, TimeMode,
+  TimerValues, Trigger, TTLConfig, ValueState}
 import org.apache.spark.sql.streaming.OutputMode.Update
 import org.apache.spark.sql.streaming.util.GlobalSingletonManualClock
 import org.apache.spark.sql.test.TestSparkSession
 import org.apache.spark.util.SystemClock
 
+private class KafkaRunningCountStatefulProcessor
+  extends StatefulProcessor[String, String, (String, Long)] {
+
+  @transient private var countState: ValueState[Long] = _
+
+  override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {
+    countState = getHandle.getValueState(
+      "count", Encoders.scalaLong, TTLConfig.NONE)
+  }
+
+  override def handleInputRows(
+      key: String,
+      inputRows: Iterator[String],
+      timerValues: TimerValues): Iterator[(String, Long)] = {
+    inputRows.map { _ =>
+      val count = Option(countState.get()).getOrElse(0L) + 1L
+      countState.update(count)
+      (key, count)
+    }
+  }
+}
+
 class KafkaRealTimeModeSuite
   extends KafkaSourceTest
     with Matchers {
@@ -164,6 +191,97 @@ class KafkaRealTimeModeSuite
       WaitUntilCurrentBatchProcessed)
   }
 
+  test("transformWithState uses a pipelined shuffle and recovers RocksDB changelog state") {
+    withSQLConf(
+        SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+        "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" -> "true") {
+      val topic = newTopic()
+      testUtils.createTopic(topic, partitions = 2)
+      testUtils.sendMessages(topic, Array("a", "a"), Some(0))
+      testUtils.sendMessages(topic, Array("b"), Some(1))
+
+      val counts = spark.readStream
+        .format("kafka")
+        .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+        .option("subscribe", topic)
+        .option("startingOffsets", "earliest")
+        .load()
+        .selectExpr("CAST(value AS STRING)")
+        .as[String]
+        .groupByKey(identity)
+        .transformWithState(
+          new KafkaRunningCountStatefulProcessor,
+          TimeMode.None(),
+          Update)
+
+      testStream(counts, Update, sink = new ContinuousMemorySink())(
+        StartStream(),
+        CheckAnswerWithTimeout(60000, ("a", 1L), ("a", 2L), ("b", 1L)),
+        Execute { q =>
+          val plan = q.lastExecution.executedPlan
+          val statefulOperators = plan.collect { case t: TransformWithStateExec => t }
+          assert(statefulOperators.size == 1, plan)
+          assert(statefulOperators.head.isRealTimeMode, plan)
+
+          val exchanges = plan.collect { case s: ShuffleExchangeExec => s }
+          assert(exchanges.nonEmpty, plan)
+          assert(exchanges.forall(_.pipelined), plan)
+        },
+        WaitUntilCurrentBatchProcessed,
+        StopStream,
+        new ExternalAction() {
+          override def runAction(): Unit = {
+            testUtils.sendMessages(topic, Array("a"), Some(0))
+            testUtils.sendMessages(topic, Array("b", "b"), Some(1))
+          }
+        },
+        StartStream(),
+        CheckAnswerWithTimeout(
+          60000,
+          ("a", 1L), ("a", 2L), ("b", 1L),
+          ("a", 3L), ("b", 2L), ("b", 3L)),
+        WaitUntilCurrentBatchProcessed)
+    }
+  }
+
+  test("transformWithState remains in real-time mode when latestOffset returns null") {
+    val topic = newTopic()
+
+    val counts = spark.readStream
+      .format("kafka")
+      .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+      .option("kafka.metadata.max.age.ms", "1")
+      .option("subscribePattern", "^" + Pattern.quote(topic) + "$")
+      .option("startingOffsets", "earliest")
+      .load()
+      .selectExpr("CAST(value AS STRING)")
+      .as[String]
+      .groupByKey(identity)
+      .transformWithState(
+        new KafkaRunningCountStatefulProcessor,
+        TimeMode.None(),
+        Update)
+
+    testStream(counts, Update, sink = new ContinuousMemorySink())(
+      StartStream(),
+      WaitUntilBatchProcessed(0),
+      Execute { q =>
+        val plan = q.lastExecution.executedPlan
+        assert(plan.collect { case _: RealTimeStreamScanExec => true }.isEmpty, plan)
+        val statefulOperators = plan.collect { case t: TransformWithStateExec => t }
+        assert(statefulOperators.size == 1, plan)
+        assert(statefulOperators.head.isRealTimeMode, plan)
+      },
+      new ExternalAction() {
+        override def runAction(): Unit = {
+          testUtils.createTopic(topic, partitions = 1)
+          testUtils.sendMessages(topic, Array("a"))
+        }
+      },
+      CheckAnswerWithTimeout(60000, ("a", 1L)),
+      StopStream)
+  }
+
   // A simple unit test that reads from Kakfa source, does a simple map and writes to memory
   // sink. Make sure there is no data for a whole batch. Also, after restart the first batch
   // has no data.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala

@@ -657,6 +657,18 @@ object UnsupportedOperationChecker extends Logging {
     if (outputMode != InternalOutputModes.Update) {
       throwRealTimeError("OUTPUT_MODE_NOT_SUPPORTED", Map("outputMode" -> outputMode.toString))
     }
+
+    plan.foreachUp {
+      case u: Union =>
+        // Block stateful operators before union
+        u.foreachUp {
+          case statefulOp @ (_: Aggregate | _: TransformWithState |
+               _: TransformWithStateInPySpark | _: Deduplicate) if statefulOp.isStateful =>
+            throwRealTimeError("STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED", Map.empty)
+          case _ =>
+        }
+      case _ =>
+    }
   }
 
   private def throwRealTimeError(subClass: String, args: Map[String, String]): Unit = {

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -3627,13 +3627,12 @@ object SQLConf {
   val STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1 =
     buildConf("spark.sql.streaming.realTimeMode.dangerouslyAllowCheckpointV1.enabled")
       .internal()
-      .doc("Whether to allow a Real-Time Mode query to start on a checkpoint whose commit log " +
-        "is at version 1. Real-Time Mode re-executes a failed batch, and with checkpoint format " +
+      .doc("Whether to allow a Real-Time Mode query to start with state store checkpoint format " +
+        "version 1. Real-Time Mode re-executes a failed batch, and with checkpoint format " +
         "version 1 the re-execution can reuse the state file names of the partially-written " +
         "failed batch, so starting on a version 1 checkpoint exposes the query to data loss on " +
-        "failure. Format version 2 avoids this with per-batch state store checkpoint ids, which " +
-        "only a commit log at version 2 or above can persist. Escape hatch only; prefer a fresh " +
-        "checkpoint location.")
+        "failure. Format version 2 avoids this with per-batch state store checkpoint ids. " +
+        "Escape hatch only; prefer a fresh checkpoint location.")
       .version("4.3.0")
       .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
       .booleanConf
@@ -4056,6 +4055,16 @@ object SQLConf {
     .booleanConf
     .createWithDefault(true)
 
+  val STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS = buildConf(
+    "spark.sql.streaming.realTimeMode.transformWithState.ttlEvictionIntervalMs")
+    .internal()
+    .doc("The threshold in milliseconds to perform eviction of TTL when using the JVM " +
+      "transformWithState operator with real-time mode.")
+    .version("4.3.0")
+    .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+    .longConf
+    .createWithDefault(1 * 1000)
+
   val STREAMING_ASYNC_PROGRESS_TRACKING_REAL_TIME_MODE_ENABLED_BY_DEFAULT = buildConf(
     "spark.sql.streaming.realTimeMode.asyncProgressTrackingByDefault.enabled")
     .internal()

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationsSuite.scala

@@ -30,8 +30,8 @@ import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.streaming.{GroupStateTimeout, OutputMode}
-import org.apache.spark.sql.types.{IntegerType, LongType, MetadataBuilder}
+import org.apache.spark.sql.streaming.{GroupStateTimeout, OutputMode, StatefulProcessor, TimeMode, TimerValues}
+import org.apache.spark.sql.types.{IntegerType, LongType, MetadataBuilder, StructType}
 
 /** A dummy command for testing unsupported operations. */
 case class DummyCommand() extends LeafCommand
@@ -967,6 +967,58 @@ class UnsupportedOperationsSuite extends SparkFunSuite with SQLHelper {
     )
   }
 
+  assertSupportedForRealTime(
+    "real-time with Scala transformWithState - update mode",
+    scalaTransformWithState(streamRelation),
+    Update
+  )
+
+  assertNotSupportedForRealTime(
+    "real-time with Scala transformWithState on both sides of union - update mode",
+    scalaTransformWithState(streamRelation)
+      .union(scalaTransformWithState(new TestStreamingRelation(attribute.newInstance()))),
+    Update,
+    "STREAMING_REAL_TIME_MODE.STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED"
+  )
+
+  assertSupportedForRealTime(
+    "real-time with batch aggregate before union - update mode",
+    streamRelation
+      .join(Aggregate(Nil, aggExprs("c"), batchRelation), joinType = Inner)
+      .select(attribute)
+      .union(new TestStreamingRelation(attribute.newInstance())),
+    Update
+  )
+
+  private def scalaTransformWithState(child: LogicalPlan): TransformWithState = {
+    val statefulProcessor = new StatefulProcessor[Any, Any, Any] {
+      override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {}
+
+      override def handleInputRows(
+          key: Any,
+          inputRows: Iterator[Any],
+          timerValues: TimerValues): Iterator[Any] = Iterator.empty
+    }
+    val keyEncoder = ExpressionEncoder(new StructType().add("a", IntegerType))
+      .asInstanceOf[ExpressionEncoder[Any]]
+    new TransformWithState(
+      keyDeserializer = attribute,
+      valueDeserializer = attribute,
+      groupingAttributes = Seq(attribute),
+      dataAttributes = Seq(attribute),
+      statefulProcessor = statefulProcessor,
+      timeMode = NoTime,
+      outputMode = Update,
+      keyEncoder = keyEncoder,
+      outputObjAttr = attribute,
+      child = child,
+      hasInitialState = false,
+      initialStateGroupingAttrs = Seq(attribute),
+      initialStateDataAttrs = Seq(attribute),
+      initialStateDeserializer = attribute,
+      initialState = LocalRelation(Seq.empty[Attribute]))
+  }
+
   /*
     =======================================================================================
                                      TESTING FUNCTIONS

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

@@ -86,6 +86,19 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
     }
   }
 
+  /**
+   * Whether this plan reads a streaming source in Real-Time Mode, which is the case when the
+   * relation carries a real-time mode duration -- the same signal that decides whether to plan
+   * a [[org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec]] for it.
+   */
+  private def isRealTimeMode(plan: LogicalPlan): Boolean = {
+    plan.collectLeaves().exists {
+      case s: StreamingDataSourceV2ScanRelation =>
+        s.relation.realTimeModeDuration.isDefined
+      case _ => false
+    }
+  }
+
   /**
    * Plans special cases of limit operators.
    */
@@ -560,20 +573,6 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * [[org.apache.spark.sql.execution.streaming.StreamExecution]]
    */
   object StatefulAggregationStrategy extends Strategy {
-
-    /**
-     * Whether this plan reads a streaming source in Real-Time Mode, which is the case when the
-     * relation carries a real-time mode duration -- the same signal that decides whether to plan
-     * a [[org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec]] for it.
-     */
-    private def isRealTimeMode(plan: LogicalPlan): Boolean = {
-      plan.collectLeaves().exists {
-        case s: StreamingDataSourceV2ScanRelation =>
-          s.relation.realTimeModeDuration.isDefined
-        case _ => false
-      }
-    }
-
     override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case _ if !plan.isStreaming => Nil
 
@@ -928,6 +927,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
           eventTimeWatermarkForEviction = None,
           planLater(child),
           isStreaming = true,
+          isRealTimeMode = isRealTimeMode(plan),
           hasInitialState,
           initialStateGroupingAttrs,
           initialStateDataAttrs,

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

@@ -1626,16 +1626,20 @@ trait SchemaValidationUtils extends Logging {
       stateSchemaDir: Path,
       session: SparkSession,
       operatorStateMetadataVersion: Int = 2,
-      stateStoreEncodingFormat: String = StateStoreEncoding.UnsafeRow.toString
+      stateStoreEncodingFormat: String = StateStoreEncoding.UnsafeRow.toString,
+      isRealTimeMode: Boolean = false
   ): List[StateSchemaValidationResult] = {
     assert(stateSchemaVersion >= 3)
     val usingAvro = stateStoreEncodingFormat == StateStoreEncoding.Avro.toString
     val newSchemas = getColFamilySchemas(usingAvro)
     val newStateSchemaFilePath =
       new Path(stateSchemaDir, s"${batchId}_${UUID.randomUUID().toString}")
     val metadataPath = new Path(info.checkpointLocation, s"${info.operatorId}")
+    // RTM can leave metadata for an uncommitted attempt of the current batch, so only metadata
+    // from the previous committed batch is safe to use for schema validation.
+    val metadataBatchId = if (isRealTimeMode) batchId - 1 else batchId
     val metadataReader = OperatorStateMetadataReader.createReader(
-      metadataPath, hadoopConf, operatorStateMetadataVersion, batchId)
+      metadataPath, hadoopConf, operatorStateMetadataVersion, metadataBatchId)
     val operatorStateMetadata = try {
       metadataReader.read()
     } catch {