You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ka...@apache.org on 2025/04/08 06:10:32 UTC
(spark) branch HeartSaVioR/SPARK-XXXXX-python-transform-with-state created (now 5a709080dd01)
This is an automated email from the ASF dual-hosted git repository.
kabhwan pushed a change to branch HeartSaVioR/SPARK-XXXXX-python-transform-with-state
in repository https://gitbox.apache.org/repos/asf/spark.git
at 5a709080dd01 fix conflict
This branch includes the following new commits:
new fc6a2d5a3a5b WIP SPARK-XXXXX still working on python transformWithState
new 5a709080dd01 fix conflict
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org
(spark) 01/02: WIP SPARK-XXXXX still working on python transformWithState
Posted by ka...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
kabhwan pushed a commit to branch HeartSaVioR/SPARK-XXXXX-python-transform-with-state
in repository https://gitbox.apache.org/repos/asf/spark.git
commit fc6a2d5a3a5b1e5e6e9e648587cc39867b5c20b4
Author: Jungtaek Lim <ka...@gmail.com>
AuthorDate: Tue Apr 8 14:15:42 2025 +0900
WIP SPARK-XXXXX still working on python transformWithState
---
.../org/apache/spark/api/python/PythonRunner.scala | 2 +
python/pyspark/sql/connect/group.py | 4 +-
python/pyspark/sql/pandas/_typing/__init__.pyi | 2 +
python/pyspark/sql/pandas/functions.py | 2 +
python/pyspark/sql/pandas/group_ops.py | 66 +-
python/pyspark/sql/pandas/serializers.py | 111 +-
.../sql/streaming/stateful_processor_util.py | 58 +-
.../transform_with_state_driver_worker.py | 6 +-
.../sql/tests/pandas/test_transform_with_state.py | 2054 ++++++++++++++++++++
python/pyspark/util.py | 7 +
python/pyspark/worker.py | 134 +-
.../analysis/UnsupportedOperationChecker.scala | 2 +-
.../plans/logical/pythonLogicalOperators.scala | 13 +-
.../org/apache/spark/sql/internal/SQLConf.scala | 17 +-
.../sql/classic/RelationalGroupedDataset.scala | 53 +-
.../spark/sql/execution/SparkStrategies.scala | 24 +-
... TransformWithStateInPySparkDeserializer.scala} | 6 +-
...scala => TransformWithStateInPySparkExec.scala} | 58 +-
... TransformWithStateInPySparkPythonRunner.scala} | 42 +-
...> TransformWithStateInPySparkStateServer.scala} | 20 +-
.../execution/streaming/IncrementalExecution.scala | 18 +-
...nsformWithStateInPySparkStateServerSuite.scala} | 46 +-
22 files changed, 2585 insertions(+), 160 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 84701ee593c1..1876618bc4d2 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -64,6 +64,8 @@ private[spark] object PythonEvalType {
val SQL_COGROUPED_MAP_ARROW_UDF = 210
val SQL_TRANSFORM_WITH_STATE_PANDAS_UDF = 211
val SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF = 212
+ val SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF = 213
+ val SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF = 214
val SQL_TABLE_UDF = 300
val SQL_ARROW_TABLE_UDF = 301
diff --git a/python/pyspark/sql/connect/group.py b/python/pyspark/sql/connect/group.py
index 79cf930fd8ee..a0efa9d4dc42 100644
--- a/python/pyspark/sql/connect/group.py
+++ b/python/pyspark/sql/connect/group.py
@@ -373,9 +373,9 @@ class GroupedData:
) -> "DataFrame":
from pyspark.sql.connect.udf import UserDefinedFunction
from pyspark.sql.connect.dataframe import DataFrame
- from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasUdfUtils
+ from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkUdfUtils
- udf_util = TransformWithStateInPandasUdfUtils(statefulProcessor, timeMode)
+ udf_util = TransformWithStateInPySparkUdfUtils(statefulProcessor, timeMode)
if initialState is None:
udf_obj = UserDefinedFunction(
udf_util.transformWithStateUDF,
diff --git a/python/pyspark/sql/pandas/_typing/__init__.pyi b/python/pyspark/sql/pandas/_typing/__init__.pyi
index 750e3f52709c..193d368f6ebd 100644
--- a/python/pyspark/sql/pandas/_typing/__init__.pyi
+++ b/python/pyspark/sql/pandas/_typing/__init__.pyi
@@ -57,6 +57,8 @@ ArrowGroupedMapUDFType = Literal[209]
ArrowCogroupedMapUDFType = Literal[210]
PandasGroupedMapUDFTransformWithStateType = Literal[211]
PandasGroupedMapUDFTransformWithStateInitStateType = Literal[212]
+GroupedMapUDFTransformWithStateType = Literal[213]
+GroupedMapUDFTransformWithStateInitStateType = Literal[214]
class PandasVariadicScalarToScalarFunction(Protocol):
def __call__(self, *_: DataFrameOrSeriesLike_) -> DataFrameOrSeriesLike_: ...
diff --git a/python/pyspark/sql/pandas/functions.py b/python/pyspark/sql/pandas/functions.py
index c9433285982e..810cd7546ba0 100644
--- a/python/pyspark/sql/pandas/functions.py
+++ b/python/pyspark/sql/pandas/functions.py
@@ -415,6 +415,8 @@ def pandas_udf(f=None, returnType=None, functionType=None):
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF,
PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF,
None,
diff --git a/python/pyspark/sql/pandas/group_ops.py b/python/pyspark/sql/pandas/group_ops.py
index cd384000f859..80b1c3e9c377 100644
--- a/python/pyspark/sql/pandas/group_ops.py
+++ b/python/pyspark/sql/pandas/group_ops.py
@@ -484,7 +484,7 @@ class PandasGroupedOpsMixin:
from pyspark.sql import GroupedData
from pyspark.sql.functions import pandas_udf
- from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasUdfUtils
+ from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkUdfUtils
assert isinstance(self, GroupedData)
if initialState is not None:
@@ -493,7 +493,7 @@ class PandasGroupedOpsMixin:
outputStructType = cast(StructType, self._df._session._parse_ddl(outputStructType))
df = self._df
- udf_util = TransformWithStateInPandasUdfUtils(statefulProcessor, timeMode)
+ udf_util = TransformWithStateInPySparkUdfUtils(statefulProcessor, timeMode)
if initialState is None:
initial_state_java_obj = None
@@ -522,6 +522,68 @@ class PandasGroupedOpsMixin:
)
return DataFrame(jdf, self.session)
+ # FIXME: documentation
+ def transformWithState(
+ self,
+ statefulProcessor: StatefulProcessor,
+ outputStructType: Union[StructType, str],
+ outputMode: str,
+ timeMode: str,
+ initialState: Optional["GroupedData"] = None,
+ eventTimeColumnName: str = "",
+ ) -> "DataFrame":
+ from pyspark.sql import GroupedData
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkUdfUtils
+
+ assert isinstance(self, GroupedData)
+ if initialState is not None:
+ assert isinstance(initialState, GroupedData)
+ if isinstance(outputStructType, str):
+ outputStructType = cast(StructType, self._df._session._parse_ddl(outputStructType))
+
+ df = self._df
+ udf_util = TransformWithStateInPySparkUdfUtils(statefulProcessor, timeMode)
+
+ if initialState is None:
+ initial_state_java_obj = None
+ udf = pandas_udf(
+ udf_util.transformWithStateUDF, # type: ignore
+ returnType=outputStructType,
+ functionType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF,
+ )
+ else:
+ initial_state_java_obj = initialState._jgd
+ udf = pandas_udf(
+ udf_util.transformWithStateWithInitStateUDF, # type: ignore
+ returnType=outputStructType,
+ functionType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF,
+ )
+
+ udf_column = udf(*[df[col] for col in df.columns])
+
+ jdf = self._jgd.transformWithState(
+ udf_column._jc,
+ self.session._jsparkSession.parseDataType(outputStructType.json()),
+ outputMode,
+ timeMode,
+ initial_state_java_obj,
+ eventTimeColumnName,
+ )
+ return DataFrame(jdf, self.session)
+
+ # FIXME: refactor!
+ # def __transformWithState(
+ # self,
+ # statefulProcessor: StatefulProcessor,
+ # outputStructType: Union[StructType, str],
+ # outputMode: str,
+ # timeMode: str,
+ # initialState: Optional["GroupedData"] = None,
+ # eventTimeColumnName: str = "",
+ #
+ # ) -> "DataFrame":
+
def applyInArrow(
self, func: "ArrowGroupedMapFunction", schema: Union[StructType, str]
) -> "DataFrame":
diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py
index de9020f633c5..de58dc09354a 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -1233,7 +1233,7 @@ class TransformWithStateInPandasSerializer(ArrowStreamPandasUDFSerializer):
this function works in overall.
"""
import pyarrow as pa
- from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+ from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkFuncMode
def generate_data_batches(batches):
"""
@@ -1258,11 +1258,11 @@ class TransformWithStateInPandasSerializer(ArrowStreamPandasUDFSerializer):
data_batches = generate_data_batches(_batches)
for k, g in groupby(data_batches, key=lambda x: x[0]):
- yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g)
+ yield (TransformWithStateInPySparkFuncMode.PROCESS_DATA, k, g)
- yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None)
+ yield (TransformWithStateInPySparkFuncMode.PROCESS_TIMER, None, None)
- yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None)
+ yield (TransformWithStateInPySparkFuncMode.COMPLETE, None, None)
def dump_stream(self, iterator, stream):
"""
@@ -1299,7 +1299,7 @@ class TransformWithStateInPandasInitStateSerializer(TransformWithStateInPandasSe
def load_stream(self, stream):
import pyarrow as pa
- from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+ from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkFuncMode
def generate_data_batches(batches):
"""
@@ -1360,8 +1360,103 @@ class TransformWithStateInPandasInitStateSerializer(TransformWithStateInPandasSe
data_batches = generate_data_batches(_batches)
for k, g in groupby(data_batches, key=lambda x: x[0]):
- yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g)
+ yield (TransformWithStateInPySparkFuncMode.PROCESS_DATA, k, g)
- yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None)
+ yield (TransformWithStateInPySparkFuncMode.PROCESS_TIMER, None, None)
- yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None)
+ yield (TransformWithStateInPySparkFuncMode.COMPLETE, None, None)
+
+
+class TransformWithStateInPySparkRowSerializer(ArrowStreamUDFSerializer):
+ """
+ Serializer used by Python worker to evaluate UDF for
+ :meth:`pyspark.sql.GroupedData.transformWithState`.
+
+ Parameters
+ ----------
+ timezone : str
+ A timezone to respect when handling timestamp values
+ safecheck : bool
+ If True, conversion from Arrow to Pandas checks for overflow/truncation
+ assign_cols_by_name : bool
+ If True, then Pandas DataFrames will get columns by name
+ arrow_max_records_per_batch : int
+ Limit of the number of records that can be written to a single ArrowRecordBatch in memory.
+ """
+
+ def __init__(self, timezone, safecheck, assign_cols_by_name, arrow_max_records_per_batch):
+ # FIXME: which constructor params are needed?
+ super(TransformWithStateInPySparkRowSerializer, self).__init__(
+ # timezone, safecheck, assign_cols_by_name
+ )
+ self.arrow_max_records_per_batch = arrow_max_records_per_batch
+ self.key_offsets = None
+
+ def load_stream(self, stream):
+ """
+ Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunk, and
+ convert the data into a list of pandas.Series.
+
+ Please refer the doc of inner function `generate_data_batches` for more details how
+ this function works in overall.
+ """
+ import pyarrow as pa
+ from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkFuncMode
+
+ def generate_data_batches(batches):
+ # FIXME: documentation
+ """
+ Deserialize ArrowRecordBatches and return a generator of pandas.Series list.
+
+ The deserialization logic assumes that Arrow RecordBatches contain the data with the
+ ordering that data chunks for same grouping key will appear sequentially.
+
+ This function must avoid materializing multiple Arrow RecordBatches into memory at the
+ same time. And data chunks from the same grouping key should appear sequentially.
+ """
+ for batch in batches:
+ DataRow = Row(*(batch.column_names))
+
+ # This is supposed to be the same.
+ batch_key = tuple(batch[o][0].as_py() for o in self.key_offsets)
+ rows = []
+ for row_idx in range(batch.num_rows):
+ row = DataRow(
+ *(batch.column(i)[row_idx].as_py() for i in range(batch.num_columns))
+ )
+ rows.append(row)
+
+ yield (batch_key, rows)
+
+ _batches = super(ArrowStreamUDFSerializer, self).load_stream(stream)
+ data_batches = generate_data_batches(_batches)
+
+ for k, g in groupby(data_batches, key=lambda x: x[0]):
+ yield (TransformWithStateInPySparkFuncMode.PROCESS_DATA, k, g)
+
+ yield (TransformWithStateInPySparkFuncMode.PROCESS_TIMER, None, None)
+
+ yield (TransformWithStateInPySparkFuncMode.COMPLETE, None, None)
+
+ def dump_stream(self, iterator, stream):
+ """
+ Read through an iterator of (iterator of Row), serialize them to Arrow
+ RecordBatches, and write batches to stream.
+ """
+ import pyarrow as pa
+
+ def flatten_iterator():
+ # iterator: iter[list[(iter[Row], pdf_type)]]
+ for packed in iterator:
+ iter_row_with_type = packed[0]
+ iter_row = iter_row_with_type[0]
+ pdf_type = iter_row_with_type[1]
+ rows_as_dict = []
+ for row in iter_row:
+ row_as_dict = row.asDict(True)
+ # FIXME: Should this respect the column name based association?
+ rows_as_dict.append(row_as_dict)
+ record_batch = pa.RecordBatch.from_pylist(rows_as_dict, schema=pdf_type)
+ yield (record_batch, pdf_type)
+
+ return ArrowStreamUDFSerializer.dump_stream(self, flatten_iterator(), stream)
diff --git a/python/pyspark/sql/streaming/stateful_processor_util.py b/python/pyspark/sql/streaming/stateful_processor_util.py
index fbc3093f8709..1d63756f4f8f 100644
--- a/python/pyspark/sql/streaming/stateful_processor_util.py
+++ b/python/pyspark/sql/streaming/stateful_processor_util.py
@@ -17,7 +17,7 @@
from enum import Enum
import itertools
-from typing import Any, Iterator, Optional, TYPE_CHECKING
+from typing import Any, Iterator, Optional, TYPE_CHECKING, Union
from pyspark.sql.streaming.stateful_processor_api_client import (
StatefulProcessorApiClient,
StatefulProcessorHandleState,
@@ -28,18 +28,20 @@ from pyspark.sql.streaming.stateful_processor import (
StatefulProcessorHandle,
TimerValues,
)
+from pyspark.sql.types import Row
if TYPE_CHECKING:
from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike
-# This file places the utilities for transformWithStateInPandas; we have a separate file to avoid
-# putting internal classes to the stateful_processor.py file which contains public APIs.
+# This file places the utilities for transformWithState in PySpark (Row, and Pandas); we have
+# a separate file to avoid putting internal classes to the stateful_processor.py file which
+# contains public APIs.
-class TransformWithStateInPandasFuncMode(Enum):
+class TransformWithStateInPySparkFuncMode(Enum):
"""
- Internal mode for python worker UDF mode for transformWithStateInPandas; external mode are in
- `StatefulProcessorHandleState` for public use purposes.
+ Internal mode for python worker UDF mode for transformWithState in PySpark; external mode are
+ in `StatefulProcessorHandleState` for public use purposes.
"""
PROCESS_DATA = 1
@@ -48,10 +50,10 @@ class TransformWithStateInPandasFuncMode(Enum):
PRE_INIT = 4
-class TransformWithStateInPandasUdfUtils:
+class TransformWithStateInPySparkUdfUtils:
"""
- Internal Utility class used for python worker UDF for transformWithStateInPandas. This class is
- shared for both classic and spark connect mode.
+ Internal Utility class used for python worker UDF for transformWithState in PySpark. This class
+ is shared for both classic and spark connect mode.
"""
def __init__(self, stateful_processor: StatefulProcessor, time_mode: str):
@@ -61,11 +63,11 @@ class TransformWithStateInPandasUdfUtils:
def transformWithStateUDF(
self,
stateful_processor_api_client: StatefulProcessorApiClient,
- mode: TransformWithStateInPandasFuncMode,
+ mode: TransformWithStateInPySparkFuncMode,
key: Any,
- input_rows: Iterator["PandasDataFrameLike"],
- ) -> Iterator["PandasDataFrameLike"]:
- if mode == TransformWithStateInPandasFuncMode.PRE_INIT:
+ input_rows: Union[Iterator["PandasDataFrameLike"], Iterator[Row]],
+ ) -> Union[Iterator["PandasDataFrameLike"], Iterator[Row]]:
+ if mode == TransformWithStateInPySparkFuncMode.PRE_INIT:
return self._handle_pre_init(stateful_processor_api_client)
handle = StatefulProcessorHandle(stateful_processor_api_client)
@@ -74,13 +76,13 @@ class TransformWithStateInPandasUdfUtils:
self._stateful_processor.init(handle)
stateful_processor_api_client.set_handle_state(StatefulProcessorHandleState.INITIALIZED)
- if mode == TransformWithStateInPandasFuncMode.PROCESS_TIMER:
+ if mode == TransformWithStateInPySparkFuncMode.PROCESS_TIMER:
stateful_processor_api_client.set_handle_state(
StatefulProcessorHandleState.DATA_PROCESSED
)
result = self._handle_expired_timers(stateful_processor_api_client)
return result
- elif mode == TransformWithStateInPandasFuncMode.COMPLETE:
+ elif mode == TransformWithStateInPySparkFuncMode.COMPLETE:
stateful_processor_api_client.set_handle_state(
StatefulProcessorHandleState.TIMER_PROCESSED
)
@@ -89,18 +91,18 @@ class TransformWithStateInPandasUdfUtils:
stateful_processor_api_client.set_handle_state(StatefulProcessorHandleState.CLOSED)
return iter([])
else:
- # mode == TransformWithStateInPandasFuncMode.PROCESS_DATA
+ # mode == TransformWithStateInPySparkFuncMode.PROCESS_DATA
result = self._handle_data_rows(stateful_processor_api_client, key, input_rows)
return result
def transformWithStateWithInitStateUDF(
self,
stateful_processor_api_client: StatefulProcessorApiClient,
- mode: TransformWithStateInPandasFuncMode,
+ mode: TransformWithStateInPySparkFuncMode,
key: Any,
- input_rows: Iterator["PandasDataFrameLike"],
- initial_states: Optional[Iterator["PandasDataFrameLike"]] = None,
- ) -> Iterator["PandasDataFrameLike"]:
+ input_rows: Union[Iterator["PandasDataFrameLike"], Iterator[Row]],
+ initial_states: Optional[Union[Iterator["PandasDataFrameLike"], Iterator[Row]]] = None,
+ ) -> Union[Iterator["PandasDataFrameLike"], Iterator[Row]]:
"""
UDF for TWS operator with non-empty initial states. Possible input combinations
of inputRows and initialStates iterator:
@@ -113,7 +115,7 @@ class TransformWithStateInPandasUdfUtils:
- `initialStates` is None, while `inputRows` is not empty. This is not first batch.
`initialStates` is initialized to the positional value as None.
"""
- if mode == TransformWithStateInPandasFuncMode.PRE_INIT:
+ if mode == TransformWithStateInPySparkFuncMode.PRE_INIT:
return self._handle_pre_init(stateful_processor_api_client)
handle = StatefulProcessorHandle(stateful_processor_api_client)
@@ -122,19 +124,19 @@ class TransformWithStateInPandasUdfUtils:
self._stateful_processor.init(handle)
stateful_processor_api_client.set_handle_state(StatefulProcessorHandleState.INITIALIZED)
- if mode == TransformWithStateInPandasFuncMode.PROCESS_TIMER:
+ if mode == TransformWithStateInPySparkFuncMode.PROCESS_TIMER:
stateful_processor_api_client.set_handle_state(
StatefulProcessorHandleState.DATA_PROCESSED
)
result = self._handle_expired_timers(stateful_processor_api_client)
return result
- elif mode == TransformWithStateInPandasFuncMode.COMPLETE:
+ elif mode == TransformWithStateInPySparkFuncMode.COMPLETE:
stateful_processor_api_client.remove_implicit_key()
self._stateful_processor.close()
stateful_processor_api_client.set_handle_state(StatefulProcessorHandleState.CLOSED)
return iter([])
else:
- # mode == TransformWithStateInPandasFuncMode.PROCESS_DATA
+ # mode == TransformWithStateInPySparkFuncMode.PROCESS_DATA
batch_timestamp, watermark_timestamp = stateful_processor_api_client.get_timestamps(
self._time_mode
)
@@ -166,7 +168,7 @@ class TransformWithStateInPandasUdfUtils:
def _handle_pre_init(
self, stateful_processor_api_client: StatefulProcessorApiClient
- ) -> Iterator["PandasDataFrameLike"]:
+ ) -> Union[Iterator["PandasDataFrameLike"], Iterator[Row]]:
# Driver handle is different from the handle used on executors;
# On JVM side, we will use `DriverStatefulProcessorHandleImpl` for driver handle which
# will only be used for handling init() and get the state schema on the driver.
@@ -186,8 +188,8 @@ class TransformWithStateInPandasUdfUtils:
self,
stateful_processor_api_client: StatefulProcessorApiClient,
key: Any,
- input_rows: Optional[Iterator["PandasDataFrameLike"]] = None,
- ) -> Iterator["PandasDataFrameLike"]:
+ input_rows: Optional[Union[Iterator["PandasDataFrameLike"], Iterator[Row]]] = None,
+ ) -> Union[Iterator["PandasDataFrameLike"], Iterator[Row]]:
stateful_processor_api_client.set_implicit_key(key)
batch_timestamp, watermark_timestamp = stateful_processor_api_client.get_timestamps(
@@ -206,7 +208,7 @@ class TransformWithStateInPandasUdfUtils:
def _handle_expired_timers(
self,
stateful_processor_api_client: StatefulProcessorApiClient,
- ) -> Iterator["PandasDataFrameLike"]:
+ ) -> Union[Iterator["PandasDataFrameLike"], Iterator[Row]]:
batch_timestamp, watermark_timestamp = stateful_processor_api_client.get_timestamps(
self._time_mode
)
diff --git a/python/pyspark/sql/streaming/transform_with_state_driver_worker.py b/python/pyspark/sql/streaming/transform_with_state_driver_worker.py
index 99d386f07b5b..680e3057aa94 100644
--- a/python/pyspark/sql/streaming/transform_with_state_driver_worker.py
+++ b/python/pyspark/sql/streaming/transform_with_state_driver_worker.py
@@ -31,7 +31,7 @@ from pyspark.util import handle_worker_exception
from typing import IO
from pyspark.worker_util import check_python_version
from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
-from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkFuncMode
from pyspark.sql.types import StructType
if TYPE_CHECKING:
@@ -51,7 +51,7 @@ def main(infile: IO, outfile: IO) -> None:
def process(
processor: StatefulProcessorApiClient,
- mode: TransformWithStateInPandasFuncMode,
+ mode: TransformWithStateInPySparkFuncMode,
key: Any,
input: Iterator["PandasDataFrameLike"],
) -> None:
@@ -81,7 +81,7 @@ def main(infile: IO, outfile: IO) -> None:
stateful_processor_api_client = StatefulProcessorApiClient(state_server_port, key_schema)
process(
stateful_processor_api_client,
- TransformWithStateInPandasFuncMode.PRE_INIT,
+ TransformWithStateInPySparkFuncMode.PRE_INIT,
None,
iter([]),
)
diff --git a/python/pyspark/sql/tests/pandas/test_transform_with_state.py b/python/pyspark/sql/tests/pandas/test_transform_with_state.py
new file mode 100644
index 000000000000..660a8d4c057d
--- /dev/null
+++ b/python/pyspark/sql/tests/pandas/test_transform_with_state.py
@@ -0,0 +1,2054 @@
+#
+# 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.
+#
+
+import json
+import os
+import time
+import tempfile
+from pyspark.sql.streaming import StatefulProcessor, StatefulProcessorHandle
+from typing import Iterator
+
+import unittest
+from typing import cast
+
+from pyspark import SparkConf
+from pyspark.errors import PySparkRuntimeError
+from pyspark.sql.functions import array_sort, col, explode, split
+from pyspark.sql.types import (
+ StringType,
+ StructType,
+ StructField,
+ Row,
+ IntegerType,
+ TimestampType,
+ LongType,
+ BooleanType,
+ FloatType,
+)
+from pyspark.testing import assertDataFrameEqual
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pyarrow,
+ cast(str, pyarrow_requirement_message),
+)
+class TransformWithStateInPySparkTestsMixin:
+ @classmethod
+ def conf(cls):
+ cfg = SparkConf()
+ cfg.set("spark.sql.shuffle.partitions", "5")
+ cfg.set(
+ "spark.sql.streaming.stateStore.providerClass",
+ "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider",
+ )
+ cfg.set("spark.sql.execution.arrow.transformWithStateInPySpark.maxRecordsPerBatch", "2")
+ cfg.set("spark.sql.session.timeZone", "UTC")
+ # TODO SPARK-49046 this config is to stop query from FEB sink gracefully
+ cfg.set("spark.sql.streaming.noDataMicroBatches.enabled", "false")
+ return cfg
+
+ def _prepare_input_data(self, input_path, col1, col2):
+ with open(input_path, "w") as fw:
+ for e1, e2 in zip(col1, col2):
+ fw.write(f"{e1}, {e2}\n")
+
+ def _prepare_test_resource1(self, input_path):
+ self._prepare_input_data(input_path + "/text-test1.txt", [0, 0, 1, 1], [123, 46, 146, 346])
+
+ def _prepare_test_resource2(self, input_path):
+ self._prepare_input_data(
+ input_path + "/text-test2.txt", [0, 0, 0, 1, 1], [123, 223, 323, 246, 6]
+ )
+
+ def _prepare_test_resource3(self, input_path):
+ self._prepare_input_data(input_path + "/text-test3.txt", [0, 1], [123, 6])
+
+ def _prepare_test_resource4(self, input_path):
+ self._prepare_input_data(input_path + "/text-test4.txt", [0, 1], [123, 6])
+
+ def _prepare_test_resource5(self, input_path):
+ self._prepare_input_data(input_path + "/text-test5.txt", [0, 1], [123, 6])
+
+ def _build_test_df(self, input_path):
+ df = self.spark.readStream.format("text").option("maxFilesPerTrigger", 1).load(input_path)
+ df_split = df.withColumn("split_values", split(df["value"], ","))
+ df_final = df_split.select(
+ df_split.split_values.getItem(0).alias("id").cast("string"),
+ df_split.split_values.getItem(1).alias("temperature").cast("int"),
+ )
+ return df_final
+
+ def _prepare_input_data_with_3_cols(self, input_path, col1, col2, col3):
+ with open(input_path, "w") as fw:
+ for e1, e2, e3 in zip(col1, col2, col3):
+ fw.write(f"{e1},{e2},{e3}\n")
+
+ def build_test_df_with_3_cols(self, input_path):
+ df = self.spark.readStream.format("text").option("maxFilesPerTrigger", 1).load(input_path)
+ df_split = df.withColumn("split_values", split(df["value"], ","))
+ df_final = df_split.select(
+ df_split.split_values.getItem(0).alias("id1").cast("string"),
+ df_split.split_values.getItem(1).alias("temperature").cast("int"),
+ df_split.split_values.getItem(2).alias("id2").cast("string"),
+ )
+ return df_final
+
+ def _test_transform_with_state_basic(
+ self,
+ stateful_processor,
+ check_results,
+ single_batch=False,
+ timeMode="None",
+ checkpoint_path=None,
+ initial_state=None,
+ ):
+ input_path = tempfile.mkdtemp()
+ if checkpoint_path is None:
+ checkpoint_path = tempfile.mkdtemp()
+ self._prepare_test_resource1(input_path)
+ if not single_batch:
+ time.sleep(2)
+ self._prepare_test_resource2(input_path)
+
+ df = self._build_test_df(input_path)
+
+ for q in self.spark.streams.active:
+ q.stop()
+ self.assertTrue(df.isStreaming)
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("countAsString", StringType(), True),
+ ]
+ )
+
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=stateful_processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode=timeMode,
+ initialState=initial_state,
+ )
+ .writeStream.queryName("this_query")
+ .option("checkpointLocation", checkpoint_path)
+ .foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+
+ self.assertEqual(q.name, "this_query")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+
+ def test_transform_with_state_basic(self):
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ self._test_transform_with_state_basic(SimpleStatefulProcessor(), check_results)
+
+ def test_transform_with_state_non_exist_value_state(self):
+ def check_results(batch_df, _):
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="0"),
+ Row(id="1", countAsString="0"),
+ }
+
+ self._test_transform_with_state_basic(
+ InvalidSimpleStatefulProcessor(), check_results, True
+ )
+
+ def test_transform_with_state_query_restarts(self):
+ root_path = tempfile.mkdtemp()
+ input_path = root_path + "/input"
+ os.makedirs(input_path, exist_ok=True)
+ checkpoint_path = root_path + "/checkpoint"
+ output_path = root_path + "/output"
+
+ self._prepare_test_resource1(input_path)
+
+ df = self._build_test_df(input_path)
+
+ for q in self.spark.streams.active:
+ q.stop()
+ self.assertTrue(df.isStreaming)
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("countAsString", StringType(), True),
+ ]
+ )
+
+ base_query = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=SimpleStatefulProcessor(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+ .writeStream.queryName("this_query")
+ .format("parquet")
+ .outputMode("append")
+ .option("checkpointLocation", checkpoint_path)
+ .option("path", output_path)
+ )
+ q = base_query.start()
+ self.assertEqual(q.name, "this_query")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+
+ # Verify custom metrics.
+ self.assertTrue(q.lastProgress.stateOperators[0].customMetrics["numValueStateVars"] > 0)
+ self.assertTrue(q.lastProgress.stateOperators[0].customMetrics["numDeletedStateVars"] > 0)
+
+ q.stop()
+
+ self._prepare_test_resource2(input_path)
+
+ q = base_query.start()
+ self.assertEqual(q.name, "this_query")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+ result_df = self.spark.read.parquet(output_path)
+ assert set(result_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ def test_transform_with_state_list_state(self):
+ def check_results(batch_df, _):
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ self._test_transform_with_state_basic(
+ ListStateProcessor(), check_results, True, "processingTime"
+ )
+
+ # test list state with ttl has the same behavior as list state when state doesn't expire.
+ def test_transform_with_state_list_state_large_ttl(self):
+ def check_results(batch_df, batch_id):
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ self._test_transform_with_state_basic(
+ ListStateLargeTTLProcessor(), check_results, True, "processingTime"
+ )
+
+ def test_transform_with_state_map_state(self):
+ def check_results(batch_df, _):
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ self._test_transform_with_state_basic(MapStateProcessor(), check_results, True)
+
+ # test map state with ttl has the same behavior as map state when state doesn't expire.
+ def test_transform_with_state_map_state_large_ttl(self):
+ def check_results(batch_df, batch_id):
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ self._test_transform_with_state_basic(
+ MapStateLargeTTLProcessor(), check_results, True, "processingTime"
+ )
+
+ # test value state with ttl has the same behavior as value state when
+ # state doesn't expire.
+ def test_value_state_ttl_basic(self):
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ self._test_transform_with_state_basic(
+ SimpleTTLStatefulProcessor(), check_results, False, "processingTime"
+ )
+
+ # TODO SPARK-50908 holistic fix for TTL suite
+ @unittest.skip("test is flaky and it is only a timing issue, skipping until we can resolve")
+ def test_value_state_ttl_expiration(self):
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assertDataFrameEqual(
+ batch_df,
+ [
+ Row(id="ttl-count-0", count=1),
+ Row(id="count-0", count=1),
+ Row(id="ttl-list-state-count-0", count=1),
+ Row(id="ttl-map-state-count-0", count=1),
+ Row(id="ttl-count-1", count=1),
+ Row(id="count-1", count=1),
+ Row(id="ttl-list-state-count-1", count=1),
+ Row(id="ttl-map-state-count-1", count=1),
+ ],
+ )
+ elif batch_id == 1:
+ assertDataFrameEqual(
+ batch_df,
+ [
+ Row(id="ttl-count-0", count=2),
+ Row(id="count-0", count=2),
+ Row(id="ttl-list-state-count-0", count=3),
+ Row(id="ttl-map-state-count-0", count=2),
+ Row(id="ttl-count-1", count=2),
+ Row(id="count-1", count=2),
+ Row(id="ttl-list-state-count-1", count=3),
+ Row(id="ttl-map-state-count-1", count=2),
+ ],
+ )
+ else:
+ # ttl-count-0 expire and restart from count 0.
+ # The TTL for value state ttl_count_state gets reset in batch 1 because of the
+ # update operation and ttl-count-1 keeps the state.
+ # ttl-list-state-count-0 expire and restart from count 0.
+ # The TTL for list state ttl_list_state gets reset in batch 1 because of the
+ # put operation and ttl-list-state-count-1 keeps the state.
+ # non-ttl state never expires
+ assertDataFrameEqual(
+ batch_df,
+ [
+ Row(id="ttl-count-0", count=1),
+ Row(id="count-0", count=3),
+ Row(id="ttl-list-state-count-0", count=1),
+ Row(id="ttl-map-state-count-0", count=1),
+ Row(id="ttl-count-1", count=3),
+ Row(id="count-1", count=3),
+ Row(id="ttl-list-state-count-1", count=7),
+ Row(id="ttl-map-state-count-1", count=3),
+ ],
+ )
+
+ if batch_id == 0 or batch_id == 1:
+ time.sleep(4)
+
+ input_dir = tempfile.TemporaryDirectory()
+ input_path = input_dir.name
+ try:
+ df = self._build_test_df(input_path)
+ self._prepare_input_data(input_path + "/batch1.txt", [1, 0], [0, 0])
+ self._prepare_input_data(input_path + "/batch2.txt", [1, 0], [0, 0])
+ self._prepare_input_data(input_path + "/batch3.txt", [1, 0], [0, 0])
+ for q in self.spark.streams.active:
+ q.stop()
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("count", IntegerType(), True),
+ ]
+ )
+
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=TTLStatefulProcessor(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="processingTime",
+ )
+ .writeStream.foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.stop()
+ q.awaitTermination()
+ self.assertTrue(q.exception() is None)
+ finally:
+ input_dir.cleanup()
+
+ def _test_transform_with_state_proc_timer(self, stateful_processor, check_results):
+ input_path = tempfile.mkdtemp()
+ self._prepare_test_resource3(input_path)
+ time.sleep(2)
+ self._prepare_test_resource1(input_path)
+ time.sleep(2)
+ self._prepare_test_resource2(input_path)
+
+ df = self._build_test_df(input_path)
+
+ for q in self.spark.streams.active:
+ q.stop()
+ self.assertTrue(df.isStreaming)
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("countAsString", StringType(), True),
+ StructField("timeValues", StringType(), True),
+ ]
+ )
+
+ query_name = "processing_time_test_query"
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=stateful_processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="processingtime",
+ )
+ .writeStream.queryName(query_name)
+ .foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+
+ self.assertEqual(q.name, query_name)
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+
+ def test_transform_with_state_proc_timer(self):
+ def check_results(batch_df, batch_id):
+ # helper function to check expired timestamp is smaller than current processing time
+ def check_timestamp(batch_df):
+ expired_df = (
+ batch_df.filter(batch_df["countAsString"] == "-1")
+ .select("id", "timeValues")
+ .withColumnRenamed("timeValues", "expiredTimestamp")
+ )
+ count_df = (
+ batch_df.filter(batch_df["countAsString"] != "-1")
+ .select("id", "timeValues")
+ .withColumnRenamed("timeValues", "countStateTimestamp")
+ )
+ joined_df = expired_df.join(count_df, on="id")
+ for row in joined_df.collect():
+ assert row["expiredTimestamp"] < row["countStateTimestamp"]
+
+ if batch_id == 0:
+ assert set(batch_df.sort("id").select("id", "countAsString").collect()) == {
+ Row(id="0", countAsString="1"),
+ Row(id="1", countAsString="1"),
+ }
+ elif batch_id == 1:
+ # for key 0, the accumulated count is emitted before the count state is cleared
+ # during the timer process
+ assert set(batch_df.sort("id").select("id", "countAsString").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="0", countAsString="-1"),
+ Row(id="1", countAsString="3"),
+ }
+ check_timestamp(batch_df)
+
+ else:
+ assert set(batch_df.sort("id").select("id", "countAsString").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="0", countAsString="-1"),
+ Row(id="1", countAsString="5"),
+ }
+
+ self._test_transform_with_state_proc_timer(
+ ProcTimeStatefulProcessor(), check_results
+ )
+
+ def _test_transform_with_state_event_time(self, stateful_processor, check_results):
+ import pyspark.sql.functions as f
+
+ input_path = tempfile.mkdtemp()
+
+ def prepare_batch1(input_path):
+ with open(input_path + "/text-test3.txt", "w") as fw:
+ fw.write("a, 20\n")
+
+ def prepare_batch2(input_path):
+ with open(input_path + "/text-test1.txt", "w") as fw:
+ fw.write("a, 4\n")
+
+ def prepare_batch3(input_path):
+ with open(input_path + "/text-test2.txt", "w") as fw:
+ fw.write("a, 11\n")
+ fw.write("a, 13\n")
+ fw.write("a, 15\n")
+
+ prepare_batch1(input_path)
+ time.sleep(2)
+ prepare_batch2(input_path)
+ time.sleep(2)
+ prepare_batch3(input_path)
+
+ df = self._build_test_df(input_path)
+ df = df.select(
+ "id", f.from_unixtime(f.col("temperature")).alias("eventTime").cast("timestamp")
+ ).withWatermark("eventTime", "10 seconds")
+
+ for q in self.spark.streams.active:
+ q.stop()
+ self.assertTrue(df.isStreaming)
+
+ output_schema = StructType(
+ [StructField("id", StringType(), True), StructField("timestamp", StringType(), True)]
+ )
+
+ query_name = "event_time_test_query"
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=stateful_processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="eventtime",
+ )
+ .writeStream.queryName(query_name)
+ .foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+
+ self.assertEqual(q.name, query_name)
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+
+ def test_transform_with_state_event_time(self):
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ # watermark for late event = 0
+ # watermark for eviction = 0
+ # timer is registered with expiration time = 0, hence expired at the same batch
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="a", timestamp="20"),
+ Row(id="a-expired", timestamp="0"),
+ }
+ elif batch_id == 1:
+ # watermark for late event = 0
+ # watermark for eviction = 10 (20 - 10)
+ # timer is registered with expiration time = 10, hence expired at the same batch
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="a", timestamp="4"),
+ Row(id="a-expired", timestamp="10000"),
+ }
+ else:
+ # watermark for late event = 10
+ # watermark for eviction = 10 (unchanged as 4 < 10)
+ # timer is registered with expiration time = 10, hence expired at the same batch
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="a", timestamp="15"),
+ Row(id="a-expired", timestamp="10000"),
+ }
+
+ self._test_transform_with_state_event_time(
+ EventTimeStatefulProcessor(), check_results
+ )
+
+ def _test_transform_with_state_init_state(
+ self,
+ stateful_processor,
+ check_results,
+ time_mode="None",
+ checkpoint_path=None,
+ initial_state=None,
+ ):
+ input_path = tempfile.mkdtemp()
+ if checkpoint_path is None:
+ checkpoint_path = tempfile.mkdtemp()
+ self._prepare_test_resource1(input_path)
+ time.sleep(2)
+ self._prepare_input_data(input_path + "/text-test2.txt", [0, 3], [67, 12])
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ df = self._build_test_df(input_path)
+ self.assertTrue(df.isStreaming)
+
+ output_schema = "id string, value string"
+
+ if initial_state is None:
+ data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(data, "id string, initVal int").groupBy("id")
+
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=stateful_processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode=time_mode,
+ initialState=initial_state,
+ )
+ .writeStream.queryName("this_query")
+ .option("checkpointLocation", checkpoint_path)
+ .foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+
+ self.assertEqual(q.name, "this_query")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+
+ def test_transform_with_state_init_state(self):
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ # for key 0, initial state was processed and it was only processed once;
+ # for key 1, it did not appear in the initial state df;
+ # for key 3, it did not appear in the first batch of input keys
+ # so it won't be emitted
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", value=str(789 + 123 + 46)),
+ Row(id="1", value=str(146 + 346)),
+ }
+ else:
+ # for key 0, verify initial state was only processed once in the first batch;
+ # for key 3, verify init state was processed and reflected in the accumulated value
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", value=str(789 + 123 + 46 + 67)),
+ Row(id="3", value=str(987 + 12)),
+ }
+
+ self._test_transform_with_state_init_state(
+ SimpleStatefulProcessorWithInitialState(), check_results
+ )
+
+ def _test_transform_with_state_non_contiguous_grouping_cols(
+ self, stateful_processor, check_results, initial_state=None
+ ):
+ input_path = tempfile.mkdtemp()
+ self._prepare_input_data_with_3_cols(
+ input_path + "/text-test1.txt", [0, 0, 1, 1], [123, 46, 146, 346], [1, 1, 2, 2]
+ )
+
+ df = self.build_test_df_with_3_cols(input_path)
+
+ for q in self.spark.streams.active:
+ q.stop()
+ self.assertTrue(df.isStreaming)
+
+ output_schema = StructType(
+ [
+ StructField("id1", StringType(), True),
+ StructField("id2", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+
+ q = (
+ df.groupBy("id1", "id2")
+ .transformWithState(
+ statefulProcessor=stateful_processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=initial_state,
+ )
+ .writeStream.queryName("this_query")
+ .foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+
+ self.assertEqual(q.name, "this_query")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+ self.assertTrue(q.exception() is None)
+
+ def test_transform_with_state_non_contiguous_grouping_cols(self):
+ def check_results(batch_df, batch_id):
+ assert set(batch_df.collect()) == {
+ Row(id1="0", id2="1", value=str(123 + 46)),
+ Row(id1="1", id2="2", value=str(146 + 346)),
+ }
+
+ self._test_transform_with_state_non_contiguous_grouping_cols(
+ SimpleStatefulProcessorWithInitialState(), check_results
+ )
+
+ def test_transform_with_state_non_contiguous_grouping_cols_with_init_state(self):
+ def check_results(batch_df, batch_id):
+ # initial state for key (0, 1) is processed
+ assert set(batch_df.collect()) == {
+ Row(id1="0", id2="1", value=str(789 + 123 + 46)),
+ Row(id1="1", id2="2", value=str(146 + 346)),
+ }
+
+ # grouping key of initial state is also not starting from the beginning of attributes
+ data = [(789, "0", "1"), (987, "3", "2")]
+ initial_state = self.spark.createDataFrame(
+ data, "initVal int, id1 string, id2 string"
+ ).groupBy("id1", "id2")
+
+ self._test_transform_with_state_non_contiguous_grouping_cols(
+ SimpleStatefulProcessorWithInitialState(), check_results, initial_state
+ )
+
+ def _test_transform_with_state_chaining_ops(
+ self, stateful_processor, check_results, timeMode="None", grouping_cols=["outputTimestamp"]
+ ):
+ import pyspark.sql.functions as f
+
+ input_path = tempfile.mkdtemp()
+ self._prepare_input_data(input_path + "/text-test3.txt", ["a", "b"], [10, 15])
+ time.sleep(2)
+ self._prepare_input_data(input_path + "/text-test4.txt", ["a", "c"], [11, 25])
+ time.sleep(2)
+ self._prepare_input_data(input_path + "/text-test1.txt", ["a"], [5])
+
+ df = self._build_test_df(input_path)
+ df = df.select(
+ "id", f.from_unixtime(f.col("temperature")).alias("eventTime").cast("timestamp")
+ ).withWatermark("eventTime", "5 seconds")
+
+ for q in self.spark.streams.active:
+ q.stop()
+ self.assertTrue(df.isStreaming)
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("outputTimestamp", TimestampType(), True),
+ ]
+ )
+
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=stateful_processor,
+ outputStructType=output_schema,
+ outputMode="Append",
+ timeMode=timeMode,
+ eventTimeColumnName="outputTimestamp",
+ )
+ .groupBy(grouping_cols)
+ .count()
+ .writeStream.queryName("chaining_ops_query")
+ .foreachBatch(check_results)
+ .outputMode("append")
+ .start()
+ )
+
+ self.assertEqual(q.name, "chaining_ops_query")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+
+ def test_transform_with_state_chaining_ops(self):
+ def check_results(batch_df, batch_id):
+ import datetime
+
+ if batch_id == 0:
+ assert batch_df.isEmpty()
+ elif batch_id == 1:
+ # eviction watermark = 15 - 5 = 10 (max event time from batch 0),
+ # late event watermark = 0 (eviction event time from batch 0)
+ assert set(
+ batch_df.sort("outputTimestamp").select("outputTimestamp", "count").collect()
+ ) == {
+ Row(outputTimestamp=datetime.datetime(1970, 1, 1, 0, 0, 10), count=1),
+ }
+ elif batch_id == 2:
+ # eviction watermark = 25 - 5 = 20, late event watermark = 10;
+ # row with watermark=5<10 is dropped so it does not show up in the results;
+ # row with eventTime<=20 are finalized and emitted
+ assert set(
+ batch_df.sort("outputTimestamp").select("outputTimestamp", "count").collect()
+ ) == {
+ Row(outputTimestamp=datetime.datetime(1970, 1, 1, 0, 0, 11), count=1),
+ Row(outputTimestamp=datetime.datetime(1970, 1, 1, 0, 0, 15), count=1),
+ }
+
+ self._test_transform_with_state_chaining_ops(
+ StatefulProcessorChainingOps(), check_results, "eventTime"
+ )
+ self._test_transform_with_state_chaining_ops(
+ StatefulProcessorChainingOps(), check_results, "eventTime", ["outputTimestamp", "id"]
+ )
+
+ def test_transform_with_state_init_state_with_timers(self):
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ # timers are registered and handled in the first batch for
+ # rows in initial state; For key=0 and key=3 which contains
+ # expired timers, both should be handled by handleExpiredTimers
+ # regardless of whether key exists in the data rows or not
+ expired_df = batch_df.filter(batch_df["id"].contains("expired"))
+ data_df = batch_df.filter(~batch_df["id"].contains("expired"))
+ assert set(expired_df.sort("id").select("id").collect()) == {
+ Row(id="0-expired"),
+ Row(id="3-expired"),
+ }
+ assert set(data_df.sort("id").collect()) == {
+ Row(id="0", value=str(789 + 123 + 46)),
+ Row(id="1", value=str(146 + 346)),
+ }
+ else:
+ # handleInitialState is only processed in the first batch,
+ # no more timer is registered so no more expired timers
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", value=str(789 + 123 + 46 + 67)),
+ Row(id="3", value=str(987 + 12)),
+ }
+
+ self._test_transform_with_state_init_state(
+ StatefulProcessorWithInitialStateTimers(), check_results, "processingTime"
+ )
+
+ def test_transform_with_state_batch_query(self):
+ data = [("0", 123), ("0", 46), ("1", 146), ("1", 346)]
+ df = self.spark.createDataFrame(data, "id string, temperature int")
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("countAsString", StringType(), True),
+ ]
+ )
+ batch_result = df.groupBy("id").transformWithState(
+ statefulProcessor=MapStateProcessor(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+ assert set(batch_result.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ def test_transform_with_state_batch_query_initial_state(self):
+ data = [("0", 123), ("0", 46), ("1", 146), ("1", 346)]
+ df = self.spark.createDataFrame(data, "id string, temperature int")
+
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, initVal int").groupBy(
+ "id"
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ batch_result = df.groupBy("id").transformWithState(
+ statefulProcessor=SimpleStatefulProcessorWithInitialState(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=initial_state,
+ )
+ assert set(batch_result.sort("id").collect()) == {
+ Row(id="0", value=str(789 + 123 + 46)),
+ Row(id="1", value=str(146 + 346)),
+ }
+
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ @unittest.skipIf(
+ "COVERAGE_PROCESS_START" in os.environ, "Flaky with coverage enabled, skipping for now."
+ )
+ def test_transform_with_map_state_metadata(self):
+ self._test_transform_with_map_state_metadata(None)
+
+ def test_transform_with_map_state_metadata_with_init_state(self):
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_map_state_metadata(initial_state)
+
+ def _test_transform_with_map_state_metadata(self, initial_state):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = batch_df.sparkSession.read.format("state-metadata").load(
+ checkpoint_path
+ )
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ ).sort("groupingKey").collect() == [
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ ]
+
+ # check for map state with flatten option
+ map_state_df_non_flatten = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert map_state_df_non_flatten.select(
+ "key.id", explode(col("map_value")).alias("map_key", "map_value")
+ ).selectExpr(
+ "id AS groupingKey",
+ "map_key.name AS mapKey",
+ "map_value.value.count AS mapValue",
+ ).sort(
+ "groupingKey"
+ ).collect() == [
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ ]
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ self._test_transform_with_state_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=initial_state,
+ )
+
+ # This test covers multiple list state variables and flatten option
+ def test_transform_with_list_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = batch_df.sparkSession.read.format("state-metadata").load(
+ checkpoint_path
+ )
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 4
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["listState1", "listState2", "listStateTimestamp"]:
+ assert state_var["stateVariableType"] == "ListState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ list_state_1_df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .option("flattenCollectionTypes", True)
+ .load()
+ )
+ assert list_state_1_df.selectExpr(
+ "key.id AS groupingKey",
+ "list_element.temperature AS listElement",
+ ).sort("groupingKey", "listElement").collect() == [
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=111),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=111),
+ Row(groupingKey="1", listElement=120),
+ Row(groupingKey="1", listElement=120),
+ ]
+
+ list_state_2_df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState2")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert list_state_2_df.selectExpr(
+ "key.id AS groupingKey", "list_value.temperature AS valueList"
+ ).sort("groupingKey").withColumn(
+ "valueSortedList", array_sort(col("valueList"))
+ ).select(
+ "groupingKey", "valueSortedList"
+ ).collect() == [
+ Row(groupingKey="0", valueSortedList=[20, 20, 120, 120, 222]),
+ Row(groupingKey="1", valueSortedList=[20, 20, 120, 120, 222]),
+ ]
+
+ self._test_transform_with_state_basic(
+ ListStateProcessor(),
+ check_results,
+ False,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # This test covers value state variable and read change feed,
+ # snapshotStartBatchId related options
+ def test_transform_with_value_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ # check for state metadata source
+ metadata_df = batch_df.sparkSession.read.format("state-metadata").load(
+ checkpoint_path
+ )
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["numViolations", "tempState"]:
+ state_var["stateVariableType"] == "ValueState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and readChangeFeed
+ value_state_df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("readChangeFeed", True)
+ .option("changeStartBatchId", 0)
+ .load()
+ ).selectExpr(
+ "change_type", "key.id AS groupingKey", "value.value AS value", "partition_id"
+ )
+
+ assert value_state_df.select("change_type", "groupingKey", "value").sort(
+ "groupingKey"
+ ).collect() == [
+ Row(change_type="update", groupingKey="0", value=1),
+ Row(change_type="update", groupingKey="1", value=2),
+ ]
+
+ partition_id_list = [
+ row["partition_id"] for row in value_state_df.select("partition_id").collect()
+ ]
+
+ for partition_id in partition_id_list:
+ # check for state data source and snapshotStartBatchId options
+ state_snapshot_df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("snapshotPartitionId", partition_id)
+ .option("snapshotStartBatchId", 0)
+ .load()
+ )
+
+ assert (
+ value_state_df.select("partition_id", "groupingKey", "value")
+ .filter(value_state_df["partition_id"] == partition_id)
+ .sort("groupingKey")
+ .collect()
+ == state_snapshot_df.selectExpr(
+ "partition_id", "key.id AS groupingKey", "value.value AS value"
+ )
+ .sort("groupingKey")
+ .collect()
+ )
+
+ with self.sql_conf(
+ {"spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled": "true"}
+ ):
+ self._test_transform_with_state_basic(
+ SimpleStatefulProcessor(),
+ check_results,
+ False,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ )
+
+ def test_transform_with_state_restart_with_multiple_rows_init_state(self):
+ def check_results(batch_df, _):
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+
+ def check_results_for_new_query(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", value=str(123 + 46)),
+ Row(id="1", value=str(146 + 346)),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", value=str(123 + 46 + 67)),
+ Row(id="3", value=str(12)),
+ }
+ # verify values in initial state is appended into list state for all keys
+ df = (
+ batch_df.sparkSession.read.format("statestore")
+ .option("path", new_checkpoint_path)
+ .option("stateVarName", "list_state")
+ .load()
+ ).selectExpr("key.id AS id", "list_element.value AS value")
+
+ def dataframe_to_value_list(output_df):
+ return [
+ row["value"] for row in output_df.sort("value").select("value").collect()
+ ]
+
+ assert dataframe_to_value_list(df.filter(df.id == "0")) == [20, 20, 111, 120, 120]
+ assert dataframe_to_value_list(df.filter(df.id == "1")) == [20, 20, 111, 120, 120]
+
+ # run a tws query and read state data source dataframe from its checkpoint
+ checkpoint_path = tempfile.mkdtemp()
+ self._test_transform_with_state_basic(
+ ListStateProcessor(), check_results, True, checkpoint_path=checkpoint_path
+ )
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .load()
+ ).selectExpr("key.id AS id", "list_element.temperature AS initVal")
+ init_df = list_state_df.groupBy("id")
+
+ # run a new tws query and pass state data source dataframe as initial state
+ # multiple rows exist in the initial state with the same grouping key
+ new_checkpoint_path = tempfile.mkdtemp()
+ self._test_transform_with_state_init_state(
+ StatefulProcessorWithListStateInitialState(),
+ check_results_for_new_query,
+ checkpoint_path=new_checkpoint_path,
+ initial_state=init_df,
+ )
+
+ # run the same test suites again but with single shuffle partition
+ def test_transform_with_state_with_timers_single_partition(self):
+ with self.sql_conf({"spark.sql.shuffle.partitions": "1"}):
+ self.test_transform_with_state_init_state_with_timers()
+ self.test_transform_with_state_event_time()
+ self.test_transform_with_state_proc_timer()
+ self.test_transform_with_state_restart_with_multiple_rows_init_state()
+
+ def _run_evolution_test(self, processor, checkpoint_dir, check_results, df):
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", processor.state_schema, True),
+ ]
+ )
+
+ # Stop any active streams first
+ for q in self.spark.streams.active:
+ q.stop()
+
+ q = (
+ df.groupBy("id")
+ .transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+ .writeStream.queryName("evolution_test")
+ .option("checkpointLocation", checkpoint_dir)
+ .foreachBatch(check_results)
+ .outputMode("update")
+ .start()
+ )
+
+ self.assertEqual(q.name, "evolution_test")
+ self.assertTrue(q.isActive)
+ q.processAllAvailable()
+ q.awaitTermination(10)
+
+ def test_schema_evolution_scenarios(self):
+ """Test various schema evolution scenarios"""
+ with self.sql_conf({"spark.sql.streaming.stateStore.encodingFormat": "avro"}):
+ with tempfile.TemporaryDirectory() as checkpoint_dir:
+ # Test 1: Basic state
+
+ input_path = tempfile.mkdtemp()
+ self._prepare_test_resource1(input_path)
+
+ df = self._build_test_df(input_path)
+
+ def check_basic_state(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["id"] == 0 # First ID from test data
+ assert result.value["name"] == "name-0"
+
+ self._run_evolution_test(BasicProcessor(), checkpoint_dir, check_basic_state, df)
+
+ self._prepare_test_resource2(input_path)
+
+ # Test 2: Add fields
+ def check_add_fields(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ # Check default values for existing key
+ assert result.value["id"] == 0
+ assert result.value["count"] is None
+ assert result.value["active"] is None
+ assert result.value["score"] is None
+
+ self._run_evolution_test(AddFieldsProcessor(), checkpoint_dir, check_add_fields, df)
+ self._prepare_test_resource3(input_path)
+
+ # Test 3: Remove fields
+ def check_remove_fields(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["id"] == 0 # First ID from test data
+ assert result.value["name"] == "name-00"
+
+ self._run_evolution_test(
+ RemoveFieldsProcessor(), checkpoint_dir, check_remove_fields, df
+ )
+ self._prepare_test_resource4(input_path)
+
+ # Test 4: Reorder fields
+ def check_reorder_fields(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["name"] == "name-00"
+ assert result.value["id"] == 0
+
+ self._run_evolution_test(
+ ReorderedFieldsProcessor(), checkpoint_dir, check_reorder_fields, df
+ )
+ self._prepare_test_resource5(input_path)
+
+ # Test 5: Upcast type
+ def check_upcast(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["id"] == 1
+ assert result.value["name"] == "name-0"
+
+ self._run_evolution_test(UpcastProcessor(), checkpoint_dir, check_upcast, df)
+
+ # This test case verifies that an exception is thrown when downcasting, which violates
+ # Avro's schema evolution rules
+ def test_schema_evolution_fails(self):
+ with self.sql_conf({"spark.sql.streaming.stateStore.encodingFormat": "avro"}):
+ with tempfile.TemporaryDirectory() as checkpoint_dir:
+ input_path = tempfile.mkdtemp()
+ self._prepare_test_resource1(input_path)
+
+ df = self._build_test_df(input_path)
+
+ def check_add_fields(batch_df, batch_id):
+ results = batch_df.collect()
+ assert results[0].value["count"] == 100
+ assert results[0].value["active"]
+
+ self._run_evolution_test(AddFieldsProcessor(), checkpoint_dir, check_add_fields, df)
+ self._prepare_test_resource2(input_path)
+
+ def check_upcast(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["name"] == "name-0"
+
+ # Long
+ self._run_evolution_test(UpcastProcessor(), checkpoint_dir, check_upcast, df)
+ self._prepare_test_resource3(input_path)
+
+ def check_basic_state(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["id"] == 0 # First ID from test data
+ assert result.value["name"] == "name-0"
+
+ # Int
+ try:
+ self._run_evolution_test(
+ BasicProcessor(),
+ checkpoint_dir,
+ check_basic_state,
+ df,
+ )
+ except Exception as e:
+ # we are expecting an exception, verify it's the right one
+ from pyspark.errors.exceptions.captured import StreamingQueryException
+
+ if not isinstance(e, StreamingQueryException):
+ return False
+ error_msg = str(e)
+ assert (
+ "[STREAM_FAILED]" in error_msg
+ and "[STATE_STORE_INVALID_VALUE_SCHEMA_EVOLUTION]" in error_msg
+ and "Schema evolution is not possible" in error_msg
+ )
+
+ def test_not_nullable_fails(self):
+ with self.sql_conf({"spark.sql.streaming.stateStore.encodingFormat": "avro"}):
+ with tempfile.TemporaryDirectory() as checkpoint_dir:
+ input_path = tempfile.mkdtemp()
+ self._prepare_test_resource1(input_path)
+
+ df = self._build_test_df(input_path)
+
+ def check_basic_state(batch_df, batch_id):
+ result = batch_df.collect()[0]
+ assert result.value["id"] == 0 # First ID from test data
+ assert result.value["name"] == "name-0"
+
+ try:
+ self._run_evolution_test(
+ BasicProcessorNotNullable(),
+ checkpoint_dir,
+ check_basic_state,
+ df,
+ )
+ except Exception as e:
+ # we are expecting an exception, verify it's the right one
+ from pyspark.errors.exceptions.captured import StreamingQueryException
+
+ if not isinstance(e, StreamingQueryException):
+ return False
+ error_msg = str(e)
+ assert (
+ "[TRANSFORM_WITH_STATE_SCHEMA_MUST_BE_NULLABLE]" in error_msg
+ and "column family state must be nullable" in error_msg
+ )
+
+
+class SimpleStatefulProcessorWithInitialState(StatefulProcessor):
+ # this dict is the same as input initial state dataframe
+ dict = {("0",): 789, ("3",): 987}
+
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("value", IntegerType(), True)])
+ self.value_state = handle.getValueState("value_state", state_schema)
+ self.handle = handle
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ exists = self.value_state.exists()
+ if exists:
+ value_row = self.value_state.get()
+ existing_value = value_row[0]
+ else:
+ existing_value = 0
+
+ accumulated_value = existing_value
+
+ for row in rows:
+ value = row["temperature"]
+ accumulated_value += value
+
+ self.value_state.update((accumulated_value,))
+
+ if len(key) > 1:
+ yield Row(id1=key[0], id2=key[1], value=str(accumulated_value))
+ else:
+ yield Row(id=key, value=str(accumulated_value))
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ init_val = initialState.at[0, "initVal"]
+ self.value_state.update((init_val,))
+ if len(key) == 1:
+ assert self.dict[key] == init_val
+
+ def close(self) -> None:
+ pass
+
+
+class StatefulProcessorWithInitialStateTimers(SimpleStatefulProcessorWithInitialState):
+ def handleExpiredTimer(self, key, timerValues, expiredTimerInfo) -> Iterator[Row]:
+ self.handle.deleteTimer(expiredTimerInfo.getExpiryTimeInMs())
+ str_key = f"{str(key[0])}-expired"
+ yield Row(id=str_key, value=str(expiredTimerInfo.getExpiryTimeInMs()))
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ super().handleInitialState(key, initialState, timerValues)
+ self.handle.registerTimer(timerValues.getCurrentProcessingTimeInMs() - 1)
+
+
+class StatefulProcessorWithListStateInitialState(SimpleStatefulProcessorWithInitialState):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ super().init(handle)
+ list_ele_schema = StructType([StructField("value", IntegerType(), True)])
+ self.list_state = handle.getListState("list_state", list_ele_schema)
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ for val in initialState["initVal"].tolist():
+ self.list_state.appendValue((val,))
+
+
+# A stateful processor that output the max event time it has seen. Register timer for
+# current watermark. Clear max state if timer expires.
+class EventTimeStatefulProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("value", StringType(), True)])
+ self.handle = handle
+ self.max_state = handle.getValueState("max_state", state_schema)
+
+ def handleExpiredTimer(self, key, timerValues, expiredTimerInfo) -> Iterator[Row]:
+ self.max_state.clear()
+ self.handle.deleteTimer(expiredTimerInfo.getExpiryTimeInMs())
+ str_key = f"{str(key[0])}-expired"
+ yield Row(id=str_key, timestamp=str(expiredTimerInfo.getExpiryTimeInMs()))
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ timestamp_list = []
+ for row in rows:
+ # timestamp is microsecond, restore to second
+ timestamp_list.append(row["eventTime"] // 1000000)
+
+ if self.max_state.exists():
+ cur_max = int(self.max_state.get()[0])
+ else:
+ cur_max = 0
+ max_event_time = str(max(cur_max, max(timestamp_list)))
+
+ self.max_state.update((max_event_time,))
+ self.handle.registerTimer(timerValues.getCurrentWatermarkInMs())
+
+ yield Row(id=key, timestamp=max_event_time)
+
+ def close(self) -> None:
+ pass
+
+
+# A stateful processor that output the accumulation of count of input rows; register
+# processing timer and clear the counter if timer expires.
+class ProcTimeStatefulProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("value", StringType(), True)])
+ self.handle = handle
+ self.count_state = handle.getValueState("count_state", state_schema)
+
+ def handleExpiredTimer(self, key, timerValues, expiredTimerInfo) -> Iterator[Row]:
+ # reset count state each time the timer is expired
+ timer_list_1 = [e for e in self.handle.listTimers()]
+ timer_list_2 = []
+ idx = 0
+ for e in self.handle.listTimers():
+ timer_list_2.append(e)
+ # check multiple iterator on the same grouping key works
+ assert timer_list_2[idx] == timer_list_1[idx]
+ idx += 1
+
+ if len(timer_list_1) > 0:
+ assert len(timer_list_1) == 2
+ self.count_state.clear()
+ self.handle.deleteTimer(expiredTimerInfo.getExpiryTimeInMs())
+ yield Row(
+ id=key,
+ countAsString=str(-1),
+ timeValues=str(expiredTimerInfo.getExpiryTimeInMs())
+ )
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ if not self.count_state.exists():
+ count = 0
+ else:
+ count = int(self.count_state.get()[0])
+
+ if key == ("0",):
+ self.handle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + 1)
+
+ rows_count = 0
+ for row in rows:
+ rows_count += 1
+
+ count = count + rows_count
+
+ self.count_state.update((str(count),))
+ timestamp = str(timerValues.getCurrentProcessingTimeInMs())
+
+ yield Row(id=key, countAsString=str(count), timeValues=timestamp)
+
+ def close(self) -> None:
+ pass
+
+
+class SimpleStatefulProcessor(StatefulProcessor, unittest.TestCase):
+ dict = {0: {"0": 1, "1": 2}, 1: {"0": 4, "1": 3}}
+ batch_id = 0
+
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ # Test both string type and struct type schemas
+ self.num_violations_state = handle.getValueState("numViolations", "value int")
+ state_schema = StructType([StructField("value", IntegerType(), True)])
+ self.temp_state = handle.getValueState("tempState", state_schema)
+ handle.deleteIfExists("tempState")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ with self.assertRaisesRegex(PySparkRuntimeError, "Error checking value state exists"):
+ self.temp_state.exists()
+ new_violations = 0
+ count = 0
+ key_str = key[0]
+ exists = self.num_violations_state.exists()
+ if exists:
+ existing_violations_row = self.num_violations_state.get()
+ existing_violations = existing_violations_row[0]
+ assert existing_violations == self.dict[0][key_str]
+ self.batch_id = 1
+ else:
+ existing_violations = 0
+ for row in rows:
+ # temperature should be non-NA to be counted
+ temperature = row["temperature"]
+ if temperature != None:
+ count += 1
+ if temperature > 100:
+ new_violations += 1
+ updated_violations = new_violations + existing_violations
+ assert updated_violations == self.dict[self.batch_id][key_str]
+ self.num_violations_state.update((updated_violations,))
+ yield Row(id=key, countAsString=str(count))
+
+ def close(self) -> None:
+ pass
+
+
+class StatefulProcessorChainingOps(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ pass
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ timestamp_list = []
+ for row in rows:
+ timestamp_list.append(row["eventTime"])
+ yield Row(id=key, outputTimestamp=timestamp_list[0])
+
+ def close(self) -> None:
+ pass
+
+
+# A stateful processor that inherit all behavior of SimpleStatefulProcessor except that it use
+# ttl state with a large timeout.
+class SimpleTTLStatefulProcessor(SimpleStatefulProcessor, unittest.TestCase):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("value", IntegerType(), True)])
+ self.num_violations_state = handle.getValueState("numViolations", state_schema, 30000)
+ self.temp_state = handle.getValueState("tempState", state_schema)
+ handle.deleteIfExists("tempState")
+
+
+class TTLStatefulProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("value", IntegerType(), True)])
+ user_key_schema = StructType([StructField("id", StringType(), True)])
+ self.ttl_count_state = handle.getValueState("ttl-state", state_schema, 10000)
+ self.count_state = handle.getValueState("state", state_schema)
+ self.ttl_list_state = handle.getListState("ttl-list-state", state_schema, 10000)
+ self.ttl_map_state = handle.getMapState(
+ "ttl-map-state", user_key_schema, state_schema, 10000
+ )
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ count = 0
+ ttl_count = 0
+ ttl_list_state_count = 0
+ ttl_map_state_count = 0
+ id = key[0]
+ if self.count_state.exists():
+ count = self.count_state.get()[0]
+ if self.ttl_count_state.exists():
+ ttl_count = self.ttl_count_state.get()[0]
+ if self.ttl_list_state.exists():
+ iter = self.ttl_list_state.get()
+ for s in iter:
+ ttl_list_state_count += s[0]
+ if self.ttl_map_state.exists():
+ ttl_map_state_count = self.ttl_map_state.getValue(key)[0]
+ for row in rows:
+ if row["temperature"] is not None:
+ count += 1
+ ttl_count += 1
+ ttl_list_state_count += 1
+ ttl_map_state_count += 1
+
+ self.count_state.update((count,))
+ # skip updating state for the 2nd batch so that ttl state expire
+ if not (ttl_count == 2 and id == "0"):
+ self.ttl_count_state.update((ttl_count,))
+ self.ttl_list_state.put([(ttl_list_state_count,), (ttl_list_state_count,)])
+ self.ttl_map_state.updateValue(key, (ttl_map_state_count,))
+ yield Row(
+ id=[
+ f"ttl-count-{id}",
+ f"count-{id}",
+ f"ttl-list-state-count-{id}",
+ f"ttl-map-state-count-{id}",
+ ],
+ count=[ttl_count, count, ttl_list_state_count, ttl_map_state_count]
+ )
+
+
+class InvalidSimpleStatefulProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("value", IntegerType(), True)])
+ self.num_violations_state = handle.getValueState("numViolations", state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ count = 0
+ exists = self.num_violations_state.exists()
+ assert not exists
+ # try to get a state variable with no value
+ assert self.num_violations_state.get() is None
+ self.num_violations_state.clear()
+ yield Row(id=key, countAsString=str(count))
+
+
+class ListStateProcessor(StatefulProcessor):
+ # Dict to store the expected results. The key represents the grouping key string, and the value
+ # is a dictionary of pandas dataframe index -> expected temperature value. Since we set
+ # maxRecordsPerBatch to 2, we expect the pandas dataframe dictionary to have 2 entries.
+ dict = {0: 120, 1: 20}
+
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("temperature", IntegerType(), True)])
+ timestamp_schema = StructType([StructField("time", TimestampType(), True)])
+ self.list_state1 = handle.getListState("listState1", state_schema)
+ self.list_state2 = handle.getListState("listState2", state_schema)
+ self.list_state_timestamp = handle.getListState("listStateTimestamp", timestamp_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ import datetime
+
+ count = 0
+ time_list = []
+ for row in rows:
+ list_state_rows = [(120,), (20,)]
+ self.list_state1.put(list_state_rows)
+ self.list_state2.put(list_state_rows)
+ self.list_state1.appendValue((111,))
+ self.list_state2.appendValue((222,))
+ self.list_state1.appendList(list_state_rows)
+ self.list_state2.appendList(list_state_rows)
+
+ if row["temperature"] is not None:
+ count += 1
+
+ current_processing_time = datetime.datetime.fromtimestamp(
+ timerValues.getCurrentProcessingTimeInMs() / 1000
+ )
+ stored_time = current_processing_time + datetime.timedelta(minutes=1)
+ time_list.append((stored_time,))
+ iter1 = self.list_state1.get()
+ iter2 = self.list_state2.get()
+ # Mixing the iterator to test it we can resume from the correct point
+ assert next(iter1)[0] == self.dict[0]
+ assert next(iter2)[0] == self.dict[0]
+ assert next(iter1)[0] == self.dict[1]
+ assert next(iter2)[0] == self.dict[1]
+ # Get another iterator for list_state1 to test if the 2 iterators (iter1 and iter3) don't
+ # interfere with each other.
+ iter3 = self.list_state1.get()
+ assert next(iter3)[0] == self.dict[0]
+ assert next(iter3)[0] == self.dict[1]
+ # the second arrow batch should contain the appended value 111 for list_state1 and
+ # 222 for list_state2
+ assert next(iter1)[0] == 111
+ assert next(iter2)[0] == 222
+ assert next(iter3)[0] == 111
+ # since we put another 2 rows after 111/222, check them here
+ assert next(iter1)[0] == self.dict[0]
+ assert next(iter2)[0] == self.dict[0]
+ assert next(iter3)[0] == self.dict[0]
+ assert next(iter1)[0] == self.dict[1]
+ assert next(iter2)[0] == self.dict[1]
+ assert next(iter3)[0] == self.dict[1]
+ if time_list:
+ # Validate timestamp type can work properly with arrow transmission
+ self.list_state_timestamp.put(time_list)
+ yield Row(id=key, countAsString=str(count))
+
+ def close(self) -> None:
+ pass
+
+
+class ListStateLargeTTLProcessor(ListStateProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ state_schema = StructType([StructField("temperature", IntegerType(), True)])
+ timestamp_schema = StructType([StructField("time", TimestampType(), True)])
+ self.list_state1 = handle.getListState("listState1", state_schema, 30000)
+ self.list_state2 = handle.getListState("listState2", state_schema, 30000)
+ self.list_state_timestamp = handle.getListState("listStateTimestamp", timestamp_schema)
+
+
+class MapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle):
+ # Test string type schemas
+ self.map_state = handle.getMapState("mapState", "name string", "count int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ count = 0
+ key1 = ("key1",)
+ key2 = ("key2",)
+ for row in rows:
+ if row["temperature"] is not None:
+ count += 1
+ value1 = count
+ value2 = count
+ if self.map_state.exists():
+ if self.map_state.containsKey(key1):
+ value1 += self.map_state.getValue(key1)[0]
+ if self.map_state.containsKey(key2):
+ value2 += self.map_state.getValue(key2)[0]
+ self.map_state.updateValue(key1, (value1,))
+ self.map_state.updateValue(key2, (value2,))
+ key_iter = self.map_state.keys()
+ assert next(key_iter)[0] == "key1"
+ assert next(key_iter)[0] == "key2"
+ value_iter = self.map_state.values()
+ assert next(value_iter)[0] == value1
+ assert next(value_iter)[0] == value2
+ map_iter = self.map_state.iterator()
+ assert next(map_iter)[0] == key1
+ assert next(map_iter)[1] == (value2,)
+ self.map_state.removeKey(key1)
+ assert not self.map_state.containsKey(key1)
+ yield Row(id=key, countAsString=str(count))
+
+ def close(self) -> None:
+ pass
+
+
+# A stateful processor that inherit all behavior of MapStateProcessor except that it use
+# ttl state with a large timeout.
+class MapStateLargeTTLProcessor(MapStateProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ key_schema = StructType([StructField("name", StringType(), True)])
+ value_schema = StructType([StructField("count", IntegerType(), True)])
+ self.map_state = handle.getMapState("mapState", key_schema, value_schema, 30000)
+ self.list_state = handle.getListState("listState", key_schema)
+
+
+class BasicProcessor(StatefulProcessor):
+ # Schema definitions
+ state_schema = StructType(
+ [StructField("id", IntegerType(), True), StructField("name", StringType(), True)]
+ )
+
+ def init(self, handle):
+ self.state = handle.getValueState("state", self.state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ for pdf in rows:
+ pass
+ id_val = int(key[0])
+ name = f"name-{id_val}"
+ self.state.update((id_val, name))
+ yield Row(id=key[0], value={"id": id_val, "name": name})
+
+ def close(self) -> None:
+ pass
+
+
+class BasicProcessorNotNullable(StatefulProcessor):
+ # Schema definitions
+ state_schema = StructType(
+ [StructField("id", IntegerType(), False), StructField("name", StringType(), False)]
+ )
+
+ def init(self, handle):
+ self.state = handle.getValueState("state", self.state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ for pdf in rows:
+ pass
+ id_val = int(key[0])
+ name = f"name-{id_val}"
+ self.state.update((id_val, name))
+ yield Row(id=key[0], value={"id": id_val, "name": name})
+
+ def close(self) -> None:
+ pass
+
+
+class AddFieldsProcessor(StatefulProcessor):
+ state_schema = StructType(
+ [
+ StructField("id", IntegerType(), True),
+ StructField("name", StringType(), True),
+ StructField("count", IntegerType(), True),
+ StructField("active", BooleanType(), True),
+ StructField("score", FloatType(), True),
+ ]
+ )
+
+ def init(self, handle):
+ self.state = handle.getValueState("state", self.state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ for pdf in rows:
+ pass
+ id_val = int(key[0])
+ name = f"name-{id_val}"
+
+ if self.state.exists():
+ state_data = self.state.get()
+ state_dict = {
+ "id": state_data[0],
+ "name": state_data[1],
+ "count": state_data[2],
+ "active": state_data[3],
+ "score": state_data[4],
+ }
+ else:
+ state_dict = {
+ "id": id_val,
+ "name": name,
+ "count": 100,
+ "active": True,
+ "score": 99.9,
+ }
+
+ self.state.update(
+ (
+ state_dict["id"],
+ state_dict["name"] + "0",
+ state_dict["count"],
+ state_dict["active"],
+ state_dict["score"],
+ )
+ )
+ yield Row(id=key[0], value=state_dict)
+
+ def close(self) -> None:
+ pass
+
+
+class RemoveFieldsProcessor(StatefulProcessor):
+ # Schema definitions
+ state_schema = StructType(
+ [StructField("id", IntegerType(), True), StructField("name", StringType(), True)]
+ )
+
+ def init(self, handle):
+ self.state = handle.getValueState("state", self.state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ for pdf in rows:
+ pass
+ id_val = int(key[0])
+ name = f"name-{id_val}"
+ if self.state.exists():
+ name = self.state.get()[1]
+ self.state.update((id_val, name))
+ yield Row(id=key[0], value={"id": id_val, "name": name})
+
+ def close(self) -> None:
+ pass
+
+
+class ReorderedFieldsProcessor(StatefulProcessor):
+ state_schema = StructType(
+ [
+ StructField("name", StringType(), True),
+ StructField("id", IntegerType(), True),
+ StructField("score", FloatType(), True),
+ StructField("count", IntegerType(), True),
+ StructField("active", BooleanType(), True),
+ ]
+ )
+
+ def init(self, handle):
+ self.state = handle.getValueState("state", self.state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ for pdf in rows:
+ pass
+ id_val = int(key[0])
+ name = f"name-{id_val}"
+
+ if self.state.exists():
+ state_data = self.state.get()
+ state_dict = {
+ "name": state_data[0],
+ "id": state_data[1],
+ "score": state_data[2],
+ "count": state_data[3],
+ "active": state_data[4],
+ }
+ else:
+ state_dict = {
+ "name": name,
+ "id": id_val,
+ "score": 99.9,
+ "count": 100,
+ "active": True,
+ }
+ self.state.update(
+ (
+ state_dict["name"],
+ state_dict["id"],
+ state_dict["score"],
+ state_dict["count"],
+ state_dict["active"],
+ )
+ )
+ yield Row(id=key[0], value=state_dict)
+
+ def close(self) -> None:
+ pass
+
+
+class UpcastProcessor(StatefulProcessor):
+ state_schema = StructType(
+ [
+ StructField("id", LongType(), True), # Upcast from Int to Long
+ StructField("name", StringType(), True),
+ ]
+ )
+
+ def init(self, handle):
+ self.state = handle.getValueState("state", self.state_schema)
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
+ for pdf in rows:
+ pass
+ id_val = int(key[0])
+ name = f"name-{id_val}"
+ if self.state.exists():
+ id_val += self.state.get()[0] + 1
+ self.state.update((id_val, name))
+ yield Row(id=key[0], value={"id": id_val, "name": name})
+
+ def close(self) -> None:
+ pass
+
+
+class TransformWithStateInPySparkTests(TransformWithStateInPySparkTestsMixin, ReusedSQLTestCase):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.sql.tests.pandas.test_transform_with_state import * # noqa: F401
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 5a5a8d31e77d..0b0668e071d0 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -633,10 +633,17 @@ class PythonEvalType:
SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF: "PandasGroupedMapUDFTransformWithStateInitStateType" = ( # noqa: E501
212
)
+ # FIXME: typing...
+ SQL_TRANSFORM_WITH_STATE_UDF: "GroupedMapUDFTransformWithStateType" = 213
+ # FIXME: typing...
+ SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF: "GroupedMapUDFTransformWithStateInitStateType" = ( # noqa: E501
+ 214
+ )
SQL_TABLE_UDF: "SQLTableUDFType" = 300
SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301
+
def _create_local_socket(sock_info: "JavaArray") -> "io.BufferedRWPair":
"""
Create a local socket that can be used to load deserialized data from the JVM
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 44a5d0b91131..87a1f523ec19 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -34,7 +34,7 @@ from pyspark.accumulators import (
_deserialize_accumulator,
)
from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
-from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPySparkFuncMode
from pyspark.taskcontext import BarrierTaskContext, TaskContext
from pyspark.resource import ResourceInformation
from pyspark.util import PythonEvalType, local_connect_and_auth
@@ -60,6 +60,7 @@ from pyspark.sql.pandas.serializers import (
ApplyInPandasWithStateSerializer,
TransformWithStateInPandasSerializer,
TransformWithStateInPandasInitStateSerializer,
+ TransformWithStateInPySparkRowSerializer,
)
from pyspark.sql.pandas.types import to_arrow_type
from pyspark.sql.types import (
@@ -570,6 +571,32 @@ def wrap_grouped_transform_with_state_pandas_init_state_udf(f, return_type, runn
arrow_return_type = to_arrow_type(return_type, use_large_var_types(runner_conf))
return lambda p, m, k, v: [(wrapped(p, m, k, v), arrow_return_type)]
+def wrap_grouped_transform_with_state_udf(f, return_type, runner_conf):
+ def wrapped(stateful_processor_api_client, mode, key, values):
+ result_iter = f(stateful_processor_api_client, mode, key, values)
+
+ # TODO(SPARK-XXXXX): add verification that elements in result_iter are
+ # indeed of type Row and confirm to assigned cols
+
+ return result_iter
+
+ arrow_return_type = to_arrow_type(return_type, use_large_var_types(runner_conf))
+ return lambda p, m, k, v: [(wrapped(p, m, k, v), arrow_return_type)]
+
+
+def wrap_grouped_transform_with_state_init_state_udf(f, return_type, runner_conf):
+ def wrapped(stateful_processor_api_client, mode, key, values):
+ state_values, init_states = itertools.tee(values, 2)
+ result_iter = f(stateful_processor_api_client, mode, key, state_values, init_states)
+
+ # TODO(SPARK-XXXXX): add verification that elements in result_iter are
+ # indeed of type pd.DataFrame and confirm to assigned cols
+
+ return result_iter
+
+ arrow_return_type = to_arrow_type(return_type, use_large_var_types(runner_conf))
+ return lambda p, m, k, v: [(wrapped(p, m, k, v), arrow_return_type)]
+
def wrap_grouped_map_pandas_udf_with_state(f, return_type, runner_conf):
"""
@@ -935,6 +962,14 @@ def read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index, profil
return args_offsets, wrap_grouped_transform_with_state_pandas_init_state_udf(
func, return_type, runner_conf
)
+ elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF:
+ return args_offsets, wrap_grouped_transform_with_state_udf(
+ func, return_type, runner_conf
+ )
+ elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF:
+ return args_offsets, wrap_grouped_transform_with_state_init_state_udf(
+ func, return_type, runner_conf
+ )
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
argspec = inspect.getfullargspec(chained_func) # signature was lost when wrapping it
return args_offsets, wrap_cogrouped_map_pandas_udf(func, return_type, argspec, runner_conf)
@@ -1532,6 +1567,8 @@ def read_udfs(pickleSer, infile, eval_type):
PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF,
):
# Load conf used for pandas_udf evaluation
num_conf = read_int(infile)
@@ -1546,6 +1583,8 @@ def read_udfs(pickleSer, infile, eval_type):
elif (
eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
or eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF
+ or eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF
+ or eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF
):
state_server_port = read_int(infile)
key_schema = StructType.fromJson(json.loads(utf8_deserializer.loads(infile)))
@@ -1595,6 +1634,26 @@ def read_udfs(pickleSer, infile, eval_type):
ser = TransformWithStateInPandasInitStateSerializer(
timezone, safecheck, _assign_cols_by_name, arrow_max_records_per_batch
)
+ elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF:
+ arrow_max_records_per_batch = runner_conf.get(
+ "spark.sql.execution.arrow.maxRecordsPerBatch", 10000
+ )
+ arrow_max_records_per_batch = int(arrow_max_records_per_batch)
+
+ # FIXME: do we need these constructor parameters?
+ ser = TransformWithStateInPySparkRowSerializer(
+ timezone, safecheck, _assign_cols_by_name, arrow_max_records_per_batch
+ )
+ elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF:
+ arrow_max_records_per_batch = runner_conf.get(
+ "spark.sql.execution.arrow.maxRecordsPerBatch", 10000
+ )
+ arrow_max_records_per_batch = int(arrow_max_records_per_batch)
+
+ # FIXME: add a new serializer
+ ser = TransformWithStateInPySparkInitStateSerializer(
+ timezone, safecheck, _assign_cols_by_name, arrow_max_records_per_batch
+ )
elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
ser = ArrowStreamUDFSerializer()
elif eval_type == PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF:
@@ -1777,7 +1836,7 @@ def read_udfs(pickleSer, infile, eval_type):
def mapper(a):
mode = a[0]
- if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
+ if mode == TransformWithStateInPySparkFuncMode.PROCESS_DATA:
key = a[1]
def values_gen():
@@ -1814,7 +1873,76 @@ def read_udfs(pickleSer, infile, eval_type):
def mapper(a):
mode = a[0]
- if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
+ if mode == TransformWithStateInPySparkFuncMode.PROCESS_DATA:
+ key = a[1]
+
+ def values_gen():
+ for x in a[2]:
+ retVal = [x[1][o] for o in parsed_offsets[0][1]]
+ initVal = [x[2][o] for o in parsed_offsets[1][1]]
+ yield retVal, initVal
+
+ # This must be generator comprehension - do not materialize.
+ return f(stateful_processor_api_client, mode, key, values_gen())
+ else:
+ # mode == PROCESS_TIMER or mode == COMPLETE
+ return f(stateful_processor_api_client, mode, None, iter([]))
+
+ elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_UDF:
+ # We assume there is only one UDF here because grouped map doesn't
+ # support combining multiple UDFs.
+ assert num_udfs == 1
+
+ # See TransformWithStateInPySparkExec for how arg_offsets are used to
+ # distinguish between grouping attributes and data attributes
+ arg_offsets, f = read_single_udf(
+ pickleSer, infile, eval_type, runner_conf, udf_index=0, profiler=profiler
+ )
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+ ser.key_offsets = parsed_offsets[0][0]
+ stateful_processor_api_client = StatefulProcessorApiClient(state_server_port, key_schema)
+
+ def mapper(a):
+ mode = a[0]
+
+ if mode == TransformWithStateInPySparkFuncMode.PROCESS_DATA:
+ key = a[1]
+
+ def values_gen():
+ for x in a[2]:
+ retVal = [x[1][o] for o in parsed_offsets[0][1]]
+ yield retVal
+
+ # This must be generator comprehension - do not materialize.
+ return f(stateful_processor_api_client, mode, key, values_gen())
+ else:
+ # mode == PROCESS_TIMER or mode == COMPLETE
+ return f(stateful_processor_api_client, mode, None, iter([]))
+
+ elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_INIT_STATE_UDF:
+ # We assume there is only one UDF here because grouped map doesn't
+ # support combining multiple UDFs.
+ assert num_udfs == 1
+
+ # See TransformWithStateInPandasExec for how arg_offsets are used to
+ # distinguish between grouping attributes and data attributes
+ arg_offsets, f = read_single_udf(
+ pickleSer, infile, eval_type, runner_conf, udf_index=0, profiler=profiler
+ )
+ # parsed offsets:
+ # [
+ # [groupingKeyOffsets, dedupDataOffsets],
+ # [initStateGroupingOffsets, dedupInitDataOffsets]
+ # ]
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+ ser.key_offsets = parsed_offsets[0][0]
+ ser.init_key_offsets = parsed_offsets[1][0]
+ stateful_processor_api_client = StatefulProcessorApiClient(state_server_port, key_schema)
+
+ def mapper(a):
+ mode = a[0]
+
+ if mode == TransformWithStateInPySparkFuncMode.PROCESS_DATA:
key = a[1]
def values_gen():
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala
index f7ab41bd6f96..146032f1d199 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala
@@ -105,7 +105,7 @@ object UnsupportedOperationChecker extends Logging {
case d: Deduplicate if d.isStreaming && d.keys.exists(hasEventTimeCol) => true
case d: DeduplicateWithinWatermark if d.isStreaming => true
case t: TransformWithState if t.isStreaming => true
- case t: TransformWithStateInPandas if t.isStreaming => true
+ case t: TransformWithStateInPySpark if t.isStreaming => true
case _ => false
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
index 2f9bf2b52190..76cf766f697f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
@@ -178,16 +178,18 @@ case class FlatMapGroupsInPandasWithState(
* @param outputAttrs used to define the output rows
* @param outputMode defines the output mode for the statefulProcessor
* @param timeMode the time mode semantics of the stateful processor for timers and TTL.
+ * // FIXME: document `userFacingDataType`
* @param child logical plan of the underlying data
* @param initialState logical plan of initial state
* @param initGroupingAttrsLen length of the seq of grouping attributes for initial state dataframe
*/
-case class TransformWithStateInPandas(
+case class TransformWithStateInPySpark(
functionExpr: Expression,
groupingAttributesLen: Int,
outputAttrs: Seq[Attribute],
outputMode: OutputMode,
timeMode: TimeMode,
+ userFacingDataType: TransformWithStateInPySpark.UserFacingDataType.Value,
child: LogicalPlan,
hasInitialState: Boolean,
initialState: LogicalPlan,
@@ -205,7 +207,7 @@ case class TransformWithStateInPandas(
AttributeSet(leftAttributes ++ rightAttributes ++ functionExpr.references) -- producedAttributes
override protected def withNewChildrenInternal(
- newLeft: LogicalPlan, newRight: LogicalPlan): TransformWithStateInPandas =
+ newLeft: LogicalPlan, newRight: LogicalPlan): TransformWithStateInPySpark =
copy(child = newLeft, initialState = newRight)
def leftAttributes: Seq[Attribute] = {
@@ -225,6 +227,13 @@ case class TransformWithStateInPandas(
}
}
+object TransformWithStateInPySpark {
+ object UserFacingDataType extends Enumeration {
+ val PYTHON_ROW = Value("python_row")
+ val PANDAS = Value("pandas")
+ }
+}
+
/**
* Flatmap cogroups using a udf: iter(pyarrow.RecordBatch) -> iter(pyarrow.RecordBatch)
* This is used by DataFrame.groupby().cogroup().applyInArrow().
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 1ff947a3436c..4ded7eb6f85f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -3465,11 +3465,16 @@ object SQLConf {
"than zero and less than INT_MAX.")
.createWithDefaultString("256MB")
- val ARROW_TRANSFORM_WITH_STATE_IN_PANDAS_MAX_RECORDS_PER_BATCH =
- buildConf("spark.sql.execution.arrow.transformWithStateInPandas.maxRecordsPerBatch")
- .doc("When using TransformWithStateInPandas, limit the maximum number of state records " +
- "that can be written to a single ArrowRecordBatch in memory.")
+ // FIXME: Consider consolidating this with the above.
+ val ARROW_TRANSFORM_WITH_STATE_IN_PYSPARK_MAX_RECORDS_PER_BATCH =
+ buildConf("spark.sql.execution.arrow.transformWithStateInPySpark.maxRecordsPerBatch")
+ .doc("When using TransformWithState in PySpark (both Python Row and Pandas), limit " +
+ "the maximum number of state records that can be written to a single ArrowRecordBatch " +
+ "in memory.")
.version("4.0.0")
+ // NOTE: This config was released already in Spark 4.0.0, so we should not remove the
+ // support of this.
+ .withAlternative("spark.sql.execution.arrow.transformWithStateInPandas.maxRecordsPerBatch")
.intConf
.createWithDefault(10000)
@@ -6435,8 +6440,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
def arrowMaxBytesPerBatch: Long = getConf(ARROW_EXECUTION_MAX_BYTES_PER_BATCH)
- def arrowTransformWithStateInPandasMaxRecordsPerBatch: Int =
- getConf(ARROW_TRANSFORM_WITH_STATE_IN_PANDAS_MAX_RECORDS_PER_BATCH)
+ def arrowTransformWithStateInPySparkMaxRecordsPerBatch: Int =
+ getConf(ARROW_TRANSFORM_WITH_STATE_IN_PYSPARK_MAX_RECORDS_PER_BATCH)
def arrowUseLargeVarTypes: Boolean = getConf(ARROW_EXECUTION_USE_LARGE_VAR_TYPES)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/RelationalGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/RelationalGroupedDataset.scala
index 082292145e85..a4dc20d6c09e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/classic/RelationalGroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/RelationalGroupedDataset.scala
@@ -461,17 +461,24 @@ class RelationalGroupedDataset protected[sql](
Dataset.ofRows(df.sparkSession, plan)
}
- /**
- * Applies a grouped vectorized python user-defined function to each group of data.
- * The user-defined function defines a transformation: iterator of `pandas.DataFrame` ->
- * iterator of `pandas.DataFrame`.
- * For each group, all elements in the group are passed as an iterator of `pandas.DataFrame`
- * along with corresponding state, and the results for all groups are combined into a new
- * [[DataFrame]].
- *
- * This function uses Apache Arrow as serialization format between Java executors and Python
- * workers.
- */
+ // FIXME: documentation
+ private[sql] def transformWithState(
+ func: Column,
+ outputStructType: StructType,
+ outputModeStr: String,
+ timeModeStr: String,
+ initialState: RelationalGroupedDataset,
+ eventTimeColumnName: String): DataFrame = {
+ _transformWithStateInPySpark(
+ func,
+ outputStructType,
+ outputModeStr,
+ timeModeStr,
+ initialState,
+ eventTimeColumnName,
+ TransformWithStateInPySpark.UserFacingDataType.PYTHON_ROW)
+ }
+
private[sql] def transformWithStateInPandas(
func: Column,
outputStructType: StructType,
@@ -479,6 +486,24 @@ class RelationalGroupedDataset protected[sql](
timeModeStr: String,
initialState: RelationalGroupedDataset,
eventTimeColumnName: String): DataFrame = {
+ _transformWithStateInPySpark(
+ func,
+ outputStructType,
+ outputModeStr,
+ timeModeStr,
+ initialState,
+ eventTimeColumnName,
+ TransformWithStateInPySpark.UserFacingDataType.PANDAS)
+ }
+
+ private def _transformWithStateInPySpark(
+ func: Column,
+ outputStructType: StructType,
+ outputModeStr: String,
+ timeModeStr: String,
+ initialState: RelationalGroupedDataset,
+ eventTimeColumnName: String,
+ userFacingDataType: TransformWithStateInPySpark.UserFacingDataType.Value): DataFrame = {
def exprToAttr(expr: Seq[Expression]): Seq[Attribute] = {
expr.map {
case ne: NamedExpression => ne
@@ -498,12 +523,13 @@ class RelationalGroupedDataset protected[sql](
Project(groupingAttrs ++ leftChild.output, leftChild)).analyzed
val plan: LogicalPlan = if (initialState == null) {
- TransformWithStateInPandas(
+ TransformWithStateInPySpark(
func.expr,
groupingAttrs.length,
outputAttrs,
outputMode,
timeMode,
+ userFacingDataType,
child = left,
hasInitialState = false,
/* The followings are dummy variables because hasInitialState is false */
@@ -519,12 +545,13 @@ class RelationalGroupedDataset protected[sql](
val right = initialState.df.sparkSession.sessionState.executePlan(
Project(initGroupingAttrs ++ rightChild.output, rightChild)).analyzed
- TransformWithStateInPandas(
+ TransformWithStateInPySpark(
func.expr,
groupingAttributesLen = groupingAttrs.length,
outputAttrs,
outputMode,
timeMode,
+ userFacingDataType,
child = left,
hasInitialState = true,
initialState = right,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 835ecddb45c8..5be6998f08d2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -39,7 +39,7 @@ import org.apache.spark.sql.execution.command._
import org.apache.spark.sql.execution.datasources.{WriteFiles, WriteFilesExec}
import org.apache.spark.sql.execution.exchange.{REBALANCE_PARTITIONS_BY_COL, REBALANCE_PARTITIONS_BY_NONE, REPARTITION_BY_COL, REPARTITION_BY_NUM, ShuffleExchangeExec}
import org.apache.spark.sql.execution.python._
-import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPandasExec}
+import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec}
import org.apache.spark.sql.execution.streaming._
import org.apache.spark.sql.execution.streaming.sources.MemoryPlan
import org.apache.spark.sql.internal.SQLConf
@@ -792,20 +792,21 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
}
/**
- * Strategy to convert [[TransformWithStateInPandas]] logical operator to physical operator
+ * Strategy to convert [[TransformWithStateInPySpark]] logical operator to physical operator
* in streaming plans.
*/
- object TransformWithStateInPandasStrategy extends Strategy {
+ object TransformWithStateInPySparkStrategy extends Strategy {
override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
- case t @ TransformWithStateInPandas(
- func, _, outputAttrs, outputMode, timeMode, child,
+ case t @ TransformWithStateInPySpark(
+ func, _, outputAttrs, outputMode, timeMode, userFacingDataType, child,
hasInitialState, initialState, _, initialStateSchema) =>
- val execPlan = TransformWithStateInPandasExec(
+ val execPlan = TransformWithStateInPySparkExec(
func, t.leftAttributes, outputAttrs, outputMode, timeMode,
stateInfo = None,
batchTimestampMs = None,
eventTimeWatermarkForLateEvents = None,
eventTimeWatermarkForEviction = None,
+ userFacingDataType,
planLater(child),
isStreaming = true,
hasInitialState,
@@ -976,12 +977,13 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
keyEncoder, outputObjAttr, planLater(child), hasInitialState,
initialStateGroupingAttrs, initialStateDataAttrs,
initialStateDeserializer, planLater(initialState)) :: Nil
- case t @ TransformWithStateInPandas(
- func, _, outputAttrs, outputMode, timeMode, child,
+ case t @ TransformWithStateInPySpark(
+ func, _, outputAttrs, outputMode, timeMode, userFacingDataType, child,
hasInitialState, initialState, _, initialStateSchema) =>
- TransformWithStateInPandasExec.generateSparkPlanForBatchQueries(func,
- t.leftAttributes, outputAttrs, outputMode, timeMode, planLater(child), hasInitialState,
- planLater(initialState), t.rightAttributes, initialStateSchema) :: Nil
+ TransformWithStateInPySparkExec.generateSparkPlanForBatchQueries(func,
+ t.leftAttributes, outputAttrs, outputMode, timeMode, userFacingDataType,
+ planLater(child), hasInitialState, planLater(initialState), t.rightAttributes,
+ initialStateSchema) :: Nil
case _: FlatMapGroupsInPandasWithState =>
// TODO(SPARK-40443): support applyInPandasWithState in batch query
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkDeserializer.scala
similarity index 92%
rename from sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasDeserializer.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkDeserializer.scala
index b38697aeb021..25f84f0be4c2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasDeserializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkDeserializer.scala
@@ -33,12 +33,12 @@ import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, Column
/**
* A helper class to deserialize state Arrow batches from the state socket in
- * TransformWithStateInPandas.
+ * TransformWithStateInPySpark.
*/
-class TransformWithStateInPandasDeserializer(deserializer: ExpressionEncoder.Deserializer[Row])
+class TransformWithStateInPySparkDeserializer(deserializer: ExpressionEncoder.Deserializer[Row])
extends Logging {
private val allocator = ArrowUtils.rootAllocator.newChildAllocator(
- s"stdin reader for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ s"stdin reader for transformWithStateInPySpark state socket", 0, Long.MaxValue)
/**
* Read Arrow batches from the given stream and deserialize them into rows.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala
similarity index 92%
rename from sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasExec.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala
index e77035e31ccb..7cc9e4900c60 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala
@@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Expression, PythonUDF, SortOrder}
import org.apache.spark.sql.catalyst.plans.logical.ProcessingTime
+import org.apache.spark.sql.catalyst.plans.logical.TransformWithStateInPySpark
import org.apache.spark.sql.catalyst.plans.physical.Distribution
import org.apache.spark.sql.catalyst.types.DataTypeUtils
import org.apache.spark.sql.execution.{BinaryExecNode, CoGroupedIterator, SparkPlan}
@@ -46,7 +47,7 @@ import org.apache.spark.util.{CompletionIterator, SerializableConfiguration, Uti
/**
* Physical operator for executing
- * [[org.apache.spark.sql.catalyst.plans.logical.TransformWithStateInPandas]]
+ * [[org.apache.spark.sql.catalyst.plans.logical.TransformWithStateInPySpark]]
*
* @param functionExpr function called on each group
* @param groupingAttributes used to group the data
@@ -57,6 +58,7 @@ import org.apache.spark.util.{CompletionIterator, SerializableConfiguration, Uti
* @param batchTimestampMs processing timestamp of the current batch.
* @param eventTimeWatermarkForLateEvents event time watermark for filtering late events
* @param eventTimeWatermarkForEviction event time watermark for state eviction
+ * // FIXME: document `userFacingDataType`
* @param child the physical plan for the underlying data
* @param isStreaming defines whether the query is streaming or batch
* @param hasInitialState defines whether the query has initial state
@@ -64,7 +66,7 @@ import org.apache.spark.util.{CompletionIterator, SerializableConfiguration, Uti
* @param initialStateGroupingAttrs grouping attributes for initial state
* @param initialStateSchema schema for initial state
*/
-case class TransformWithStateInPandasExec(
+case class TransformWithStateInPySparkExec(
functionExpr: Expression,
groupingAttributes: Seq[Attribute],
output: Seq[Attribute],
@@ -74,6 +76,7 @@ case class TransformWithStateInPandasExec(
batchTimestampMs: Option[Long],
eventTimeWatermarkForLateEvents: Option[Long],
eventTimeWatermarkForEviction: Option[Long],
+ userFacingDataType: TransformWithStateInPySpark.UserFacingDataType.Value,
child: SparkPlan,
isStreaming: Boolean = true,
hasInitialState: Boolean,
@@ -85,7 +88,15 @@ case class TransformWithStateInPandasExec(
with WatermarkSupport
with TransformWithStateMetadataUtils {
- override def shortName: String = "transformWithStateInPandasExec"
+ // NOTE: This is needed to comply with existing release of transformWithStateInPandas.
+ override def shortName: String = if (
+ userFacingDataType == TransformWithStateInPySpark.UserFacingDataType.PANDAS
+ ) {
+ "transformWithStateInPySparkExec"
+ } else {
+ "transformWithStateInPySparkExec"
+ }
+
private val pythonUDF = functionExpr.asInstanceOf[PythonUDF]
private val pythonFunction = pythonUDF.func
private val chainedFunc =
@@ -116,8 +127,7 @@ case class TransformWithStateInPandasExec(
override def operatorStateMetadataVersion: Int = 2
override def getColFamilySchemas(
- shouldBeNullable: Boolean
- ): Map[String, StateStoreColFamilySchema] = {
+ shouldBeNullable: Boolean): Map[String, StateStoreColFamilySchema] = {
// For Python, the user can explicitly set nullability on schema, so
// we need to throw an error if the schema is nullable
driverProcessorHandle.getColumnFamilySchemas(
@@ -172,7 +182,7 @@ case class TransformWithStateInPandasExec(
batchId: Long,
stateSchemaVersion: Int): List[StateSchemaValidationResult] = {
// Start a python runner on driver, and execute pre-init UDF on the runner
- val runner = new TransformWithStateInPandasPythonPreInitRunner(
+ val runner = new TransformWithStateInPySparkPythonPreInitRunner(
pythonFunction,
"pyspark.sql.streaming.transform_with_state_driver_worker",
sessionLocalTimeZone,
@@ -186,7 +196,7 @@ case class TransformWithStateInPandasExec(
runner.process()
} catch {
case e: Throwable =>
- throw new SparkException("TransformWithStateInPandas driver worker " +
+ throw new SparkException("TransformWithStateInPySpark driver worker " +
"exited unexpectedly (crashed)", e)
}
runner.stop()
@@ -330,7 +340,7 @@ case class TransformWithStateInPandasExec(
private def initNewStateStoreAndProcessData(
partitionId: Int,
hadoopConfBroadcast: Broadcast[SerializableConfiguration])
- (f: StateStore => Iterator[InternalRow]): Iterator[InternalRow] = {
+ (f: StateStore => Iterator[InternalRow]): Iterator[InternalRow] = {
val providerId = {
val tempDirPath = Utils.createTempDir().getAbsolutePath
@@ -391,10 +401,18 @@ case class TransformWithStateInPandasExec(
val processorHandle = new StatefulProcessorHandleImpl(store, getStateInfo.queryRunId,
groupingKeyExprEncoder, timeMode, isStreaming, batchTimestampMs, metrics)
+ val evalType = {
+ if (userFacingDataType == TransformWithStateInPySpark.UserFacingDataType.PANDAS) {
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
+ } else {
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF
+ }
+ }
+
val outputIterator = if (!hasInitialState) {
- val runner = new TransformWithStateInPandasPythonRunner(
+ val runner = new TransformWithStateInPySparkPythonRunner(
chainedFunc,
- PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
+ evalType,
Array(argOffsets),
DataTypeUtils.fromAttributes(dedupAttributes),
processorHandle,
@@ -418,9 +436,17 @@ case class TransformWithStateInPandasExec(
val groupedData: Iterator[(InternalRow, Iterator[InternalRow], Iterator[InternalRow])] =
new CoGroupedIterator(data, initData, groupingAttributes)
- val runner = new TransformWithStateInPandasPythonInitialStateRunner(
+ val evalType = {
+ if (userFacingDataType == TransformWithStateInPySpark.UserFacingDataType.PANDAS) {
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF
+ } else {
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF
+ }
+ }
+
+ val runner = new TransformWithStateInPySparkPythonInitialStateRunner(
chainedFunc,
- PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
+ evalType,
Array(argOffsets ++ initArgOffsets),
DataTypeUtils.fromAttributes(dedupAttributes),
DataTypeUtils.fromAttributes(initDedupAttributes),
@@ -457,7 +483,7 @@ case class TransformWithStateInPandasExec(
}
override protected def withNewChildrenInternal(
- newLeft: SparkPlan, newRight: SparkPlan): TransformWithStateInPandasExec =
+ newLeft: SparkPlan, newRight: SparkPlan): TransformWithStateInPySparkExec =
if (hasInitialState) {
copy(child = newLeft, initialState = newRight)
} else {
@@ -470,7 +496,7 @@ case class TransformWithStateInPandasExec(
}
// scalastyle:off argcount
-object TransformWithStateInPandasExec {
+object TransformWithStateInPySparkExec {
// Plan logical transformWithStateInPandas for batch queries
def generateSparkPlanForBatchQueries(
@@ -479,6 +505,7 @@ object TransformWithStateInPandasExec {
output: Seq[Attribute],
outputMode: OutputMode,
timeMode: TimeMode,
+ userFacingDataType: TransformWithStateInPySpark.UserFacingDataType.Value,
child: SparkPlan,
hasInitialState: Boolean = false,
initialState: SparkPlan,
@@ -494,7 +521,7 @@ object TransformWithStateInPandasExec {
stateStoreCkptIds = None
)
- new TransformWithStateInPandasExec(
+ new TransformWithStateInPySparkExec(
functionExpr,
groupingAttributes,
output,
@@ -504,6 +531,7 @@ object TransformWithStateInPandasExec {
Some(System.currentTimeMillis),
None,
None,
+ userFacingDataType,
child,
isStreaming = false,
hasInitialState,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala
similarity index 88%
rename from sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasPythonRunner.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala
index 9b2a2518a7b2..595885c1ff55 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasPythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala
@@ -31,7 +31,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.python.{BasicPythonArrowOutput, PythonArrowInput, PythonUDFRunner}
-import org.apache.spark.sql.execution.python.streaming.TransformWithStateInPandasPythonRunner.{GroupedInType, InType}
+import org.apache.spark.sql.execution.python.streaming.TransformWithStateInPySparkPythonRunner.{GroupedInType, InType}
import org.apache.spark.sql.execution.streaming.{DriverStatefulProcessorHandleImpl, StatefulProcessorHandleImpl}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
@@ -39,10 +39,10 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.spark.util.ThreadUtils
/**
- * Python runner with no initial state in TransformWithStateInPandas.
+ * Python runner with no initial state in TransformWithStateInPySpark.
* Write input data as one single InternalRow in each row in arrow batch.
*/
-class TransformWithStateInPandasPythonRunner(
+class TransformWithStateInPySparkPythonRunner(
funcs: Seq[(ChainedPythonFunctions, Long)],
evalType: Int,
argOffsets: Array[Array[Int]],
@@ -55,7 +55,7 @@ class TransformWithStateInPandasPythonRunner(
groupingKeySchema: StructType,
batchTimestampMs: Option[Long],
eventTimeWatermarkForEviction: Option[Long])
- extends TransformWithStateInPandasPythonBaseRunner[InType](
+ extends TransformWithStateInPySparkPythonBaseRunner[InType](
funcs, evalType, argOffsets, _schema, processorHandle, _timeZoneId,
initialWorkerConf, pythonMetrics, jobArtifactUUID, groupingKeySchema,
batchTimestampMs, eventTimeWatermarkForEviction)
@@ -93,10 +93,10 @@ class TransformWithStateInPandasPythonRunner(
}
/**
- * Python runner with initial state in TransformWithStateInPandas.
+ * Python runner with initial state in TransformWithStateInPySpark.
* Write input data as one InternalRow(inputRow, initialState) in each row in arrow batch.
*/
-class TransformWithStateInPandasPythonInitialStateRunner(
+class TransformWithStateInPySparkPythonInitialStateRunner(
funcs: Seq[(ChainedPythonFunctions, Long)],
evalType: Int,
argOffsets: Array[Array[Int]],
@@ -110,7 +110,7 @@ class TransformWithStateInPandasPythonInitialStateRunner(
groupingKeySchema: StructType,
batchTimestampMs: Option[Long],
eventTimeWatermarkForEviction: Option[Long])
- extends TransformWithStateInPandasPythonBaseRunner[GroupedInType](
+ extends TransformWithStateInPySparkPythonBaseRunner[GroupedInType](
funcs, evalType, argOffsets, dataSchema, processorHandle, _timeZoneId,
initialWorkerConf, pythonMetrics, jobArtifactUUID, groupingKeySchema,
batchTimestampMs, eventTimeWatermarkForEviction)
@@ -159,9 +159,9 @@ class TransformWithStateInPandasPythonInitialStateRunner(
}
/**
- * Base Python runner implementation for TransformWithStateInPandas.
+ * Base Python runner implementation for TransformWithStateInPySpark.
*/
-abstract class TransformWithStateInPandasPythonBaseRunner[I](
+abstract class TransformWithStateInPySparkPythonBaseRunner[I](
funcs: Seq[(ChainedPythonFunctions, Long)],
evalType: Int,
argOffsets: Array[Array[Int]],
@@ -178,7 +178,7 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
funcs.map(_._1), evalType, argOffsets, jobArtifactUUID, pythonMetrics)
with PythonArrowInput[I]
with BasicPythonArrowOutput
- with TransformWithStateInPandasPythonRunnerUtils
+ with TransformWithStateInPySparkPythonRunnerUtils
with Logging {
protected val sqlConf = SQLConf.get
@@ -211,9 +211,9 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
val executionContext = ExecutionContext.fromExecutor(executor)
executionContext.execute(
- new TransformWithStateInPandasStateServer(stateServerSocket, processorHandle,
+ new TransformWithStateInPySparkStateServer(stateServerSocket, processorHandle,
groupingKeySchema, timeZoneId, errorOnDuplicatedFieldNames, largeVarTypes,
- sqlConf.arrowTransformWithStateInPandasMaxRecordsPerBatch,
+ sqlConf.arrowTransformWithStateInPySparkMaxRecordsPerBatch,
batchTimestampMs, eventTimeWatermarkForEviction))
context.addTaskCompletionListener[Unit] { _ =>
@@ -231,17 +231,17 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
}
/**
- * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * TransformWithStateInPySpark driver side Python runner. Similar as executor side runner,
* will start a new daemon thread on the Python runner to run state server.
*/
-class TransformWithStateInPandasPythonPreInitRunner(
+class TransformWithStateInPySparkPythonPreInitRunner(
func: PythonFunction,
workerModule: String,
timeZoneId: String,
groupingKeySchema: StructType,
processorHandleImpl: DriverStatefulProcessorHandleImpl)
extends StreamingPythonRunner(func, "", "", workerModule)
- with TransformWithStateInPandasPythonRunnerUtils
+ with TransformWithStateInPySparkPythonRunnerUtils
with Logging {
protected val sqlConf = SQLConf.get
@@ -285,13 +285,13 @@ class TransformWithStateInPandasPythonPreInitRunner(
daemonThread = new Thread {
override def run(): Unit = {
try {
- new TransformWithStateInPandasStateServer(stateServerSocket, processorHandleImpl,
+ new TransformWithStateInPySparkStateServer(stateServerSocket, processorHandleImpl,
groupingKeySchema, timeZoneId, errorOnDuplicatedFieldNames = true,
largeVarTypes = sqlConf.arrowUseLargeVarTypes,
- sqlConf.arrowTransformWithStateInPandasMaxRecordsPerBatch).run()
+ sqlConf.arrowTransformWithStateInPySparkMaxRecordsPerBatch).run()
} catch {
case e: Exception =>
- throw new SparkException("TransformWithStateInPandas state server " +
+ throw new SparkException("TransformWithStateInPySpark state server " +
"daemon thread exited unexpectedly (crashed)", e)
}
}
@@ -303,10 +303,10 @@ class TransformWithStateInPandasPythonPreInitRunner(
}
/**
- * TransformWithStateInPandas Python runner utils functions for handling a state server
+ * TransformWithStateInPySpark Python runner utils functions for handling a state server
* in a new daemon thread.
*/
-trait TransformWithStateInPandasPythonRunnerUtils extends Logging {
+trait TransformWithStateInPySparkPythonRunnerUtils extends Logging {
protected var stateServerSocketPort: Int = 0
protected var stateServerSocket: ServerSocket = null
protected def initStateServer(): Unit = {
@@ -337,7 +337,7 @@ trait TransformWithStateInPandasPythonRunnerUtils extends Logging {
}
}
-object TransformWithStateInPandasPythonRunner {
+object TransformWithStateInPySparkPythonRunner {
type InType = (InternalRow, Iterator[InternalRow])
type GroupedInType = (InternalRow, Iterator[InternalRow], Iterator[InternalRow])
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasStateServer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
similarity index 98%
rename from sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasStateServer.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
index f46b66204383..11513218ec54 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasStateServer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
@@ -43,7 +43,7 @@ import org.apache.spark.util.Utils
/**
* This class is used to handle the state requests from the Python side. It runs on a separate
- * thread spawned by TransformWithStateInPandasStateRunner per task. It opens a dedicated socket
+ * thread spawned by TransformWithStateInPySparkStateRunner per task. It opens a dedicated socket
* to process/transfer state related info which is shut down when task finishes or there's an error
* on opening the socket. It processes following state requests and return responses to the
* Python side:
@@ -51,19 +51,19 @@ import org.apache.spark.util.Utils
* - Stateful processor requests.
* - Requests for managing state variables (e.g. valueState).
*/
-class TransformWithStateInPandasStateServer(
+class TransformWithStateInPySparkStateServer(
stateServerSocket: ServerSocket,
statefulProcessorHandle: StatefulProcessorHandleImplBase,
groupingKeySchema: StructType,
timeZoneId: String,
errorOnDuplicatedFieldNames: Boolean,
largeVarTypes: Boolean,
- arrowTransformWithStateInPandasMaxRecordsPerBatch: Int,
+ arrowTransformWithStateInPySparkMaxRecordsPerBatch: Int,
batchTimestampMs: Option[Long] = None,
eventTimeWatermarkForEviction: Option[Long] = None,
outputStreamForTest: DataOutputStream = null,
valueStateMapForTest: mutable.HashMap[String, ValueStateInfo] = null,
- deserializerForTest: TransformWithStateInPandasDeserializer = null,
+ deserializerForTest: TransformWithStateInPySparkDeserializer = null,
arrowStreamWriterForTest: BaseStreamingArrowWriter = null,
listStatesMapForTest : mutable.HashMap[String, ListStateInfo] = null,
iteratorMapForTest: mutable.HashMap[String, Iterator[Row]] = null,
@@ -253,7 +253,7 @@ class TransformWithStateInPandasStateServer(
Option(statefulProcessorHandle
.asInstanceOf[StatefulProcessorHandleImpl].getExpiredTimers(expiryTimestamp))
}
- // expiryTimestampIter could be None in the TWSPandasServerSuite
+ // expiryTimestampIter could be None in the TWSPySparkServerSuite
if (!expiryTimestampIter.isDefined || !expiryTimestampIter.get.hasNext) {
// iterator is exhausted, signal the end of iterator on python client
sendResponse(1)
@@ -464,7 +464,7 @@ class TransformWithStateInPandasStateServer(
val deserializer = if (deserializerForTest != null) {
deserializerForTest
} else {
- new TransformWithStateInPandasDeserializer(listStateInfo.deserializer)
+ new TransformWithStateInPySparkDeserializer(listStateInfo.deserializer)
}
message.getMethodCase match {
case ListStateCall.MethodCase.EXISTS =>
@@ -770,21 +770,21 @@ class TransformWithStateInPandasStateServer(
val arrowSchema = ArrowUtils.toArrowSchema(outputSchema, timeZoneId,
errorOnDuplicatedFieldNames, largeVarTypes)
val allocator = ArrowUtils.rootAllocator.newChildAllocator(
- s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ s"stdout writer for transformWithStateInPySpark state socket", 0, Long.MaxValue)
val root = VectorSchemaRoot.create(arrowSchema, allocator)
val writer = new ArrowStreamWriter(root, null, outputStream)
val arrowStreamWriter = if (arrowStreamWriterForTest != null) {
arrowStreamWriterForTest
} else {
new BaseStreamingArrowWriter(root, writer,
- arrowTransformWithStateInPandasMaxRecordsPerBatch)
+ arrowTransformWithStateInPySparkMaxRecordsPerBatch)
}
// Only write a single batch in each GET request. Stops writing row if rowCount reaches
- // the arrowTransformWithStateInPandasMaxRecordsPerBatch limit. This is to handle a case
+ // the arrowTransformWithStateInPySparkMaxRecordsPerBatch limit. This is to handle a case
// when there are multiple state variables, user tries to access a different state variable
// while the current state variable is not exhausted yet.
var rowCount = 0
- while (iter.hasNext && rowCount < arrowTransformWithStateInPandasMaxRecordsPerBatch) {
+ while (iter.hasNext && rowCount < arrowTransformWithStateInPySparkMaxRecordsPerBatch) {
val data = iter.next()
val internalRow = func(data)
arrowStreamWriter.writeRow(internalRow)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
index 246057a5a9d0..be0b3a7207e0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.execution.{CommandExecutionMode, LocalLimitExec, Que
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, MergingSessionsExec, ObjectHashAggregateExec, SortAggregateExec, UpdatingSessionsExec}
import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader
import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
-import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPandasExec}
+import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec}
import org.apache.spark.sql.execution.streaming.sources.WriteToMicroBatchDataSourceV1
import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadataReader, OperatorStateMetadataV1, OperatorStateMetadataV2, OperatorStateMetadataWriter, StateSchemaBroadcast, StateSchemaMetadata}
import org.apache.spark.sql.internal.SQLConf
@@ -91,7 +91,7 @@ class IncrementalExecution(
StreamingDeduplicationStrategy ::
StreamingGlobalLimitStrategy(outputMode) ::
StreamingTransformWithStateStrategy ::
- TransformWithStateInPandasStrategy :: Nil
+ TransformWithStateInPySparkStrategy :: Nil
}
private lazy val hadoopConf = sparkSession.sessionState.newHadoopConf()
@@ -222,7 +222,7 @@ class IncrementalExecution(
// filepath, and write this path out in the OperatorStateMetadata file
case statefulOp: StatefulOperator if isFirstBatch =>
val stateSchemaVersion = statefulOp match {
- case _: TransformWithStateExec | _: TransformWithStateInPandasExec =>
+ case _: TransformWithStateExec | _: TransformWithStateInPySparkExec =>
sparkSession.sessionState.conf.
getConf(SQLConf.STREAMING_TRANSFORM_WITH_STATE_OP_STATE_SCHEMA_VERSION)
case _ => STATE_SCHEMA_DEFAULT_VERSION
@@ -278,7 +278,7 @@ class IncrementalExecution(
case exec: TransformWithStateExec =>
exec.copy(stateInfo = Some(exec.getStateInfo.copy(
stateSchemaMetadata = Some(stateSchemaBroadcast))))
- case exec: TransformWithStateInPandasExec =>
+ case exec: TransformWithStateInPySparkExec =>
exec.copy(stateInfo = Some(exec.getStateInfo.copy(
stateSchemaMetadata = Some(stateSchemaBroadcast))))
// Add other cases if needed for different StateStoreWriter implementations
@@ -372,7 +372,7 @@ class IncrementalExecution(
hasInitialState = hasInitialState
)
- case t: TransformWithStateInPandasExec =>
+ case t: TransformWithStateInPySparkExec =>
val hasInitialState = (currentBatchId == 0L && t.hasInitialState)
t.copy(
stateInfo = Some(nextStatefulOperationStateInfo()),
@@ -474,11 +474,11 @@ class IncrementalExecution(
// UpdateEventTimeColumnExec is used to tag the eventTime column, and validate
// emitted rows adhere to watermark in the output of transformWithStateInp.
- // Hence, this node shares the same watermark value as TransformWithStateInPandasExec.
+ // Hence, this node shares the same watermark value as TransformWithStateInPySparkExec.
// This is the same as above in TransformWithStateExec.
- // The only difference is TransformWithStateInPandasExec is analysed slightly different
+ // The only difference is TransformWithStateInPySparkExec is analysed slightly different
// with no SerializeFromObjectExec wrapper.
- case UpdateEventTimeColumnExec(eventTime, delay, None, t: TransformWithStateInPandasExec)
+ case UpdateEventTimeColumnExec(eventTime, delay, None, t: TransformWithStateInPySparkExec)
if t.stateInfo.isDefined =>
val stateInfo = t.stateInfo.get
val iwLateEvents = inputWatermarkForLateEvents(stateInfo)
@@ -496,7 +496,7 @@ class IncrementalExecution(
eventTimeWatermarkForEviction = inputWatermarkForEviction(t.stateInfo.get)
)
- case t: TransformWithStateInPandasExec if t.stateInfo.isDefined =>
+ case t: TransformWithStateInPySparkExec if t.stateInfo.isDefined =>
t.copy(
eventTimeWatermarkForLateEvents = inputWatermarkForLateEvents(t.stateInfo.get),
eventTimeWatermarkForEviction = inputWatermarkForEviction(t.stateInfo.get)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasStateServerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala
similarity index 93%
rename from sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasStateServerSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala
index 305a520f6af8..a67c733ef992 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPandasStateServerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala
@@ -36,7 +36,7 @@ import org.apache.spark.sql.execution.streaming.state.StateMessage.{AppendList,
import org.apache.spark.sql.streaming.{ListState, MapState, TTLConfig, ValueState}
import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
-class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with BeforeAndAfterEach {
+class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with BeforeAndAfterEach {
val stateName = "test"
val iteratorId = "testId"
val serverSocket: ServerSocket = mock(classOf[ServerSocket])
@@ -54,10 +54,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
var valueState: ValueState[Row] = _
var listState: ListState[Row] = _
var mapState: MapState[Row, Row] = _
- var stateServer: TransformWithStateInPandasStateServer = _
+ var stateServer: TransformWithStateInPySparkStateServer = _
var stateDeserializer: ExpressionEncoder.Deserializer[Row] = _
var stateSerializer: ExpressionEncoder.Serializer[Row] = _
- var transformWithStateInPandasDeserializer: TransformWithStateInPandasDeserializer = _
+ var transformWithStateInPySparkDeserializer: TransformWithStateInPySparkDeserializer = _
var arrowStreamWriter: BaseStreamingArrowWriter = _
var batchTimestampMs: Option[Long] = _
var eventTimeWatermarkForEviction: Option[Long] = _
@@ -92,16 +92,16 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val keyValueIteratorMap = mutable.HashMap[String, Iterator[(Row, Row)]](iteratorId ->
Iterator((testRow, testRow)))
listTimerMap = mutable.HashMap[String, Iterator[Long]](iteratorId -> Iterator(1L))
- transformWithStateInPandasDeserializer = mock(classOf[TransformWithStateInPandasDeserializer])
+ transformWithStateInPySparkDeserializer = mock(classOf[TransformWithStateInPySparkDeserializer])
arrowStreamWriter = mock(classOf[BaseStreamingArrowWriter])
batchTimestampMs = mock(classOf[Option[Long]])
eventTimeWatermarkForEviction = mock(classOf[Option[Long]])
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false, 2,
batchTimestampMs, eventTimeWatermarkForEviction,
- outputStream, valueStateMap, transformWithStateInPandasDeserializer, arrowStreamWriter,
+ outputStream, valueStateMap, transformWithStateInPySparkDeserializer, arrowStreamWriter,
listStateMap, iteratorMap, mapStateMap, keyValueIteratorMap, expiryTimerIter, listTimerMap)
- when(transformWithStateInPandasDeserializer.readArrowBatches(any))
+ when(transformWithStateInPySparkDeserializer.readArrowBatches(any))
.thenReturn(Seq(getIntegerRow(1)))
when(transformWithStateInPandasDeserializer.readListElements(any, any))
.thenReturn(Seq(getIntegerRow(1)))
@@ -274,10 +274,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
.setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build()
val iteratorMap = mutable.HashMap[String, Iterator[Row]](iteratorId ->
Iterator(getIntegerRow(1), getIntegerRow(2), getIntegerRow(3), getIntegerRow(4)))
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
maxRecordsPerBatch, batchTimestampMs, eventTimeWatermarkForEviction, outputStream,
- valueStateMap, transformWithStateInPandasDeserializer, arrowStreamWriter,
+ valueStateMap, transformWithStateInPySparkDeserializer, arrowStreamWriter,
listStateMap, iteratorMap)
// First call should send 2 records.
stateServer.handleListStateRequest(message)
@@ -303,10 +303,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val message = ListStateCall.newBuilder().setStateName(stateName)
.setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build()
val iteratorMap: mutable.HashMap[String, Iterator[Row]] = mutable.HashMap()
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
maxRecordsPerBatch, batchTimestampMs, eventTimeWatermarkForEviction, outputStream,
- valueStateMap, transformWithStateInPandasDeserializer, arrowStreamWriter,
+ valueStateMap, transformWithStateInPySparkDeserializer, arrowStreamWriter,
listStateMap, iteratorMap)
when(listState.get()).thenReturn(Iterator(getIntegerRow(1), getIntegerRow(2), getIntegerRow(3)))
stateServer.handleListStateRequest(message)
@@ -324,7 +324,7 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val message = ListStateCall.newBuilder().setStateName(stateName)
.setListStatePut(ListStatePut.newBuilder().build()).build()
stateServer.handleListStateRequest(message)
- verify(transformWithStateInPandasDeserializer).readListElements(any, any)
+ verify(transformWithStateInPySparkDeserializer).readListElements(any, any)
verify(listState).put(any)
}
@@ -340,7 +340,7 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val message = ListStateCall.newBuilder().setStateName(stateName)
.setAppendList(AppendList.newBuilder().build()).build()
stateServer.handleListStateRequest(message)
- verify(transformWithStateInPandasDeserializer).readListElements(any, any)
+ verify(transformWithStateInPySparkDeserializer).readListElements(any, any)
verify(listState).appendList(any)
}
@@ -395,10 +395,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val keyValueIteratorMap = mutable.HashMap[String, Iterator[(Row, Row)]](iteratorId ->
Iterator((getIntegerRow(1), getIntegerRow(1)), (getIntegerRow(2), getIntegerRow(2)),
(getIntegerRow(3), getIntegerRow(3)), (getIntegerRow(4), getIntegerRow(4))))
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
maxRecordsPerBatch, batchTimestampMs, eventTimeWatermarkForEviction, outputStream,
- valueStateMap, transformWithStateInPandasDeserializer, arrowStreamWriter,
+ valueStateMap, transformWithStateInPySparkDeserializer, arrowStreamWriter,
listStateMap, null, mapStateMap, keyValueIteratorMap)
// First call should send 2 records.
stateServer.handleMapStateRequest(message)
@@ -419,10 +419,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val message = MapStateCall.newBuilder().setStateName(stateName)
.setIterator(StateMessage.Iterator.newBuilder().setIteratorId(iteratorId).build()).build()
val keyValueIteratorMap: mutable.HashMap[String, Iterator[(Row, Row)]] = mutable.HashMap()
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
maxRecordsPerBatch, batchTimestampMs, eventTimeWatermarkForEviction,
- outputStream, valueStateMap, transformWithStateInPandasDeserializer,
+ outputStream, valueStateMap, transformWithStateInPySparkDeserializer,
arrowStreamWriter, listStateMap, null, mapStateMap, keyValueIteratorMap)
when(mapState.iterator()).thenReturn(Iterator((getIntegerRow(1), getIntegerRow(1)),
(getIntegerRow(2), getIntegerRow(2)), (getIntegerRow(3), getIntegerRow(3))))
@@ -448,10 +448,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val message = MapStateCall.newBuilder().setStateName(stateName)
.setKeys(Keys.newBuilder().setIteratorId(iteratorId).build()).build()
val iteratorMap: mutable.HashMap[String, Iterator[Row]] = mutable.HashMap()
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
maxRecordsPerBatch, batchTimestampMs, eventTimeWatermarkForEviction,
- outputStream, valueStateMap, transformWithStateInPandasDeserializer,
+ outputStream, valueStateMap, transformWithStateInPySparkDeserializer,
arrowStreamWriter, listStateMap, iteratorMap, mapStateMap)
when(mapState.keys()).thenReturn(Iterator(getIntegerRow(1), getIntegerRow(2), getIntegerRow(3)))
stateServer.handleMapStateRequest(message)
@@ -476,10 +476,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
val message = MapStateCall.newBuilder().setStateName(stateName)
.setValues(Values.newBuilder().setIteratorId(iteratorId).build()).build()
val iteratorMap: mutable.HashMap[String, Iterator[Row]] = mutable.HashMap()
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
maxRecordsPerBatch, batchTimestampMs, eventTimeWatermarkForEviction, outputStream,
- valueStateMap, transformWithStateInPandasDeserializer,
+ valueStateMap, transformWithStateInPySparkDeserializer,
arrowStreamWriter, listStateMap, iteratorMap, mapStateMap)
when(mapState.values()).thenReturn(Iterator(getIntegerRow(1), getIntegerRow(2),
getIntegerRow(3)))
@@ -572,10 +572,10 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
.setList(ListTimers.newBuilder().setIteratorId("non-exist").build())
.build()
).build()
- stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ stateServer = new TransformWithStateInPySparkStateServer(serverSocket,
statefulProcessorHandle, groupingKeySchema, "", false, false,
2, batchTimestampMs, eventTimeWatermarkForEviction, outputStream,
- valueStateMap, transformWithStateInPandasDeserializer,
+ valueStateMap, transformWithStateInPySparkDeserializer,
arrowStreamWriter, listStateMap, null, mapStateMap, null,
null, listTimerMap)
when(statefulProcessorHandle.listTimers()).thenReturn(Iterator(1))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org
(spark) 02/02: fix conflict
Posted by ka...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
kabhwan pushed a commit to branch HeartSaVioR/SPARK-XXXXX-python-transform-with-state
in repository https://gitbox.apache.org/repos/asf/spark.git
commit 5a709080dd01163451ee9f75b5eca1c2a5e83249
Author: Jungtaek Lim <ka...@gmail.com>
AuthorDate: Tue Apr 8 15:10:10 2025 +0900
fix conflict
---
.../python/streaming/TransformWithStateInPySparkStateServer.scala | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
index 11513218ec54..ee43f7b8447c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
@@ -515,7 +515,7 @@ class TransformWithStateInPySparkStateServer(
// when there are multiple state variables, user tries to access a different state variable
// while the current state variable is not exhausted yet.
var rowCount = 0
- while (iter.hasNext && rowCount < arrowTransformWithStateInPandasMaxRecordsPerBatch) {
+ while (iter.hasNext && rowCount < arrowTransformWithStateInPySparkMaxRecordsPerBatch) {
val data = iter.next()
// Serialize the value row as a byte array
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org