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

jingz-db · 2024-09-05T23:15:56Z

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.

sql/core/src/main/java/org/apache/spark/sql/execution/streaming/StateMessage.proto

...main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala

bogao007 · 2024-09-27T21:23:47Z

...main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala

      case _ =>
        throw new IllegalArgumentException("Invalid method call")
    }
  }

+  private def handleStatefulProcessorUtilRequest(message: UtilsCallCommand): Unit = {


Should we add some scala unit tests for these 2 new APIs?

python/pyspark/sql/streaming/stateful_processor_api_client.py

...main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala

python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py

bogao007 · 2024-10-07T22:36:13Z

python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py

+        yield pd.DataFrame({"id": key, "value": str(accumulated_value)})
+
+    def handleInitialState(self, key, initialState) -> None:
+        initVal = initialState.at[0, "initVal"]


Can we add verifications on the initVal here?

bogao007 · 2024-10-11T20:58:01Z

python/pyspark/sql/pandas/group_ops.py

-        if isinstance(outputStructType, str):
-            outputStructType = cast(StructType, _parse_datatype_string(outputStructType))
+        def transformWithStateWithInitStateUDF(
+                statefulProcessorApiClient: StatefulProcessorApiClient,


Nit: fix the indent.

bogao007 · 2024-10-11T21:28:35Z

python/pyspark/sql/streaming/stateful_processor.py

@@ -234,3 +234,12 @@ def close(self) -> None:
        operations.
        """
        ...
+
+    def handleInitialState(
+            self, key: Any, initialState: "PandasDataFrameLike"


Nit: indent.

bogao007 · 2024-10-11T21:30:24Z

python/pyspark/sql/pandas/group_ops.py

@@ -402,6 +404,9 @@ def transformWithStateInPandas(
            The output mode of the stateful processor.
        timeMode : str
            The time mode semantics of the stateful processor for timers and TTL.
+        initialState: "GroupedData"


Let's use something like below to represent the actual type.

:class:`pyspark.sql.types.DataType`

bogao007 · 2024-10-11T21:50:54Z

python/pyspark/sql/pandas/group_ops.py

+                    StatefulProcessorHandleState.INITIALIZED
+                )
+
+            # if we don't have state for the given key but in initial state,


I guess you wanted to say if we only have initial state but don't have input rows for the given key, the inputRows iterator could be empty?

bogao007 · 2024-10-11T21:54:19Z

...main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala

@@ -168,13 +169,10 @@ class TransformWithStateInPandasStateServer(
        val requestedState = message.getSetHandleState.getState
        requestedState match {
          case HandleState.CREATED =>
-            logInfo(log"set handle state to Created")


Why do we need to remove these log lines?

bogao007 · 2024-10-11T21:55:01Z

python/pyspark/sql/streaming/stateful_processor_api_client.py

+            return False
+        else:
+            # TODO(SPARK-49233): Classify user facing errors.
+            raise PySparkRuntimeError(f"Error getting batch id: " f"{response_message[1]}")


Should we have a better error message?

bogao007 · 2024-10-11T21:58:12Z

...ain/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala

+  with PythonArrowInput[GroupedInType] {
+
+  override protected lazy val schema: StructType = new StructType()
+    .add("state", dataSchema)


Maybe use data or inputDate to differentiate with initState?

bogao007 · 2024-10-11T22:12:10Z

python/pyspark/sql/pandas/group_ops.py

+        ) -> Iterator["PandasDataFrameLike"]:
+            handle = StatefulProcessorHandle(statefulProcessorApiClient)
+
+            if statefulProcessorApiClient.handle_state == StatefulProcessorHandleState.CREATED:


There's something not very clear to me here, could you help me understand more?

We only call handleInitialState when handle state is CREATED, but after we processed the initial state of the first grouping key, we update the state to be INITIALIZED. Wouldn't that skip the initial state for other grouping keys?

If my understanding is correct, we should move the handleInitialState outside the handle state check, do it after the init call.

bogao007 · 2024-10-11T22:32:10Z

python/pyspark/sql/pandas/group_ops.py

+                statefulProcessorApiClient: StatefulProcessorApiClient,
+                key: Any,
+                inputRows: Iterator["PandasDataFrameLike"],
+                # for non first batch, initialStates will be None


For non first batch, would initialStates be None or empty?

bogao007 · 2024-10-11T22:37:00Z

python/pyspark/sql/pandas/group_ops.py

+                inputRows: Iterator["PandasDataFrameLike"],
+                # for non first batch, initialStates will be None
+                initialStates: Iterator["PandasDataFrameLike"] = None
+        ) -> Iterator["PandasDataFrameLike"]:


Can we add some commentss on the possible input combinations that we need to handle in this udf for people to understand easier? IIUC there should be 3 cases:

Both inputRows and initialStates contain data. This would only happen in the first batch and the associated grouping key contains both input data and initial state.

Only inputRows contains data. This could happen when either the grouping key doesn't have any initial state to process or it's non first batch.

Only initialStates contains data. This could happen when the grouping key doesn't have any associated input data but it has initial state to process.

github-actions bot added SQL STRUCTURED STREAMING PYTHON labels Sep 5, 2024

jingz-db changed the title ~~[SS][PYTHON] Support Initial State for TransformWithStateInPandas~~ [SS][PYTHON] Support Initial State Handling for TransformWithStateInPandas Sep 5, 2024

jingz-db marked this pull request as ready for review September 9, 2024 16:48

github-actions bot added the DOCS label Sep 9, 2024

jingz-db changed the title ~~[SS][PYTHON] Support Initial State Handling for TransformWithStateInPandas~~ [SPARK-49601][SS][PYTHON] Support Initial State Handling for TransformWithStateInPandas Sep 11, 2024

resolve conflicts with master

099d827

jingz-db force-pushed the python-init-state-impl branch from 253e56d to 099d827 Compare September 12, 2024 19:01

propagate grouping attributes correctly

6fb43b5

bogao007 reviewed Sep 27, 2024

View reviewed changes

jingz-db and others added 7 commits October 2, 2024 16:14

Merge branch 'master' into python-init-state-impl

5cace96

restore protoc gen files

6bf6e5a

resolve conflicts in stateful_processor_api_client

569a7f1

add scala unit tests

41cc007

add logs

6c9b94e

read init state per partition

fd8cd0f

Merge branch 'master' into python-init-state-impl

42829cb

bogao007 reviewed Oct 7, 2024

View reviewed changes

jingz-db added 3 commits October 9, 2024 15:50

fix hanging, send single init state per row

ff73636

multiple batches

6db48d9

new python runner compiling, init data is missing

9e3bc77

github-actions bot added the CORE label Oct 10, 2024

test passed

2876875

jingz-db requested a review from bogao007 October 11, 2024 09:27

bogao007 reviewed Oct 11, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

jingz-db commented Sep 5, 2024 •

edited

Loading

bogao007 Sep 27, 2024

bogao007 Oct 7, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

bogao007 Oct 11, 2024

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

Are you sure you want to change the base?

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

Conversation

jingz-db commented Sep 5, 2024 • edited Loading

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?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jingz-db commented Sep 5, 2024 •

edited

Loading