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

python/pyspark/sql/streaming/stateful_processor_api_client.py
649 строк28 KB
Kousuke Saruta
[SPARK-57760][SS][PYTHON][FOLLOWUP] Fix mypy union-attr error in _normalize_state_value
10 июл 2026, 01:57
10 июл 2026, 01:5713fbf14
#
# 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 datetime import datetime
from enum import Enum
import json
import os
import socket
from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast

from pyspark.serializers import write_int, read_int, UTF8Deserializer
from pyspark.sql.pandas.serializers import ArrowStreamSerializer
from pyspark.sql.types import (
    StructType,
    Row,
)
from pyspark.sql.pandas.types import convert_pandas_using_numpy_type
from pyspark.serializers import PickleSerializer
from pyspark.errors import PySparkRuntimeError
import uuid

__all__ = ["StatefulProcessorApiClient", "StatefulProcessorHandleState"]

# None means not yet checked; True/False after _load_numpy() is called.
has_numpy: Optional[bool] = None
np = None

SCALAR_TYPES = (bool, int, float, str, bytes, datetime, type(None))


def _load_numpy() -> None:
    """Lazily resolve numpy availability without importing it at module load time.

    Importing numpy at the top level would slow down ``import pyspark``
    (see test_import_spark_libraries).
    """
    global has_numpy, np
    try:
        import numpy

        np = numpy
        has_numpy = True
    except ImportError:
        has_numpy = False


def _normalize_state_value(v: Any) -> Any:
    if type(v) in SCALAR_TYPES:  # Fast path for common scalar values.
        return v
    # Convert NumPy scalar values to Python primitive values.
    if np is not None and isinstance(v, np.generic):
        return v.tolist()
    # Named tuples (collections.namedtuple or typing.NamedTuple) and Row both
    # require positional arguments and cannot be instantiated with a generator expression.
    if isinstance(v, Row) or (isinstance(v, tuple) and hasattr(v, "_fields")):
        return type(v)(*map(_normalize_state_value, v))
    # List / tuple: recursively normalize each element.
    if isinstance(v, (list, tuple)):
        return type(v)(map(_normalize_state_value, v))
    # Dict: normalize both keys and values.
    if isinstance(v, dict):
        return {_normalize_state_value(k): _normalize_state_value(val) for k, val in v.items()}
    # Address a couple of pandas dtypes too.
    if hasattr(v, "to_pytimedelta"):
        return v.to_pytimedelta()
    if hasattr(v, "to_pydatetime"):
        return v.to_pydatetime()
    return v


class StatefulProcessorHandleState(Enum):
    PRE_INIT = 0
    CREATED = 1
    INITIALIZED = 2
    DATA_PROCESSED = 3
    TIMER_PROCESSED = 4
    CLOSED = 5


class StatefulProcessorApiClient:
    def __init__(
        self, state_server_port: Union[int, str], key_schema: StructType, is_driver: bool = False
    ) -> None:
        self.key_schema = key_schema
        if isinstance(state_server_port, str):
            self._client_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            self._client_socket.connect(state_server_port)
        else:
            self._client_socket = socket.socket()
            self._client_socket.connect(("localhost", state_server_port))

            # SPARK-51667: We have a pattern of sending messages continuously from one side
            # (Python -> JVM, and vice versa) before getting response from other side. Since most
            # messages we are sending are small, this triggers the bad combination of Nagle's
            # algorithm and delayed ACKs, which can cause a significant delay on the latency.
            # See SPARK-51667 for more details on how this can be a problem.
            #
            # Disabling either would work, but it's more common to disable Nagle's algorithm; there
            # is lot less reference to disabling delayed ACKs, while there are lots of resources to
            # disable Nagle's algorithm.
            self._client_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

        self.sockfile = self._client_socket.makefile(
            "rwb", int(os.environ.get("SPARK_BUFFER_SIZE", 65536))
        )
        if is_driver:
            self.handle_state = StatefulProcessorHandleState.PRE_INIT
        else:
            self.handle_state = StatefulProcessorHandleState.CREATED
        self.utf8_deserializer = UTF8Deserializer()
        self.pickleSer = PickleSerializer()
        self.serializer = ArrowStreamSerializer()
        # Dictionaries to store the mapping between iterator id and a tuple of data batch
        # and the index of the last row that was read.
        self.list_timer_iterator_cursors: Dict[str, Tuple[Any, int, bool]] = {}
        self.expiry_timer_iterator_cursors: Dict[str, Tuple[Any, int, bool]] = {}

        # statefulProcessorApiClient is initialized per batch per partition,
        # so we will have new timestamps for a new batch
        self._batch_timestamp = -1
        self._watermark_timestamp = -1

    def set_handle_state(self, state: StatefulProcessorHandleState) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        if state == StatefulProcessorHandleState.PRE_INIT:
            proto_state = stateMessage.PRE_INIT
        elif state == StatefulProcessorHandleState.CREATED:
            proto_state = stateMessage.CREATED
        elif state == StatefulProcessorHandleState.INITIALIZED:
            proto_state = stateMessage.INITIALIZED
        elif state == StatefulProcessorHandleState.DATA_PROCESSED:
            proto_state = stateMessage.DATA_PROCESSED
        elif state == StatefulProcessorHandleState.TIMER_PROCESSED:
            proto_state = stateMessage.TIMER_PROCESSED
        else:
            proto_state = stateMessage.CLOSED
        set_handle_state = stateMessage.SetHandleState(state=proto_state)
        handle_call = stateMessage.StatefulProcessorCall(setHandleState=set_handle_state)
        message = stateMessage.StateRequest(statefulProcessorCall=handle_call)

        self._send_proto_message(message.SerializeToString())

        response_message = self._receive_proto_message()
        status = response_message[0]
        if status == 0:
            self.handle_state = state
        else:
            # TODO(SPARK-49233): Classify errors thrown by internal methods.
            raise PySparkRuntimeError(f"Error setting handle state: {response_message[1]}")

    def set_implicit_key(self, key: Tuple) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        key_bytes = self._serialize_to_bytes(self.key_schema, key)
        set_implicit_key = stateMessage.SetImplicitKey(key=key_bytes)
        request = stateMessage.ImplicitGroupingKeyRequest(setImplicitKey=set_implicit_key)
        message = stateMessage.StateRequest(implicitGroupingKeyRequest=request)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify errors thrown by internal methods.
            raise PySparkRuntimeError(f"Error setting implicit key: {response_message[1]}")

    def remove_implicit_key(self) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        remove_implicit_key = stateMessage.RemoveImplicitKey()
        request = stateMessage.ImplicitGroupingKeyRequest(removeImplicitKey=remove_implicit_key)
        message = stateMessage.StateRequest(implicitGroupingKeyRequest=request)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify errors thrown by internal methods.
            raise PySparkRuntimeError(f"Error removing implicit key: {response_message[1]}")

    def get_value_state(
        self, state_name: str, schema: Union[StructType, str], ttl_duration_ms: Optional[int]
    ) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        if isinstance(schema, str):
            schema = self._parse_string_schema(schema)

        state_call_command = stateMessage.StateCallCommand()
        state_call_command.stateName = state_name
        state_call_command.schema = schema.json()
        if ttl_duration_ms is not None:
            state_call_command.ttl.durationMs = ttl_duration_ms
        call = stateMessage.StatefulProcessorCall(getValueState=state_call_command)
        message = stateMessage.StateRequest(statefulProcessorCall=call)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error initializing value state: {response_message[1]}")

    def get_list_state(
        self, state_name: str, schema: Union[StructType, str], ttl_duration_ms: Optional[int]
    ) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        if isinstance(schema, str):
            schema = self._parse_string_schema(schema)

        state_call_command = stateMessage.StateCallCommand()
        state_call_command.stateName = state_name
        state_call_command.schema = schema.json()
        if ttl_duration_ms is not None:
            state_call_command.ttl.durationMs = ttl_duration_ms
        call = stateMessage.StatefulProcessorCall(getListState=state_call_command)
        message = stateMessage.StateRequest(statefulProcessorCall=call)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error initializing list state: {response_message[1]}")

    def register_timer(self, expiry_time_stamp_ms: int) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        register_call = stateMessage.RegisterTimer(expiryTimestampMs=expiry_time_stamp_ms)
        state_call_command = stateMessage.TimerStateCallCommand(register=register_call)
        call = stateMessage.StatefulProcessorCall(timerStateCall=state_call_command)
        message = stateMessage.StateRequest(statefulProcessorCall=call)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error register timer: {response_message[1]}")

    def delete_timer(self, expiry_time_stamp_ms: int) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        delete_call = stateMessage.DeleteTimer(expiryTimestampMs=expiry_time_stamp_ms)
        state_call_command = stateMessage.TimerStateCallCommand(delete=delete_call)
        call = stateMessage.StatefulProcessorCall(timerStateCall=state_call_command)
        message = stateMessage.StateRequest(statefulProcessorCall=call)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error deleting timer: {response_message[1]}")

    def get_list_timer_row(self, iterator_id: str) -> Tuple[int, bool]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        if iterator_id in self.list_timer_iterator_cursors:
            # if the iterator is already in the dictionary, return the next row
            data_batch, index, require_next_fetch = self.list_timer_iterator_cursors[iterator_id]
        else:
            list_call = stateMessage.ListTimers(iteratorId=iterator_id)
            state_call_command = stateMessage.TimerStateCallCommand(list=list_call)
            call = stateMessage.StatefulProcessorCall(timerStateCall=state_call_command)
            message = stateMessage.StateRequest(statefulProcessorCall=call)

            self._send_proto_message(message.SerializeToString())
            response_message = self._receive_proto_message_with_timers()
            status = response_message[0]
            if status == 0:
                data_batch = list(map(lambda x: x.timestampMs, response_message[2]))
                require_next_fetch = response_message[3]
                index = 0
            else:
                raise StopIteration()

        is_last_row = False
        new_index = index + 1
        if new_index < len(data_batch):
            # Update the index in the dictionary.
            self.list_timer_iterator_cursors[iterator_id] = (
                data_batch,
                new_index,
                require_next_fetch,
            )
        else:
            # If the index is at the end of the data batch, remove the state from the dictionary.
            self.list_timer_iterator_cursors.pop(iterator_id, None)
            is_last_row = True

        is_last_row_from_iterator = is_last_row and not require_next_fetch
        timestamp = data_batch[index]
        return (timestamp, is_last_row_from_iterator)

    def get_expiry_timers_iterator(
        self, iterator_id: str, expiry_timestamp: int
    ) -> Tuple[Tuple, int, bool]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        if iterator_id in self.expiry_timer_iterator_cursors:
            # If the state is already in the dictionary, return the next row.
            data_batch, index, require_next_fetch = self.expiry_timer_iterator_cursors[iterator_id]
        else:
            expiry_timer_call = stateMessage.ExpiryTimerRequest(
                expiryTimestampMs=expiry_timestamp, iteratorId=iterator_id
            )
            timer_request = stateMessage.TimerRequest(expiryTimerRequest=expiry_timer_call)
            message = stateMessage.StateRequest(timerRequest=timer_request)

            self._send_proto_message(message.SerializeToString())
            response_message = self._receive_proto_message_with_timers()
            status = response_message[0]
            if status == 0:
                data_batch = list(
                    map(
                        lambda x: (self._deserialize_from_bytes(x.key), x.timestampMs),
                        response_message[2],
                    )
                )
                require_next_fetch = response_message[3]
                index = 0
            else:
                raise StopIteration()

        is_last_row = False
        new_index = index + 1
        if new_index < len(data_batch):
            # Update the index in the dictionary.
            self.expiry_timer_iterator_cursors[iterator_id] = (
                data_batch,
                new_index,
                require_next_fetch,
            )
        else:
            # If the index is at the end of the data batch, remove the state from the dictionary.
            self.expiry_timer_iterator_cursors.pop(iterator_id, None)
            is_last_row = True

        is_last_row_from_iterator = is_last_row and not require_next_fetch
        key, timestamp = data_batch[index]
        return (key, timestamp, is_last_row_from_iterator)

    def get_timestamps(self, time_mode: str) -> Tuple[int, int]:
        if time_mode.lower() == "none":
            return -1, -1
        else:
            if self._batch_timestamp == -1:
                self._batch_timestamp = self._get_batch_timestamp()
            if self._watermark_timestamp == -1:
                self._watermark_timestamp = self._get_watermark_timestamp()
        return self._batch_timestamp, self._watermark_timestamp

    def get_map_state(
        self,
        state_name: str,
        user_key_schema: Union[StructType, str],
        value_schema: Union[StructType, str],
        ttl_duration_ms: Optional[int],
    ) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        if isinstance(user_key_schema, str):
            user_key_schema = self._parse_string_schema(user_key_schema)
        if isinstance(value_schema, str):
            value_schema = self._parse_string_schema(value_schema)

        state_call_command = stateMessage.StateCallCommand()
        state_call_command.stateName = state_name
        state_call_command.schema = user_key_schema.json()
        state_call_command.mapStateValueSchema = value_schema.json()
        if ttl_duration_ms is not None:
            state_call_command.ttl.durationMs = ttl_duration_ms
        call = stateMessage.StatefulProcessorCall(getMapState=state_call_command)
        message = stateMessage.StateRequest(statefulProcessorCall=call)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error initializing map state: {response_message[1]}")

    def delete_if_exists(self, state_name: str) -> None:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        state_call_command = stateMessage.StateCallCommand()
        state_call_command.stateName = state_name
        call = stateMessage.StatefulProcessorCall(deleteIfExists=state_call_command)
        message = stateMessage.StateRequest(statefulProcessorCall=call)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error deleting state: {response_message[1]}")

    def _get_batch_timestamp(self) -> int:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        get_processing_time_call = stateMessage.GetProcessingTime()
        timer_value_call = stateMessage.TimerValueRequest(
            getProcessingTimer=get_processing_time_call
        )
        timer_request = stateMessage.TimerRequest(timerValueRequest=timer_value_call)
        message = stateMessage.StateRequest(timerRequest=timer_request)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message_with_long_value()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error getting processing timestamp: {response_message[1]}")
        else:
            timestamp = response_message[2]
            return timestamp

    def _get_watermark_timestamp(self) -> int:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        get_watermark_call = stateMessage.GetWatermark()
        timer_value_call = stateMessage.TimerValueRequest(getWatermark=get_watermark_call)
        timer_request = stateMessage.TimerRequest(timerValueRequest=timer_value_call)
        message = stateMessage.StateRequest(timerRequest=timer_request)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message_with_long_value()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error getting eventtime timestamp: {response_message[1]}")
        else:
            timestamp = response_message[2]
            return timestamp

    def _send_proto_message(self, message: bytes) -> None:
        # Writing zero here to indicate message version. This allows us to evolve the message
        # format or even changing the message protocol in the future.
        write_int(0, self.sockfile)
        write_int(len(message), self.sockfile)
        self.sockfile.write(message)
        self.sockfile.flush()

    def _receive_proto_message(self) -> Tuple[int, str, bytes]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponse()
        message.ParseFromString(bytes)
        return message.statusCode, message.errorMessage, message.value

    def _receive_proto_message_with_long_value(self) -> Tuple[int, str, int]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponseWithLongTypeVal()
        message.ParseFromString(bytes)
        return message.statusCode, message.errorMessage, message.value

    def _receive_proto_message_with_string_value(self) -> Tuple[int, str, str]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponseWithStringTypeVal()
        message.ParseFromString(bytes)
        return message.statusCode, message.errorMessage, message.value

    # The third return type is RepeatedScalarFieldContainer[bytes], which is protobuf's container
    # type. We simplify it to Any here to avoid unnecessary complexity.
    def _receive_proto_message_with_list_get(self) -> Tuple[int, str, Any, bool]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponseWithListGet()
        message.ParseFromString(bytes)

        return message.statusCode, message.errorMessage, message.value, message.requireNextFetch

    # The third return type is RepeatedScalarFieldContainer[bytes], which is protobuf's container
    # type. We simplify it to Any here to avoid unnecessary complexity.
    def _receive_proto_message_with_map_keys_values(self) -> Tuple[int, str, Any, bool]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponseWithMapKeysOrValues()
        message.ParseFromString(bytes)

        return message.statusCode, message.errorMessage, message.value, message.requireNextFetch

    # The third return type is RepeatedScalarFieldContainer[KeyAndValuePair], which is protobuf's
    # container type. We simplify it to Any here to avoid unnecessary complexity.
    def _receive_proto_message_with_map_pairs(self) -> Tuple[int, str, Any, bool]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponseWithMapIterator()
        message.ParseFromString(bytes)

        return message.statusCode, message.errorMessage, message.kvPair, message.requireNextFetch

    # The third return type is RepeatedScalarFieldContainer[TimerInfo], which is protobuf's
    # container type. We simplify it to Any here to avoid unnecessary complexity.
    def _receive_proto_message_with_timers(self) -> Tuple[int, str, Any, bool]:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        length = read_int(self.sockfile)
        bytes = self.sockfile.read(length)
        message = stateMessage.StateResponseWithTimer()
        message.ParseFromString(bytes)

        return message.statusCode, message.errorMessage, message.timer, message.requireNextFetch

    def _receive_str(self) -> str:
        return self.utf8_deserializer.loads(self.sockfile)

    def _serialize_to_bytes(self, schema: StructType, data: Tuple) -> bytes:
        if has_numpy is None:
            _load_numpy()
        if has_numpy:
            converted = tuple(map(_normalize_state_value, data))
        else:
            converted = data
        return self.pickleSer.dumps(schema.toInternal(converted))

    def _deserialize_from_bytes(self, value: bytes) -> Any:
        return self.pickleSer.loads(value)

    def _send_arrow_state(self, schema: StructType, state: List[Tuple]) -> None:
        import pyarrow as pa
        import pandas as pd

        column_names = [field.name for field in schema.fields]
        pandas_df = convert_pandas_using_numpy_type(
            pd.DataFrame(state, columns=column_names), schema
        )
        batch = pa.RecordBatch.from_pandas(pandas_df)
        self.serializer.dump_stream([batch], cast(IO[bytes], self.sockfile))
        self.sockfile.flush()

    def _read_arrow_state(self) -> Any:
        return self.serializer.load_stream(cast(IO[bytes], self.sockfile))

    def _send_list_state(self, schema: StructType, state: List[Tuple]) -> None:
        for value in state:
            bytes = self._serialize_to_bytes(schema, value)
            length = len(bytes)
            write_int(length, self.sockfile)
            self.sockfile.write(bytes)

        write_int(-1, self.sockfile)
        self.sockfile.flush()

    def _read_list_state(self) -> List[Any]:
        data_array = []
        while True:
            length = read_int(self.sockfile)
            if length < 0:
                break
            bytes = self.sockfile.read(length)
            data_array.append(self._deserialize_from_bytes(bytes))
        return data_array

    # Parse a string schema into a StructType schema. This method will perform an API call to
    # JVM side to parse the schema string.
    def _parse_string_schema(self, schema: str) -> StructType:
        import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage

        parse_string_schema_call = stateMessage.ParseStringSchema(schema=schema)
        utils_request = stateMessage.UtilsRequest(parseStringSchema=parse_string_schema_call)
        message = stateMessage.StateRequest(utilsRequest=utils_request)

        self._send_proto_message(message.SerializeToString())
        response_message = self._receive_proto_message_with_string_value()
        status = response_message[0]
        if status != 0:
            # TODO(SPARK-49233): Classify user facing errors.
            raise PySparkRuntimeError(f"Error parsing string schema: {response_message[1]}")
        else:
            return StructType.fromJson(json.loads(response_message[2]))


class ListTimerIterator:
    def __init__(self, stateful_processor_api_client: StatefulProcessorApiClient):
        # Generate a unique identifier for the iterator to make sure iterators on the
        # same partition won't interfere with each other
        self.iterator_id = str(uuid.uuid4())
        self.stateful_processor_api_client = stateful_processor_api_client
        self.iterator_fully_consumed = False

    def __iter__(self) -> Iterator[int]:
        return self

    def __next__(self) -> int:
        if self.iterator_fully_consumed:
            raise StopIteration()

        ts, is_last_row = self.stateful_processor_api_client.get_list_timer_row(self.iterator_id)
        if is_last_row:
            self.iterator_fully_consumed = True

        return ts


class ExpiredTimerIterator:
    def __init__(
        self, stateful_processor_api_client: StatefulProcessorApiClient, expiry_timestamp: int
    ):
        # Generate a unique identifier for the iterator to make sure iterators on the
        # same partition won't interfere with each other
        self.iterator_id = str(uuid.uuid4())
        self.stateful_processor_api_client = stateful_processor_api_client
        self.expiry_timestamp = expiry_timestamp
        self.iterator_fully_consumed = False

    def __iter__(self) -> Iterator[Tuple[Tuple, int]]:
        return self

    def __next__(self) -> Tuple[Tuple, int]:
        if self.iterator_fully_consumed:
            raise StopIteration()

        key, ts, is_last_row = self.stateful_processor_api_client.get_expiry_timers_iterator(
            self.iterator_id, self.expiry_timestamp
        )
        if is_last_row:
            self.iterator_fully_consumed = True

        return (key, ts)