[SPARK-50856][SS][PYTHON][CONNECT] Spark Connect Support for TransformWithStateInPandas In Python

dev/sparktestsupport/modules.py

@@ -1095,6 +1095,7 @@ def __hash__(self):
         "pyspark.sql.tests.connect.pandas.test_parity_pandas_udf_scalar",
         "pyspark.sql.tests.connect.pandas.test_parity_pandas_udf_grouped_agg",
         "pyspark.sql.tests.connect.pandas.test_parity_pandas_udf_window",
+        "pyspark.sql.tests.connect.pandas.test_parity_pandas_transform_with_state",
     ],
     excluded_python_implementations=[
         "PyPy"  # Skip these tests under PyPy since they require numpy, pandas, and pyarrow and

python/pyspark/sql/connect/group.py

@@ -41,6 +41,7 @@
 from pyspark.sql.column import Column
 from pyspark.sql.connect.functions import builtin as F
 from pyspark.errors import PySparkNotImplementedError, PySparkTypeError
+from pyspark.sql.streaming.stateful_processor import StatefulProcessor
 
 if TYPE_CHECKING:
     from pyspark.sql.connect._typing import (
@@ -361,6 +362,56 @@ def applyInPandasWithState(
 
     applyInPandasWithState.__doc__ = PySparkGroupedData.applyInPandasWithState.__doc__
 
+    def transformWithStateInPandas(
+        self,
+        statefulProcessor: StatefulProcessor,
+        outputStructType: Union[StructType, str],
+        outputMode: str,
+        timeMode: str,
+        initialState: Optional["GroupedData"] = None,
+        eventTimeColumnName: str = "",
+    ) -> "DataFrame":
+        from pyspark.sql.connect.udf import UserDefinedFunction
+        from pyspark.sql.connect.dataframe import DataFrame
+        from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasUdfUtils
+
+        udf_util = TransformWithStateInPandasUdfUtils(statefulProcessor, timeMode)
+        if initialState is None:
+            udf_obj = UserDefinedFunction(
+                udf_util.transformWithStateUDF,
+                returnType=outputStructType,
+                evalType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
+            )
+            initial_state_plan = None
+            initial_state_grouping_cols = None
+        else:
+            self._df._check_same_session(initialState._df)
+            udf_obj = UserDefinedFunction(
+                udf_util.transformWithStateWithInitStateUDF,
+                returnType=outputStructType,
+                evalType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
+            )
+            initial_state_plan = initialState._df._plan
+            initial_state_grouping_cols = initialState._grouping_cols
+
+        return DataFrame(
+            plan.TransformWithStateInPandas(
+                child=self._df._plan,
+                grouping_cols=self._grouping_cols,
+                function=udf_obj,
+                output_schema=outputStructType,
+                output_mode=outputMode,
+                time_mode=timeMode,
+                event_time_col_name=eventTimeColumnName,
+                cols=self._df.columns,
+                initial_state_plan=initial_state_plan,
+                initial_state_grouping_cols=initial_state_grouping_cols,
+            ),
+            session=self._df._session,
+        )
+
+    transformWithStateInPandas.__doc__ = PySparkGroupedData.transformWithStateInPandas.__doc__
+
     def applyInArrow(
         self, func: "ArrowGroupedMapFunction", schema: Union[StructType, str]
     ) -> "DataFrame":

python/pyspark/sql/connect/plan.py

@@ -2546,6 +2546,74 @@ def plan(self, session: "SparkConnectClient") -> proto.Relation:
         return self._with_relations(plan, session)
 
 
+class TransformWithStateInPandas(LogicalPlan):
+    """Logical plan object for a TransformWithStateInPandas."""
+
+    def __init__(
+        self,
+        child: Optional["LogicalPlan"],
+        grouping_cols: Sequence[Column],
+        function: "UserDefinedFunction",
+        output_schema: Union[DataType, str],
+        output_mode: str,
+        time_mode: str,
+        event_time_col_name: str,
+        cols: List[str],
+        initial_state_plan: Optional["LogicalPlan"],
+        initial_state_grouping_cols: Optional[Sequence[Column]],
+    ):
+        assert isinstance(grouping_cols, list) and all(isinstance(c, Column) for c in grouping_cols)
+        if initial_state_plan is not None:
+            assert isinstance(initial_state_grouping_cols, list) and all(
+                isinstance(c, Column) for c in initial_state_grouping_cols
+            )
+            super().__init__(
+                child, self._collect_references(grouping_cols + initial_state_grouping_cols)
+            )
+        else:
+            super().__init__(child, self._collect_references(grouping_cols))
+        self._grouping_cols = grouping_cols
+        self._output_schema: DataType = (
+            UnparsedDataType(output_schema) if isinstance(output_schema, str) else output_schema
+        )
+        self._output_mode = output_mode
+        self._time_mode = time_mode
+        self._event_time_col_name = event_time_col_name
+        self._function = function._build_common_inline_user_defined_function(*cols)
+        self._initial_state_plan = initial_state_plan
+        self._initial_state_grouping_cols = initial_state_grouping_cols
+
+    def plan(self, session: "SparkConnectClient") -> proto.Relation:
+        assert self._child is not None
+        plan = self._create_proto_relation()
+        plan.group_map.input.CopyFrom(self._child.plan(session))
+        plan.group_map.grouping_expressions.extend(
+            [c.to_plan(session) for c in self._grouping_cols]
+        )
+        plan.group_map.output_mode = self._output_mode
+
+        # fill in initial state related fields
+        if self._initial_state_plan is not None:
+            plan.group_map.initial_input.CopyFrom(self._initial_state_plan.plan(session))
+            assert self._initial_state_grouping_cols is not None
+            plan.group_map.initial_grouping_expressions.extend(
+                [c.to_plan(session) for c in self._initial_state_grouping_cols]
+            )
+
+        # fill in transformWithStateInPandas related fields
+        tws_info = proto.TransformWithStateInfo()
+        tws_info.time_mode = self._time_mode
+        tws_info.event_time_column_name = self._event_time_col_name
+        tws_info.output_schema.CopyFrom(pyspark_types_to_proto_types(self._output_schema))
+
+        plan.group_map.transform_with_state_info.CopyFrom(tws_info)
+
+        # wrap transformWithStateInPandasUdf in a function
+        plan.group_map.func.CopyFrom(self._function.to_plan_udf(session))
+
+        return self._with_relations(plan, session)
+
+
 class PythonUDTF:
     """Represents a Python user-defined table function."""
 

python/pyspark/sql/connect/proto/relations_pb2.pyi

@@ -3613,6 +3613,7 @@ class GroupMap(google.protobuf.message.Message):
     OUTPUT_MODE_FIELD_NUMBER: builtins.int
     TIMEOUT_CONF_FIELD_NUMBER: builtins.int
     STATE_SCHEMA_FIELD_NUMBER: builtins.int
+    TRANSFORM_WITH_STATE_INFO_FIELD_NUMBER: builtins.int
     @property
     def input(self) -> global___Relation:
         """(Required) Input relation for Group Map API: apply, applyInPandas."""
@@ -3654,6 +3655,11 @@ class GroupMap(google.protobuf.message.Message):
     @property
     def state_schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
         """(Optional) The schema for the grouped state."""
+    @property
+    def transform_with_state_info(self) -> global___TransformWithStateInfo:
+        """Below fields are used by TransformWithState and TransformWithStateInPandas
+        (Optional) TransformWithState related parameters.
+        """
     def __init__(
         self,
         *,
@@ -3677,6 +3683,7 @@ class GroupMap(google.protobuf.message.Message):
         output_mode: builtins.str | None = ...,
         timeout_conf: builtins.str | None = ...,
         state_schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ...,
+        transform_with_state_info: global___TransformWithStateInfo | None = ...,
     ) -> None: ...
     def HasField(
         self,
@@ -3689,6 +3696,8 @@ class GroupMap(google.protobuf.message.Message):
             b"_state_schema",
             "_timeout_conf",
             b"_timeout_conf",
+            "_transform_with_state_info",
+            b"_transform_with_state_info",
             "func",
             b"func",
             "initial_input",
@@ -3703,6 +3712,8 @@ class GroupMap(google.protobuf.message.Message):
             b"state_schema",
             "timeout_conf",
             b"timeout_conf",
+            "transform_with_state_info",
+            b"transform_with_state_info",
         ],
     ) -> builtins.bool: ...
     def ClearField(
@@ -3716,6 +3727,8 @@ class GroupMap(google.protobuf.message.Message):
             b"_state_schema",
             "_timeout_conf",
             b"_timeout_conf",
+            "_transform_with_state_info",
+            b"_transform_with_state_info",
             "func",
             b"func",
             "grouping_expressions",
@@ -3736,6 +3749,8 @@ class GroupMap(google.protobuf.message.Message):
             b"state_schema",
             "timeout_conf",
             b"timeout_conf",
+            "transform_with_state_info",
+            b"transform_with_state_info",
         ],
     ) -> None: ...
     @typing.overload
@@ -3757,9 +3772,82 @@ class GroupMap(google.protobuf.message.Message):
     def WhichOneof(
         self, oneof_group: typing_extensions.Literal["_timeout_conf", b"_timeout_conf"]
     ) -> typing_extensions.Literal["timeout_conf"] | None: ...
+    @typing.overload
+    def WhichOneof(
+        self,
+        oneof_group: typing_extensions.Literal[
+            "_transform_with_state_info", b"_transform_with_state_info"
+        ],
+    ) -> typing_extensions.Literal["transform_with_state_info"] | None: ...
 
 global___GroupMap = GroupMap
 
+class TransformWithStateInfo(google.protobuf.message.Message):
+    """Additional input parameters used for TransformWithState operator."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    TIME_MODE_FIELD_NUMBER: builtins.int
+    EVENT_TIME_COLUMN_NAME_FIELD_NUMBER: builtins.int
+    OUTPUT_SCHEMA_FIELD_NUMBER: builtins.int
+    time_mode: builtins.str
+    """(Required) Time mode string for transformWithState."""
+    event_time_column_name: builtins.str
+    """(Optional) Event time column name."""
+    @property
+    def output_schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
+        """(Optional) Schema for the output DataFrame.
+        Only required used for TransformWithStateInPandas.
+        """
+    def __init__(
+        self,
+        *,
+        time_mode: builtins.str = ...,
+        event_time_column_name: builtins.str | None = ...,
+        output_schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ...,
+    ) -> None: ...
+    def HasField(
+        self,
+        field_name: typing_extensions.Literal[
+            "_event_time_column_name",
+            b"_event_time_column_name",
+            "_output_schema",
+            b"_output_schema",
+            "event_time_column_name",
+            b"event_time_column_name",
+            "output_schema",
+            b"output_schema",
+        ],
+    ) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing_extensions.Literal[
+            "_event_time_column_name",
+            b"_event_time_column_name",
+            "_output_schema",
+            b"_output_schema",
+            "event_time_column_name",
+            b"event_time_column_name",
+            "output_schema",
+            b"output_schema",
+            "time_mode",
+            b"time_mode",
+        ],
+    ) -> None: ...
+    @typing.overload
+    def WhichOneof(
+        self,
+        oneof_group: typing_extensions.Literal[
+            "_event_time_column_name", b"_event_time_column_name"
+        ],
+    ) -> typing_extensions.Literal["event_time_column_name"] | None: ...
+    @typing.overload
+    def WhichOneof(
+        self, oneof_group: typing_extensions.Literal["_output_schema", b"_output_schema"]
+    ) -> typing_extensions.Literal["output_schema"] | None: ...
+
+global___TransformWithStateInfo = TransformWithStateInfo
+
 class CoGroupMap(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor