[SPARK-52333][SS][PYTHON] Squeeze the protocol of retrieving timers for transformWithState in PySpark

### What changes were proposed in this pull request?

This PR proposes to squeeze the protocol of retrieving timers for transformWithState in PySpark, which will help a lot on dealing with not-to-be-huge number of timers.

Here are the changes:

* StatefulProcessorHandleImpl.listTimers(), StatefulProcessorHandleImpl.getExpiredTimers() no longer requires additional request to notice there is no further data to read.
  * We inline the data into proto message, to ease of determine whether the iterator has fully consumed or not.

This change is the same mechanism we applied for ListState & MapState. We got performance improvement in the prior case, and we also see this change to be helpful on our internal benchmark.

### Why are the changes needed?

To optimize further on some timer operations.

We benchmarked the change with listing 100 timers (PR for benchmarking: apache#50952), and we saw overall performance improvements.

> Before the fix

```

 ==================== SET IMPLICIT KEY latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
78.250		141.583		184.375		635.792		962743.500
 ==================== REGISTER latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
65.375		126.125		162.792		565.833		60809.333
 ==================== DELETE latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
68.500		130.000		170.292		610.083		156733.125
 ==================== LIST latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
486.833		714.961		998.625		2695.417		167039.959

 ==================== SET IMPLICIT KEY latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
77.916		139.000		182.375		671.792		521809.958
 ==================== REGISTER latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
65.000		124.333		160.875		596.667		30860.208
 ==================== DELETE latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
67.125		127.916		170.250		740.051		64404.416
 ==================== LIST latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
482.041		710.333		1050.333		2685.500		76762.583

 ==================== SET IMPLICIT KEY latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
78.208		139.959		181.459		722.459		713788.250
 ==================== REGISTER latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
65.209		125.125		159.625		636.666		27963.167
 ==================== DELETE latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
67.417		129.000		168.875		764.602		12991.667
 ==================== LIST latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
479.000		709.584		1045.543		2776.541		92247.542
```

> After the fix

```
 ==================== SET IMPLICIT KEY latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
31.250		47.250		75.875		150.000		551557.750
 ==================== REGISTER latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
26.958		39.208		65.208		122.667		78609.292
 ==================== DELETE latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
23.500		41.125		64.542		125.958		52641.042
 ==================== LIST latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
93.125		118.542		156.500		284.625		19910.000

 ==================== SET IMPLICIT KEY latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
30.875		44.083		70.417		128.875		628912.209
 ==================== REGISTER latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
26.917		36.416		61.292		109.917		164584.666
 ==================== DELETE latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
23.333		38.375		59.542		113.839		114350.250
 ==================== LIST latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
94.125		115.208		148.917		246.292		36924.292

 ==================== SET IMPLICIT KEY latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
31.375		58.375		93.041		243.750		719545.583
 ==================== REGISTER latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
26.959		50.167		81.833		194.375		67609.583
 ==================== DELETE latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
24.208		50.834		83.000		211.018		20611.959
 ==================== LIST latency (micros) ======================
perc:50		perc:95		perc:99		perc:99.9		perc:100
95.291		132.375		183.875		427.584		36971.792
```

Worth noting that it is not only impacting the LIST operation - it also impacts other operations as well. It's not clear why it happens, but the direction of reducing round-trips is proven to be the right direction.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Existing UT.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes apache#51036 from HeartSaVioR/SPARK-52333.

Authored-by: Jungtaek Lim <kabhwan.opensource@gmail.com>
Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

Author: HeartSaVioR

python/pyspark/sql/streaming/proto/StateMessage_pb2.py

@@ -366,6 +366,74 @@ class StateResponseWithMapIterator(google.protobuf.message.Message):
 
 global___StateResponseWithMapIterator = StateResponseWithMapIterator
 
+class TimerInfo(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    KEY_FIELD_NUMBER: builtins.int
+    TIMESTAMPMS_FIELD_NUMBER: builtins.int
+    key: builtins.bytes
+    timestampMs: builtins.int
+    def __init__(
+        self,
+        *,
+        key: builtins.bytes | None = ...,
+        timestampMs: builtins.int = ...,
+    ) -> None: ...
+    def HasField(
+        self, field_name: typing_extensions.Literal["_key", b"_key", "key", b"key"]
+    ) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing_extensions.Literal[
+            "_key", b"_key", "key", b"key", "timestampMs", b"timestampMs"
+        ],
+    ) -> None: ...
+    def WhichOneof(
+        self, oneof_group: typing_extensions.Literal["_key", b"_key"]
+    ) -> typing_extensions.Literal["key"] | None: ...
+
+global___TimerInfo = TimerInfo
+
+class StateResponseWithTimer(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    STATUSCODE_FIELD_NUMBER: builtins.int
+    ERRORMESSAGE_FIELD_NUMBER: builtins.int
+    TIMER_FIELD_NUMBER: builtins.int
+    REQUIRENEXTFETCH_FIELD_NUMBER: builtins.int
+    statusCode: builtins.int
+    errorMessage: builtins.str
+    @property
+    def timer(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+        global___TimerInfo
+    ]: ...
+    requireNextFetch: builtins.bool
+    def __init__(
+        self,
+        *,
+        statusCode: builtins.int = ...,
+        errorMessage: builtins.str = ...,
+        timer: collections.abc.Iterable[global___TimerInfo] | None = ...,
+        requireNextFetch: builtins.bool = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing_extensions.Literal[
+            "errorMessage",
+            b"errorMessage",
+            "requireNextFetch",
+            b"requireNextFetch",
+            "statusCode",
+            b"statusCode",
+            "timer",
+            b"timer",
+        ],
+    ) -> None: ...
+
+global___StateResponseWithTimer = StateResponseWithTimer
+
 class StatefulProcessorCall(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
@@ -634,15 +702,21 @@ global___TimerValueRequest = TimerValueRequest
 class ExpiryTimerRequest(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
+    ITERATORID_FIELD_NUMBER: builtins.int
     EXPIRYTIMESTAMPMS_FIELD_NUMBER: builtins.int
+    iteratorId: builtins.str
     expiryTimestampMs: builtins.int
     def __init__(
         self,
         *,
+        iteratorId: builtins.str = ...,
         expiryTimestampMs: builtins.int = ...,
     ) -> None: ...
     def ClearField(
-        self, field_name: typing_extensions.Literal["expiryTimestampMs", b"expiryTimestampMs"]
+        self,
+        field_name: typing_extensions.Literal[
+            "expiryTimestampMs", b"expiryTimestampMs", "iteratorId", b"iteratorId"
+        ],
     ) -> None: ...
 
 global___ExpiryTimerRequest = ExpiryTimerRequest

python/pyspark/sql/streaming/proto/StateMessage_pb2.pyi

@@ -24,17 +24,13 @@
 from pyspark.sql.pandas.serializers import ArrowStreamSerializer
 from pyspark.sql.types import (
     StructType,
-    TYPE_CHECKING,
     Row,
 )
 from pyspark.sql.pandas.types import convert_pandas_using_numpy_type
 from pyspark.serializers import CPickleSerializer
 from pyspark.errors import PySparkRuntimeError
 import uuid
 
-if TYPE_CHECKING:
-    from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike
-
 __all__ = ["StatefulProcessorApiClient", "StatefulProcessorHandleState"]
 
 
@@ -80,9 +76,11 @@ def __init__(
         self.utf8_deserializer = UTF8Deserializer()
         self.pickleSer = CPickleSerializer()
         self.serializer = ArrowStreamSerializer()
-        # Dictionaries to store the mapping between iterator id and a tuple of pandas DataFrame
+        # 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["PandasDataFrameLike", int]] = {}
+        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
@@ -222,76 +220,93 @@ def delete_timer(self, expiry_time_stamp_ms: int) -> None:
             # TODO(SPARK-49233): Classify user facing errors.
             raise PySparkRuntimeError(f"Error deleting timer: " f"{response_message[1]}")
 
-    def get_list_timer_row(self, iterator_id: str) -> int:
+    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
-            pandas_df, index = self.list_timer_iterator_cursors[iterator_id]
+            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()
+            response_message = self._receive_proto_message_with_timers()
             status = response_message[0]
             if status == 0:
-                iterator = self._read_arrow_state()
-                # We need to exhaust the iterator here to make sure all the arrow batches are read,
-                # even though there is only one batch in the iterator. Otherwise, the stream might
-                # block further reads since it thinks there might still be some arrow batches left.
-                # We only need to read the first batch in the iterator because it's guaranteed that
-                # there would only be one batch sent from the JVM side.
-                data_batch = None
-                for batch in iterator:
-                    if data_batch is None:
-                        data_batch = batch
-                if data_batch is None:
-                    # TODO(SPARK-49233): Classify user facing errors.
-                    raise PySparkRuntimeError("Error getting map state entry.")
-                pandas_df = data_batch.to_pandas()
+                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(pandas_df):
+        if new_index < len(data_batch):
             # Update the index in the dictionary.
-            self.list_timer_iterator_cursors[iterator_id] = (pandas_df, new_index)
+            self.list_timer_iterator_cursors[iterator_id] = (
+                data_batch,
+                new_index,
+                require_next_fetch,
+            )
         else:
-            # If the index is at the end of the DataFrame, remove the state from the dictionary.
+            # 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)
-        return pandas_df.at[index, "timestamp"].item()
+            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, expiry_timestamp: int
-    ) -> Iterator[list[Tuple[Tuple, int]]]:
+        self, iterator_id: str, expiry_timestamp: int
+    ) -> Tuple[Tuple, int, bool]:
         import pyspark.sql.streaming.proto.StateMessage_pb2 as stateMessage
 
-        while True:
-            expiry_timer_call = stateMessage.ExpiryTimerRequest(expiryTimestampMs=expiry_timestamp)
+        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()
+            response_message = self._receive_proto_message_with_timers()
             status = response_message[0]
-            if status == 1:
-                break
-            elif status == 0:
-                result_list = []
-                iterator = self._read_arrow_state()
-                for batch in iterator:
-                    batch_df = batch.to_pandas()
-                    for i in range(batch.num_rows):
-                        deserialized_key = self.pickleSer.loads(batch_df.at[i, "key"])
-                        timestamp = batch_df.at[i, "timestamp"].item()
-                        result_list.append((tuple(deserialized_key), timestamp))
-                yield result_list
+            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:
-                # TODO(SPARK-49233): Classify user facing errors.
-                raise PySparkRuntimeError(f"Error getting expiry timers: " f"{response_message[1]}")
+                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":
@@ -461,6 +476,18 @@ def _receive_proto_message_with_map_pairs(self) -> Tuple[int, str, Any, bool]:
 
         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)
 
@@ -552,9 +579,44 @@ def __init__(self, stateful_processor_api_client: StatefulProcessorApiClient):
         # 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:
-        return self.stateful_processor_api_client.get_list_timer_row(self.iterator_id)
+        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)

python/pyspark/sql/streaming/stateful_processor_api_client.py

@@ -28,6 +28,7 @@
     StatefulProcessorHandle,
     TimerValues,
 )
+from pyspark.sql.streaming.stateful_processor_api_client import ExpiredTimerIterator
 from pyspark.sql.types import Row
 
 if TYPE_CHECKING:
@@ -218,24 +219,19 @@ def _handle_expired_timers(
         )
 
         if self._time_mode.lower() == "processingtime":
-            expiry_list_iter = stateful_processor_api_client.get_expiry_timers_iterator(
-                batch_timestamp
-            )
+            expiry_iter = ExpiredTimerIterator(stateful_processor_api_client, batch_timestamp)
         elif self._time_mode.lower() == "eventtime":
-            expiry_list_iter = stateful_processor_api_client.get_expiry_timers_iterator(
-                watermark_timestamp
-            )
+            expiry_iter = ExpiredTimerIterator(stateful_processor_api_client, watermark_timestamp)
         else:
-            expiry_list_iter = iter([[]])
+            expiry_iter = iter([])  # type: ignore[assignment]
 
         # process with expiry timers, only timer related rows will be emitted
-        for expiry_list in expiry_list_iter:
-            for key_obj, expiry_timestamp in expiry_list:
-                stateful_processor_api_client.set_implicit_key(key_obj)
-                for pd in self._stateful_processor.handleExpiredTimer(
-                    key=key_obj,
-                    timerValues=TimerValues(batch_timestamp, watermark_timestamp),
-                    expiredTimerInfo=ExpiredTimerInfo(expiry_timestamp),
-                ):
-                    yield pd
-                stateful_processor_api_client.delete_timer(expiry_timestamp)
+        for key_obj, expiry_timestamp in expiry_iter:
+            stateful_processor_api_client.set_implicit_key(key_obj)
+            for pd in self._stateful_processor.handleExpiredTimer(
+                key=key_obj,
+                timerValues=TimerValues(batch_timestamp, watermark_timestamp),
+                expiredTimerInfo=ExpiredTimerInfo(expiry_timestamp),
+            ):
+                yield pd
+            stateful_processor_api_client.delete_timer(expiry_timestamp)