[SPARK-47363][SS] Initial State without state reader implementation for State API v2. #45467

jingz-db · 2024-03-11T21:21:21Z

What changes were proposed in this pull request?

This PR adds support for users to provide a Dataframe that can be used to instantiate state for the query in the first batch for arbitrary state API v2.

Note that populating the initial state will only happen for the first batch of the new streaming query. Trying to re-initialize state for the same grouping key will result in an error.

Why are the changes needed?

These changes are needed to support initial state. The changes are part of the work around adding new stateful streaming operator for arbitrary state mgmt that provides a bunch of new features listed in the SPIP JIRA here - https://issues.apache.org/jira/browse/SPARK-45939

Does this PR introduce any user-facing change?

Yes.
This PR introduces a new function:

def transformWithState(
      statefulProcessor: StatefulProcessorWithInitialState[K, V, U, S],
      timeoutMode: TimeoutMode,
      outputMode: OutputMode,
      initialState: KeyValueGroupedDataset[K, S]): Dataset[U]

How was this patch tested?

Unit tests in TransformWithStateWithInitialStateSuite

Was this patch authored or co-authored using generative AI tooling?

No

anishshri-db · 2024-03-13T17:30:03Z

sql/api/src/main/scala/org/apache/spark/sql/streaming/StatefulProcessor.scala

@@ -85,3 +85,21 @@ private[sql] trait StatefulProcessor[K, I, O] extends Serializable {
    statefulProcessorHandle
  }
 }
+
+/**
+ * Similar usage as StatefulProcessor. Represents the arbitrary stateful logic that needs to


Maybe reword this - Stateful processor with support for specifying initial state. Accepts a user-defined type as initial state to be initialized in the first batch. This can be used for starting a new streaming query with existing state from a previous streaming query ?

anishshri-db · 2024-03-13T17:34:15Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

@@ -665,7 +665,8 @@ class KeyValueGroupedDataset[K, V] private[sql](
      outputMode: OutputMode = OutputMode.Append()): Dataset[U] = {
    Dataset[U](
      sparkSession,
-      TransformWithState[K, V, U](
+      // The last K type is only to silence compiler error


Any way to avoid this ?

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

anishshri-db · 2024-03-13T17:38:40Z

Error seems relevant on the MIMA checks -

problems with Sql module: 
method transformWithState(org.apache.spark.sql.streaming.StatefulProcessorWithInitialState,org.apache.spark.sql.streaming.TimeoutMode,org.apache.spark.sql.streaming.OutputMode,org.apache.spark.sql.KeyValueGroupedDataset,org.apache.spark.sql.Encoder,org.apache.spark.sql.Encoder)org.apache.spark.sql.Dataset in class org.apache.spark.sql.KeyValueGroupedDataset does not have a correspondent in client version

we probably need to update the Connect variants as well

anishshri-db · 2024-03-13T17:40:31Z

...core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateInitialStateSuite.scala

+      key: String,
+      initialState: (String, Double)): Unit = {
+    val initStateVal = initialState._2
+    _valState.update(initStateVal)


Can we simulate an actual case class for initial state that stores list/map and/or iterator for list values/iterator for map key-values ?

HeartSaVioR · 2024-03-14T05:12:35Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

+   *                          the query in the first batch.
+   *
+   */
+  def transformWithState[U: Encoder, S: Encoder](


private[sql]

We want to defer exposing the API to public till we complete the work.

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

anishshri-db · 2024-03-21T00:46:55Z

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

-      child.execute().mapPartitionsWithStateStore[InternalRow](
+    if (hasInitialState) {
+      val storeConf = new StateStoreConf(session.sqlContext.sessionState.conf)
+      val hadoopConfBroadcast = sparkContext.broadcast(


Why do we need to do this ?

I am not 100% percent sure, but this will distribute the read-only variable hadoopConf to all executors - similar as here:

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

Lines 55 to 57 in 74a9c6c

// A Hadoop Configuration can be about 10 KB, which is pretty big, so broadcast it

private val hadoopConfBroadcast = session.sparkContext.broadcast(

new SerializableConfiguration(session.sessionState.newHadoopConf()))

Yeah there is a code comment. The practice seems to be that it's better to use broadcast rather than task serialization as it could be huge.

anishshri-db · 2024-03-21T17:41:38Z

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

-      child.execute().mapPartitionsWithStateStore[InternalRow](
+    if (hasInitialState) {
+      val storeConf = new StateStoreConf(session.sqlContext.sessionState.conf)
+      val hadoopConfBroadcast =


I mean this was only needed for the batch support part right ?

We will also need this for StateStore.get here: https://github.com/apache/spark/blob/40465b6760fb120c9cc3ac1a4ee42a82843f4bc5/sql/[…]ache/spark/sql/execution/streaming/TransformWithStateExec.scala

HeartSaVioR

Not yet reviewed the test suite, though I guess Anish has reviewed in detail.

HeartSaVioR · 2024-03-27T04:54:03Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

+        child = logicalPlan,
+        initialState.groupingAttributes,
+        initialState.dataAttributes,
+        initialState.queryExecution.logical


Shall we follow the practice we did in flatMapGroupsWithState for safeness sake?

initialState.queryExecution.analyzed

HeartSaVioR · 2024-03-27T04:59:29Z

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

@@ -268,11 +268,13 @@ class IncrementalExecution(
        )

      case t: TransformWithStateExec =>
+        val hasInitialState = (isFirstBatch && t.hasInitialState)


I don't think we want to allow adding state in the middle of the query lifecycle. Here isFirstBatch does not mean batch ID = 0 but mean this is the first batch in this query run.

This should follow the above logic we did for FlatMapGroupsWithStateExec, currentBatchId == 0L.

Please let me know if this is a different functionality than we had in flatMapGroupsWithState.

HeartSaVioR · 2024-03-27T05:07:03Z

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

-      child.execute().mapPartitionsWithStateStore[InternalRow](
+    if (hasInitialState) {
+      val storeConf = new StateStoreConf(session.sqlContext.sessionState.conf)
+      val hadoopConfBroadcast = sparkContext.broadcast(


Yeah there is a code comment. The practice seems to be that it's better to use broadcast rather than task serialization as it could be huge.

HeartSaVioR · 2024-03-27T05:09:18Z

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

+            processData(store, singleIterator)
+        }
+      } else {
+        // If the query is running in batch mode, we need to create a new StateStore and instantiate


nit: apply the same practice while we are here? broadcast

HeartSaVioR · 2024-03-27T05:17:57Z

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

+              useMultipleValuesPerKey = true)
+            val store = stateStoreProvider.getStore(0)
+
+            processDataWithInitialState(store, childDataIterator, initStateIterator)


We close the state store and state store provider in batch codepath (see below). Shall we do that here as well?

Also, this is a good representation that we have duplicated code. two batch parts have similarity on spinning up state store provider and state store, and also closing them. That could be extracted out.

Good advice! Refactored duplicated codes into initNewStateStoreAndProcessData().

HeartSaVioR · 2024-03-27T05:19:56Z

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

+
+    // Check if is first batch
+    // Only process initial states for first batch
+    if (processorHandle.getQueryInfo().getBatchId == 0) {


OK I see we have multiple checks. Though still better to change the condition in IncrementalExecution as reader can misunderstand that there are inconsistency between flatMapGroupsWithState and transformWithState.

HeartSaVioR

+1 pending CI

HeartSaVioR · 2024-03-28T05:49:52Z

CI failure isn't related - only pyspark-connect failed.

HeartSaVioR · 2024-03-28T05:49:59Z

Thanks! Merging to master.

…or State API v2 ### What changes were proposed in this pull request? This PR adds support for users to provide a Dataframe that can be used to instantiate state for the query in the first batch for arbitrary state API v2. Note that populating the initial state will only happen for the first batch of the new streaming query. Trying to re-initialize state for the same grouping key will result in an error. ### Why are the changes needed? These changes are needed to support initial state. The changes are part of the work around adding new stateful streaming operator for arbitrary state mgmt that provides a bunch of new features listed in the SPIP JIRA here - https://issues.apache.org/jira/browse/SPARK-45939 ### Does this PR introduce _any_ user-facing change? Yes. This PR introduces a new function: ``` def transformWithState( statefulProcessor: StatefulProcessorWithInitialState[K, V, U, S], timeoutMode: TimeoutMode, outputMode: OutputMode, initialState: KeyValueGroupedDataset[K, S]): Dataset[U] ``` ### How was this patch tested? Unit tests in `TransformWithStateWithInitialStateSuite` ### Was this patch authored or co-authored using generative AI tooling? No Closes apache#45467 from jingz-db/initial-state-state-v2. Lead-authored-by: jingz-db <jing.zhan@databricks.com> Co-authored-by: Jing Zhan <135738831+jingz-db@users.noreply.github.com> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

…mWithStateInPandas ### What changes were proposed in this pull request? This PR adds support for users to provide a Dataframe that can be used to instantiate state for the query in the first batch for arbitrary state API v2 in Python. The Scala PR for supporting initial state is here: #45467 We propose to create a new PythonRunner that handles initial state specifically for TransformWithStateInPandas. From JVM, we coGroup input rows and initial state rows on the same grouping key. Then we create a new row that contains one row in the input rows iterator and one row in the initial state iterator, and send the new grouped row to Py4j. Inside the python worker, we deserialize the grouped row into input rows and initial state rows separately and input those into `handleInitialState` and `handleInputRows`. We will launch a python worker for each partition that has a non-empty input rows in either input rows or initial states. This will guarantee all keys in the initial state will be processed even if they do not appear in the first batch or they don't lie in the same partition with keys in the first batch. ### Why are the changes needed? We need to couple the API as we support initial state handling in Scala. ### Does this PR introduce _any_ user-facing change? Yes. This PR introduces a new API in the `StatefulProcessor` which allows users to define their own udf for processing initial state: ``` def handleInitialState( self, key: Any, initialState: "PandasDataFrameLike" ) -> None: ``` The implementation of this function is optional. If not defined, then it will act as no-op. ### How was this patch tested? Unit tests & integration tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48005 from jingz-db/python-init-state-impl. Authored-by: jingz-db <jing.zhan@databricks.com> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

github-actions bot added SQL STRUCTURED STREAMING labels Mar 11, 2024

jingz-db changed the title ~~WIP~~ [SS] Initial State without state reader implementation for State API v2. Mar 11, 2024

jingz-db force-pushed the initial-state-state-v2 branch from 8ec379a to 8aac855 Compare March 12, 2024 01:21

jingz-db marked this pull request as ready for review March 12, 2024 01:21

jingz-db changed the title ~~[SS] Initial State without state reader implementation for State API v2.~~ [SS][SPARK-47363] Initial State without state reader implementation for State API v2. Mar 12, 2024

jingz-db force-pushed the initial-state-state-v2 branch from cd8b827 to 8fbd501 Compare March 13, 2024 00:24

anishshri-db reviewed Mar 13, 2024

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Mar 13, 2024

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Mar 13, 2024

View reviewed changes

HeartSaVioR reviewed Mar 14, 2024

View reviewed changes

jingz-db added 14 commits March 14, 2024 11:52

WIP

190bdbb

everything but batch works

ad2a3bd

process init state, need to support batch

bd9d2f9

batch with initial state succeed

5587b08

rebase on master

8a863a1

change in IncrementalExecution

4aebe30

fix batch

fd31d2b

fix scala style

c8cf3c3

add test variable

bd5d908

refactor test suite

e2b1875

rebase on mapstate

0e5e997

partially resolve comments, will refactor test case

7e56085

fix mima?

6465521

resolve conflicts

9f18601

jingz-db force-pushed the initial-state-state-v2 branch from 25d7bab to 9f18601 Compare March 14, 2024 19:33

jingz-db requested a review from anishshri-db March 14, 2024 19:51

jingz-db requested a review from anishshri-db March 20, 2024 22:58

anishshri-db reviewed Mar 21, 2024

View reviewed changes

common/utils/src/main/resources/error/error-classes.json Outdated Show resolved Hide resolved

anishshri-db reviewed Mar 21, 2024

View reviewed changes

resolve comments

40465b6

jingz-db requested a review from anishshri-db March 21, 2024 17:40

anishshri-db reviewed Mar 21, 2024

View reviewed changes

anishshri-db approved these changes Mar 21, 2024

View reviewed changes

jingz-db and others added 2 commits March 21, 2024 11:25

hadoop broadcast

d91817a

Merge branch 'master' into initial-state-state-v2

b3394d0

jingz-db requested a review from HeartSaVioR March 26, 2024 21:11

HeartSaVioR reviewed Mar 27, 2024

View reviewed changes

beliefer changed the title ~~[SS][SPARK-47363] Initial State without state reader implementation for State API v2.~~ [SPARK-47363][SS] Initial State without state reader implementation for State API v2. Mar 27, 2024

github-actions bot added the BUILD label Mar 27, 2024

jingz-db force-pushed the initial-state-state-v2 branch from e528dfd to b3394d0 Compare March 27, 2024 20:59

github-actions bot removed the BUILD label Mar 27, 2024

jingz-db added 3 commits March 27, 2024 14:06

resolve comments

695a02b

reuse code in batch process

1ca648f

scala style

88d8de7

jingz-db requested a review from HeartSaVioR March 27, 2024 22:53

HeartSaVioR approved these changes Mar 28, 2024

View reviewed changes

HeartSaVioR closed this in 4d72be3 Mar 28, 2024

dongjoon-hyun mentioned this pull request Jul 31, 2024

[SPARK-49070][SS][SQL] TransformWithStateExec.initialState is rewritten incorrectly to produce invalid query plan #47546

Closed

jingz-db mentioned this pull request Sep 10, 2024

[SPARK-49601][SS][PYTHON] Support Initial State Handling for TransformWithStateInPandas #48005

Closed

jingz-db mentioned this pull request Oct 29, 2024

[SPARK-50152][SS] Support handleInitialState with state data source reader #48686

Closed

	// A Hadoop Configuration can be about 10 KB, which is pretty big, so broadcast it
	private val hadoopConfBroadcast = session.sparkContext.broadcast(
	new SerializableConfiguration(session.sessionState.newHadoopConf()))

[SPARK-47363][SS] Initial State without state reader implementation for State API v2. #45467

[SPARK-47363][SS] Initial State without state reader implementation for State API v2. #45467

Uh oh!

Conversation

jingz-db commented Mar 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

anishshri-db commented Mar 13, 2024

Uh oh!

anishshri-db Mar 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HeartSaVioR left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HeartSaVioR left a comment

Choose a reason for hiding this comment

Uh oh!

HeartSaVioR commented Mar 28, 2024

Uh oh!

HeartSaVioR commented Mar 28, 2024

Uh oh!

Uh oh!

jingz-db commented Mar 11, 2024 •

edited

Loading

anishshri-db Mar 13, 2024 •

edited

Loading