python/pyspark/sql/streaming/tws_tester.py - master - githubmirror/spark

python/pyspark/sql/streaming/tws_tester.py
689 строк27 KB
Dmytro Fedoriaka
[SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark
05 фев 2026, 07:08
05 фев 2026, 07:08612ade4
#
# 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.
#
from __future__ import annotations

from typing import Any, Callable, cast, Iterator, Optional, TYPE_CHECKING, Union

from pyspark.sql.streaming.stateful_processor import (
    ExpiredTimerInfo,
    ListState,
    MapState,
    StatefulProcessor,
    StatefulProcessorHandle,
    TimerValues,
    ValueState,
)
from pyspark.sql.types import Row, StructType
from pyspark.errors import PySparkValueError, PySparkAssertionError
from pyspark.errors.exceptions.base import IllegalArgumentException

if TYPE_CHECKING:
    from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike

__all__ = ["TwsTester"]


class _InMemoryTimers:
    """In-memory implementation of timers."""

    def __init__(self) -> None:
        self._key_to_timers: dict[Any, set[int]] = {}

    def registerTimer(self, grouping_key: Any, expiryTimestampMs: int) -> None:
        if grouping_key not in self._key_to_timers:
            self._key_to_timers[grouping_key] = set()
        self._key_to_timers[grouping_key].add(expiryTimestampMs)

    def deleteTimer(self, grouping_key: Any, expiryTimestampMs: int) -> None:
        if grouping_key in self._key_to_timers:
            self._key_to_timers[grouping_key].discard(expiryTimestampMs)
            if not self._key_to_timers[grouping_key]:
                del self._key_to_timers[grouping_key]

    def listTimers(self, grouping_key: Any) -> Iterator[int]:
        if grouping_key in self._key_to_timers:
            return iter(sorted(self._key_to_timers[grouping_key]))
        return iter([])

    def getAllKeysWithTimers(self) -> Iterator[Any]:
        return iter(self._key_to_timers.keys())


class _InMemoryValueState(ValueState):
    """In-memory implementation of ValueState."""

    def __init__(
        self,
        handle: _InMemoryStatefulProcessorHandle,
    ) -> None:
        self.handle = handle
        self.state: dict[Any, tuple] = dict()

    def _get_value(self) -> Optional[tuple]:
        key = self.handle.grouping_key
        return self.state.get(key, None)

    def exists(self) -> bool:
        return self._get_value() is not None

    def get(self) -> Optional[tuple]:
        return self._get_value()

    def update(self, newValue: tuple) -> None:
        key = self.handle.grouping_key
        self.state[key] = newValue

    def clear(self) -> None:
        if self.handle.grouping_key in self.state:
            del self.state[self.handle.grouping_key]


class _InMemoryListState(ListState):
    """In-memory implementation of ListState."""

    def __init__(
        self,
        handle: _InMemoryStatefulProcessorHandle,
    ) -> None:
        self.handle = handle
        self.state = dict()  # type: dict[Any, list[tuple]]

    def _get_list(self) -> Optional[list[tuple]]:
        key = self.handle.grouping_key
        return self.state.get(key, None)

    def exists(self) -> bool:
        return self._get_list() is not None

    def get(self) -> Iterator[tuple]:
        result = self._get_list()
        return iter(result if result is not None else [])

    def put(self, newState: list[tuple]) -> None:
        key = self.handle.grouping_key
        self.state[key] = newState

    def appendValue(self, newState: tuple) -> None:
        key = self.handle.grouping_key
        if not self.exists():
            self.state[key] = []
        self.state[key].append(newState)

    def appendList(self, newState: list[tuple]) -> None:
        key = self.handle.grouping_key
        if not self.exists():
            self.state[key] = []
        self.state[key].extend(newState)

    def clear(self) -> None:
        if self.handle.grouping_key in self.state:
            del self.state[self.handle.grouping_key]


class _InMemoryMapState(MapState):
    """In-memory implementation of MapState."""

    def __init__(
        self,
        handle: _InMemoryStatefulProcessorHandle,
    ) -> None:
        self.handle = handle
        self.state = dict()  # type: dict[Any, dict[tuple, tuple]]

    def _get_map(self) -> Optional[dict[tuple, tuple]]:
        key = self.handle.grouping_key
        return self.state.get(key, None)

    def exists(self) -> bool:
        return self._get_map() is not None

    def getValue(self, key: tuple) -> Optional[tuple]:
        map_state = self._get_map()
        if map_state is None:
            return None
        return map_state.get(key, None)

    def containsKey(self, key: tuple) -> bool:
        map_state = self._get_map()
        if map_state is None:
            return False
        return key in map_state

    def updateValue(self, key: tuple, value: tuple) -> None:
        grouping_key = self.handle.grouping_key
        if not self.exists():
            self.state[grouping_key] = {}
        self.state[grouping_key][key] = value

    def iterator(self) -> Iterator[tuple[tuple, tuple]]:
        map_state = self._get_map()
        if map_state is None:
            return iter([])
        return iter(map_state.items())

    def keys(self) -> Iterator[tuple]:
        map_state = self._get_map()
        if map_state is None:
            return iter([])
        return iter(map_state.keys())

    def values(self) -> Iterator[tuple]:
        map_state = self._get_map()
        if map_state is None:
            return iter([])
        return iter(map_state.values())

    def removeKey(self, key: tuple) -> None:
        map_state = self._get_map()
        if map_state is not None and key in map_state:
            del map_state[key]

    def clear(self) -> None:
        if self.handle.grouping_key in self.state:
            del self.state[self.handle.grouping_key]


class _InMemoryStatefulProcessorHandle(StatefulProcessorHandle):
    """In-memory implementation of StatefulProcessorHandle."""

    def __init__(self, timeMode: str = "None") -> None:
        self.grouping_key = None
        self.value_states: dict[Any, _InMemoryValueState] = dict()
        self.list_states: dict[Any, _InMemoryListState] = dict()
        self.map_states: dict[Any, _InMemoryMapState] = dict()
        self._timeMode = timeMode.lower()
        self._timers = _InMemoryTimers()
        self._currentWatermarkMs: int = 0

    def setGroupingKey(self, key: Any) -> None:
        self.grouping_key = key

    def setWatermark(self, watermarkMs: int) -> None:
        """Updates the current watermark. Used by TwsTester to sync watermark state."""
        self._currentWatermarkMs = watermarkMs

    def getValueState(
        self,
        stateName: str,
        schema: Union[StructType, str],
        ttlDurationMs: Optional[int] = None,
    ) -> ValueState:
        if stateName not in self.value_states:
            self.value_states[stateName] = _InMemoryValueState(self)
        return self.value_states[stateName]

    def getListState(
        self,
        stateName: str,
        schema: Union[StructType, str],
        ttlDurationMs: Optional[int] = None,
    ) -> ListState:
        if stateName not in self.list_states:
            self.list_states[stateName] = _InMemoryListState(self)
        return self.list_states[stateName]

    def getMapState(
        self,
        stateName: str,
        userKeySchema: Union[StructType, str],
        valueSchema: Union[StructType, str],
        ttlDurationMs: Optional[int] = None,
    ) -> MapState:
        if stateName not in self.map_states:
            self.map_states[stateName] = _InMemoryMapState(self)
        return self.map_states[stateName]

    def registerTimer(self, expiryTimestampMs: int) -> None:
        if self._timeMode == "none":
            raise PySparkValueError(
                errorClass="UNSUPPORTED_OPERATION",
                messageParameters={"operation": "Timers with TimeMode.None"},
            )
        if self._timeMode == "eventtime":
            if expiryTimestampMs <= self._currentWatermarkMs:
                raise IllegalArgumentException(
                    message=f"Cannot register timer with expiry {expiryTimestampMs} ms "
                    f"which is not later than the current watermark {self._currentWatermarkMs} ms."
                )
        self._timers.registerTimer(self.grouping_key, expiryTimestampMs)

    def deleteTimer(self, expiryTimestampMs: int) -> None:
        if self._timeMode == "none":
            raise PySparkValueError(
                errorClass="UNSUPPORTED_OPERATION",
                messageParameters={"operation": "Timers with TimeMode.None"},
            )
        self._timers.deleteTimer(self.grouping_key, expiryTimestampMs)

    def listTimers(self) -> Iterator[int]:
        if self._timeMode == "none":
            raise PySparkValueError(
                errorClass="UNSUPPORTED_OPERATION",
                messageParameters={"operation": "Timers with TimeMode.None"},
            )
        return self._timers.listTimers(self.grouping_key)

    def deleteIfExists(self, stateName: str) -> None:
        if stateName in self.value_states:
            del self.value_states[stateName]
        if stateName in self.list_states:
            del self.list_states[stateName]
        if stateName in self.map_states:
            del self.map_states[stateName]

    def deleteState(self, stateName: str) -> None:
        if stateName in self.value_states:
            self.value_states[stateName].clear()
        elif stateName in self.list_states:
            self.list_states[stateName].clear()
        elif stateName in self.map_states:
            self.map_states[stateName].clear()
        else:
            raise PySparkAssertionError(
                errorClass="STATE_NOT_EXISTS",
                messageParameters={},
            )


def _asRowList(x: list[Any]) -> list[Row]:
    return cast(list[Row], x)


def _asPandasDataFrame(x: list[Any]) -> "PandasDataFrameLike":
    import pandas

    if len(x) == 0:
        return pandas.DataFrame()
    else:
        return pandas.concat(cast(list[pandas.DataFrame], x), ignore_index=True)


class TwsTester:
    """
    Testing utility for transformWithState stateful processors.

    This class enables unit testing of StatefulProcessor business logic by simulating the
    behavior of transformWithState. It processes input rows and returns output rows equivalent
    to those that would be produced by the processor in an actual Spark streaming query.

    **Supported:**
        - Processing input rows and producing output rows via :meth:`test` (Row mode) or
          :meth:`testInPandas` (Pandas mode).
        - Initial state setup via constructor parameters ``initialStateRow`` or
          ``initialStatePandas``.
        - Direct state manipulation via :meth:`updateValueState`, :meth:`updateListState`,
          :meth:`updateMapState`.
        - Direct state inspection via :meth:`peekValueState`, :meth:`peekListState`,
          :meth:`peekMapState`.
        - Timers in ProcessingTime mode (use :meth:`setProcessingTime` to fire timers).
        - Timers in EventTime mode (use :meth:`setWatermark` to manually set the watermark
          and fire expired timers).
        - Late event filtering in EventTime mode.

    **Not Supported:**
        - **TTL**. States persist indefinitely, even if TTLConfig is set.
        - **Automatic watermark propagation**: In production Spark streaming, the watermark is
          computed from event times and propagated at the end of each microbatch. TwsTester does
          not simulate this behavior because it processes keys individually rather than in batches.
          To test watermark-dependent logic, use :meth:`setWatermark` to manually set the watermark
          to the desired value before calling :meth:`test`.

    **Use Cases:**
        - **Primary**: Unit testing business logic in ``handleInputRows`` implementations.
        - **Not recommended**: End-to-end testing or performance testing - use actual Spark
          streaming queries for those scenarios.

    Parameters
    ----------
    processor : :class:`StatefulProcessor`
        The stateful processor to test.
    initialStateRow : list of tuple, optional
        Initial state for each key as a list of (key, Row) tuples (for Row mode processors).
        Cannot be specified together with ``initialStatePandas``.
    initialStatePandas : list of tuple, optional
        Initial state for each key as a list of (key, DataFrame) tuples (for Pandas mode
        processors). Cannot be specified together with ``initialStateRow``.
    timeMode : str, default "None"
        Time mode for the stateful processor. Valid values are "None", "ProcessingTime",
        or "EventTime" (case-insensitive).
    outputMode : str, default "Append"
        Output mode for the stateful processor. Valid values are "Append", "Update",
        or "Complete" (case-insensitive).
    eventTimeExtractor : callable, optional
        Function to extract event time (in milliseconds) from input rows. Required when
        ``timeMode`` is "EventTime". Used for late event filtering. For Row mode, the
        function takes a Row and returns an int. For Pandas mode, the function takes
        a DataFrame row and returns an int.

    Examples
    --------
    **Example 1: Testing a Row mode processor**

    >>> processor = RunningCountStatefulProcessorFactory().row()
    >>> tester = TwsTester(processor)
    >>> result = tester.test("key1", [Row(value="a"), Row(value="b")])
    >>> # Result: [Row(key="key1", count=2)]

    **Example 2: Testing a Pandas mode processor**

    >>> processor = RunningCountStatefulProcessorFactory().pandas()
    >>> tester = TwsTester(processor)
    >>> input_df = DataFrame({"value": ["a", "b"]})
    >>> result = tester.testInPandas("key1", input_df)
    >>> # Result: DataFrame({"key": ["key1"], "count": [2]})

    **Example 3: Testing with initial state**

    >>> processor = RunningCountStatefulProcessorFactory().row()
    >>> tester = TwsTester(
    ...     processor,
    ...     initialStateRow=[("key1", Row(initial_count=10))]
    ... )
    >>> result = tester.test("key1", [Row(value="a")])
    >>> # Result: [Row(key="key1", count=11)]

    **Example 4: Direct state manipulation**

    >>> processor = RunningCountStatefulProcessorFactory().row()
    >>> tester = TwsTester(processor)
    >>> tester.updateValueState("count", "key1", (5,))
    >>> result = tester.test("key1", [Row(value="a")])
    >>> tester.peekValueState("count", "key1")
    >>> # Returns: (6,)

    .. versionadded:: 4.2.0
    """

    def __init__(
        self,
        processor: StatefulProcessor,
        initialStateRow: Optional[list[tuple[Any, Row]]] = None,
        initialStatePandas: Optional[list[tuple[Any, "PandasDataFrameLike"]]] = None,
        timeMode: str = "None",
        outputMode: str = "Append",
        eventTimeExtractor: Optional[Callable[[Any], int]] = None,
    ) -> None:
        self.processor = processor
        self._timeMode = timeMode.lower()
        self._outputMode = outputMode
        self._eventTimeExtractor = eventTimeExtractor
        self._usingPandas = False

        # Timer-related state
        self._currentProcessingTimeMs: int = 0
        self._currentWatermarkMs: int = 0

        # Create handle
        self.handle = _InMemoryStatefulProcessorHandle(timeMode=timeMode)

        assert self._timeMode in ["none", "eventtime", "processingtime"]
        if self._timeMode == "eventtime" and eventTimeExtractor is None:
            raise PySparkValueError(
                errorClass="ARGUMENT_REQUIRED",
                messageParameters={
                    "arg_name": "eventTimeExtractor",
                    "condition": "timeMode is EventTime",
                },
            )

        self.processor.init(self.handle)

        if initialStateRow is not None:
            assert initialStatePandas is None, "Cannot specify both Row and Pandas initial states."
            for key, row in initialStateRow:
                self.handle.setGroupingKey(key)
                self.processor.handleInitialState((key,), row, self._getTimerValues())
        elif initialStatePandas is not None:
            self._usingPandas = True
            for key, df in initialStatePandas:
                self.handle.setGroupingKey(key)
                self.processor.handleInitialState((key,), df, self._getTimerValues())

    def test(self, key: Any, input: list[Row]) -> list[Row]:
        """
        Processes input rows for a single key through the stateful processor.

        This method is used for testing **Row mode** processors (transformWithState).
        In EventTime mode, late events (where event time <= current watermark) are filtered out
        before reaching the processor, matching the behavior of real Spark streaming.

        The watermark is not automatically advanced based on event times. Use :meth:`setWatermark`
        to manually set the watermark before calling :meth:`test`.

        Parameters
        ----------
        key : Any
            The grouping key for the input rows.
        input : list of :class:`Row`
            Input rows to process for the given key.

        Returns
        -------
        list of :class:`Row`
            All output rows produced by the processor for this key.
        """
        self._usingPandas = False
        self.handle.setGroupingKey(key)
        filtered_input = self._filterLateEventsRow(input)
        result_iter = self.processor.handleInputRows(
            (key,), iter(filtered_input), self._getTimerValues()
        )
        result: list[Row] = _asRowList(list(result_iter))
        return result

    def testInPandas(self, key: Any, input: "PandasDataFrameLike") -> "PandasDataFrameLike":
        """
        Processes input data for a single key through the stateful processor.

        This method does the same thing as :meth:`test`, but input and output are pandas
        DataFrames. It is used for testing **Pandas mode** processors - those intended to be
        passed to ``transformWithStateInPandas`` in a real streaming query.

        In EventTime mode, late events (where event time <= current watermark) are filtered out
        before reaching the processor, matching the behavior of real Spark streaming.

        The watermark is not automatically advanced based on event times. Use :meth:`setWatermark`
        to manually set the watermark before calling :meth:`testInPandas`.

        Parameters
        ----------
        key : Any
            The grouping key for the input data.
        input : :class:`pandas.DataFrame`
            Input data to process for the given key.

        Returns
        -------
        :class:`pandas.DataFrame`
            All output data produced by the processor for this key.
        """
        self._usingPandas = True
        self.handle.setGroupingKey(key)
        filtered_input = self._filterLateEventsPandas(input)
        result_iter = self.processor.handleInputRows(
            (key,), iter([filtered_input]), self._getTimerValues()
        )
        return _asPandasDataFrame(list(result_iter))

    def updateValueState(self, stateName: str, key: Any, value: tuple) -> None:
        """Sets the value state for a given key."""
        assert stateName in self.handle.value_states, f"State {stateName} has not been initialized."
        self.handle.value_states[stateName].state[key] = value

    def peekValueState(self, stateName: str, key: Any) -> Optional[tuple]:
        """
        Peek at a value state for a given key without modifying it.

        Returns None if the state does not exist for the given key.
        """
        assert stateName in self.handle.value_states, f"State {stateName} has not been initialized."
        self.handle.setGroupingKey(key)
        return self.handle.value_states[stateName].get()

    def updateListState(self, stateName: str, key: Any, value: list[tuple]) -> None:
        """Sets the list state for a given key."""
        assert stateName in self.handle.list_states, f"State {stateName} has not been initialized."
        self.handle.list_states[stateName].state[key] = value

    def peekListState(self, stateName: str, key: Any) -> list[tuple]:
        """
        Peek at a list state for a given key without modifying it.

        Returns an empty list if the state does not exist for the given key.
        """
        assert stateName in self.handle.list_states, f"State {stateName} has not been initialized."
        self.handle.setGroupingKey(key)
        return list(self.handle.list_states[stateName].get())

    def updateMapState(self, stateName: str, key: Any, value: dict) -> None:
        """Sets the map state for a given key."""
        assert stateName in self.handle.map_states, f"State {stateName} has not been initialized."
        self.handle.map_states[stateName].state[key] = dict(value)

    def peekMapState(self, stateName: str, key: Any) -> dict:
        """
        Peek at a map state for a given key without modifying it.

        Returns an empty dict if the state does not exist for the given key.
        """
        assert stateName in self.handle.map_states, f"State {stateName} has not been initialized."
        self.handle.setGroupingKey(key)
        return dict(self.handle.map_states[stateName].iterator())

    def deleteState(self, stateName: str, key: Any) -> None:
        """Deletes state for a given key."""
        self.handle.setGroupingKey(key)
        self.handle.deleteState(stateName)

    def _getTimerValues(self) -> TimerValues:
        """Creates TimerValues based on current time state."""
        if self._timeMode == "none":
            return TimerValues(-1, -1)
        elif self._timeMode == "processingtime":
            return TimerValues(self._currentProcessingTimeMs, -1)
        else:  # EventTime
            return TimerValues(self._currentProcessingTimeMs, self._currentWatermarkMs)

    def _handleExpiredTimers(self) -> list[Row] | "PandasDataFrameLike":
        """Handles expired timers and returns emitted rows."""
        if self._timeMode == "none":
            return []

        timer_values = self._getTimerValues()
        if self._timeMode == "processingtime":
            expiry_threshold = self._currentProcessingTimeMs
        else:  # eventtime
            expiry_threshold = self._currentWatermarkMs

        result: list = []
        for key in list(self.handle._timers.getAllKeysWithTimers()):
            self.handle.setGroupingKey(key)
            expired_timers = [
                t for t in self.handle._timers.listTimers(key) if t <= expiry_threshold
            ]
            for timer_expiry_ms in expired_timers:
                expired_timer_info = ExpiredTimerInfo(timer_expiry_ms)
                output_iter = self.processor.handleExpiredTimer(
                    (key,), timer_values, expired_timer_info
                )
                result.extend(list(output_iter))
                self.handle._timers.deleteTimer(key, timer_expiry_ms)

        if self._usingPandas:
            return _asPandasDataFrame(result)
        else:
            return _asRowList(result)

    def _filterLateEventsRow(self, values: list[Row]) -> list[Row]:
        """Filters out late events based on watermark and eventTimeExtractor."""
        if self._timeMode != "eventtime" or self._eventTimeExtractor is None:
            return values
        return [v for v in values if self._eventTimeExtractor(v) > self._currentWatermarkMs]

    def _filterLateEventsPandas(self, df: "PandasDataFrameLike") -> "PandasDataFrameLike":
        """Filters out late events based on watermark and eventTimeExtractor."""
        if self._timeMode != "eventtime" or self._eventTimeExtractor is None:
            return df
        extractor = self._eventTimeExtractor
        mask = df.apply(lambda row: extractor(row) > self._currentWatermarkMs, axis=1)
        return df[mask].reset_index(drop=True)

    def setProcessingTime(self, currentTimeMs: int) -> list[Row] | "PandasDataFrameLike":
        """
        Sets the simulated processing time and fires all expired timers.

        Call this after :meth:`test` to simulate time passage and trigger any timers registered
        with ``registerTimer()``. Timers with expiry time <= current processing time will fire,
        invoking ``handleExpiredTimer`` for each. This mirrors Spark's behavior where timers
        are processed after input data within a microbatch.

        Parameters
        ----------
        currentTimeMs : int
            The processing time to set in milliseconds.

        Returns
        -------
        list of :class:`Row` or :class:`pandas.DataFrame`
            Output rows emitted by ``handleExpiredTimer`` for all fired timers.
        """
        if self._timeMode != "processingtime":
            raise PySparkValueError(
                errorClass="UNSUPPORTED_OPERATION",
                messageParameters={
                    "operation": "setProcessingTime with TimeMode other than ProcessingTime"
                },
            )
        if currentTimeMs <= self._currentProcessingTimeMs:
            raise IllegalArgumentException(
                message=f"Processing time must move forward. "
                f"Current: {self._currentProcessingTimeMs} ms, provided: {currentTimeMs} ms."
            )
        self._currentProcessingTimeMs = currentTimeMs
        return self._handleExpiredTimers()

    def setWatermark(self, currentWatermarkMs: int) -> list[Row] | "PandasDataFrameLike":
        """
        Sets the watermark and fires all expired event-time timers.

        Use this in EventTime mode to manually set the watermark. This is the only way to
        set the watermark in TwsTester, as automatic watermark propagation based on event
        times is not supported. Timers with expiry time <= new watermark will fire.

        Parameters
        ----------
        currentWatermarkMs : int
            The watermark to set in milliseconds.

        Returns
        -------
        list of :class:`Row` or :class:`pandas.DataFrame`
            Output rows emitted by ``handleExpiredTimer`` for all fired timers.
        """
        if self._timeMode != "eventtime":
            raise PySparkValueError(
                errorClass="UNSUPPORTED_OPERATION",
                messageParameters={"operation": "setWatermark with TimeMode other than EventTime"},
            )
        if currentWatermarkMs <= self._currentWatermarkMs:
            raise IllegalArgumentException(
                message=f"Watermark must move forward. "
                f"Current: {self._currentWatermarkMs} ms, provided: {currentWatermarkMs} ms."
            )
        self._currentWatermarkMs = currentWatermarkMs
        self.handle.setWatermark(currentWatermarkMs)
        return self._handleExpiredTimers()