Refactor _as_wall_clock wrapper

Author: gwaramadze

quixstreams/dataframe/utils.py

@@ -1,4 +1,4 @@
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Union
 
 
@@ -22,3 +22,7 @@ def ensure_milliseconds(delta: Union[int, timedelta]) -> int:
             f'Timedelta must be either "int" representing milliseconds '
             f'or "datetime.timedelta", got "{type(delta)}"'
         )
+
+
+def now() -> int:
+    return int(datetime.now().timestamp() * 1000)

quixstreams/dataframe/windows/base.py

@@ -21,7 +21,6 @@
     Stream,
     TransformExpandedCallback,
     TransformFunction,
-    TransformWallClockExpandedCallback,
 )
 from quixstreams.core.stream.exceptions import InvalidOperation
 from quixstreams.models.topics.manager import TopicManager
@@ -47,7 +46,7 @@
     Iterable[Message],
 ]
 
-WallClockCallback = Callable[[int, WindowedPartitionTransaction], Iterable[Message]]
+WallClockCallback = Callable[[WindowedPartitionTransaction], Iterable[Message]]
 
 
 class Window(abc.ABC):
@@ -79,7 +78,6 @@ def process_window(
     @abstractmethod
     def process_wall_clock(
         self,
-        timestamp_ms: int,
         transaction: WindowedPartitionTransaction,
     ) -> Iterable[WindowKeyResult]:
         pass
@@ -122,7 +120,7 @@ def _apply_window(
             func=windowed_func, expand=True
         )
         wall_clock_stream = Stream(
-            TransformFunction(wall_clock_transform_func, expand=True, wall_clock=True)
+            func=TransformFunction(wall_clock_transform_func, expand=True)
         )
         sdf = self._dataframe.__dataframe_clone__(stream=windowed_stream)
         return sdf.concat_wall_clock(wall_clock_stream)
@@ -169,10 +167,10 @@ def window_callback(
                 yield (window, key, window["start"], None)
 
         def wall_clock_callback(
-            timestamp: int, transaction: WindowedPartitionTransaction
+            transaction: WindowedPartitionTransaction,
         ) -> Iterable[Message]:
             # TODO: Check if this will work for sliding windows
-            for key, window in self.process_wall_clock(timestamp, transaction):
+            for key, window in self.process_wall_clock(transaction):
                 yield (window, key, window["start"], None)
 
         return self._apply_window(
@@ -225,7 +223,7 @@ def window_callback(
                 yield (window, key, window["start"], None)
 
         def wall_clock_callback(
-            timestamp: int, transaction: WindowedPartitionTransaction
+            transaction: WindowedPartitionTransaction,
         ) -> Iterable[Message]:
             # TODO: Implement wall_clock callback
             return []
@@ -475,17 +473,19 @@ def _as_wall_clock(
     processing_context: "ProcessingContext",
     store_name: str,
     stream_id: str,
-) -> TransformWallClockExpandedCallback:
+) -> TransformExpandedCallback:
     @functools.wraps(func)
-    def wrapper(timestamp: int) -> Iterable[Message]:
+    def wrapper(
+        value: Any, key: Any, timestamp: int, headers: Any
+    ) -> Iterable[Message]:
         ctx = message_context()
         transaction = cast(
             WindowedPartitionTransaction,
             processing_context.checkpoint.get_store_transaction(
                 stream_id=stream_id, partition=ctx.partition, store_name=store_name
             ),
         )
-        return func(timestamp, transaction)
+        return func(transaction)
 
     return wrapper
 

quixstreams/dataframe/windows/count_based.py

@@ -191,7 +191,6 @@ def process_window(
 
     def process_wall_clock(
         self,
-        timestamp_ms: int,
         transaction: WindowedPartitionTransaction,
     ) -> Iterable[WindowKeyResult]:
         return []

quixstreams/dataframe/windows/time_based.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Any, Iterable, Literal, Optional
 
 from quixstreams.context import message_context
+from quixstreams.dataframe.utils import now
 from quixstreams.state import WindowedPartitionTransaction, WindowedState
 
 from .base import (
@@ -202,16 +203,12 @@ def process_window(
 
     def process_wall_clock(
         self,
-        timestamp_ms: int,
         transaction: WindowedPartitionTransaction,
     ) -> Iterable[WindowKeyResult]:
-        latest_expired_window_end = transaction.get_latest_expired(prefix=b"")
-        latest_timestamp = max(timestamp_ms, latest_expired_window_end)
-        max_expired_window_end = latest_timestamp - self._grace_ms
         return self.expire_by_partition(
-            transaction,
-            max_expired_window_end,
-            self.collect,
+            transaction=transaction,
+            max_expired_end=now() - self._grace_ms,
+            collect=self.collect,
             advance_last_expired_timestamp=False,
         )