Integrate with before/after update triggers

Author: gwaramadze

docs/windowing.md

@@ -594,7 +594,6 @@ if __name__ == '__main__':
 
 
 ### Early window expiration with triggers
-!!! info New in v3.24.0
 
 To expire windows before their natural expiration time based on custom conditions, you can pass `before_update` or `after_update` callbacks to `.tumbling_window()` and `.hopping_window()` methods.
 

quixstreams/app.py

@@ -45,7 +45,7 @@
 )
 from .platforms.quix.env import QUIX_ENVIRONMENT
 from .processing import ProcessingContext
-from .processing.watermarking import WatermarkManager
+from .processing.watermarking import WatermarkManager, WatermarkMessage
 from .runtracker import RunTracker
 from .sinks import SinkManager
 from .sources import BaseSource, SourceException, SourceManager
@@ -1008,7 +1008,9 @@ def _process_message(self, dataframe_composed: dict[str, VoidExecutor]):
         )
 
         if topic_name == self._watermark_manager.watermarks_topic.name:
-            watermark = self._watermark_manager.receive(message=first_row.value)
+            watermark = self._watermark_manager.receive(
+                message=cast(WatermarkMessage, first_row.value)
+            )
             if watermark is None:
                 return
 
@@ -1073,12 +1075,12 @@ def _process_message(self, dataframe_composed: dict[str, VoidExecutor]):
 
         # Store the message offset after it's successfully processed
         self._processing_context.store_offset(
-            topic=topic_name, partition=partition, offset=offset
+            topic=topic_name, partition=partition, offset=offset or 0
         )
         self._run_tracker.set_message_consumed(True)
 
         if self._on_message_processed is not None:
-            self._on_message_processed(topic_name, partition, offset)
+            self._on_message_processed(topic_name, partition, offset or 0)
 
     def _on_assign(self, _, topic_partitions: List[TopicPartition]):
         """
@@ -1104,6 +1106,7 @@ def _on_assign(self, _, topic_partitions: List[TopicPartition]):
             )
             for i in range(
                 self._watermark_manager.watermarks_topic.broker_config.num_partitions
+                or 1
             )
         ]
         # TODO: The set is used because the watermark tp can already be present in the "topic_partitions"

quixstreams/core/stream/functions/apply.py

@@ -48,7 +48,6 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
                 is_watermark: bool = False,
-                on_watermark=self.on_watermark,
             ) -> None:
                 # Execute a function on a single value and wrap results into a list
                 # to expand them downstream

quixstreams/core/stream/functions/base.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any
+from typing import Any, Optional
 
 from quixstreams.utils.pickle import pickle_copier
 
@@ -18,9 +18,11 @@ class StreamFunction(abc.ABC):
 
     expand: bool = False
 
-    def __init__(self, func: StreamCallback):
+    def __init__(
+        self, func: StreamCallback, on_watermark: Optional[StreamCallback] = None
+    ):
         self.func = func
-        self.on_watermark = None
+        self.on_watermark = on_watermark
 
     @abc.abstractmethod
     def get_executor(self, *child_executors: VoidExecutor) -> VoidExecutor:

quixstreams/core/stream/functions/transform.py

@@ -65,12 +65,14 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
                 is_watermark: bool = False,
-                on_watermark=self.on_watermark,
             ):
                 if is_watermark:
-                    if on_watermark is not None:
+                    if self.on_watermark is not None:
                         # React on the new watermark if "on_watermark" is defined
-                        result = self.on_watermark(None, None, timestamp, ())
+                        watermark_func = cast(
+                            TransformExpandedCallback, self.on_watermark
+                        )
+                        result = watermark_func(None, None, timestamp, ())
                         for new_value, new_key, new_timestamp, new_headers in result:
                             child_executor(
                                 new_value,
@@ -102,13 +104,13 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
                 is_watermark: bool = False,
-                on_watermark=self.on_watermark,
             ):
                 if is_watermark:
-                    if on_watermark is not None:
+                    if self.on_watermark is not None:
                         # React on the new watermark if "on_watermark" is defined
-                        new_value, new_key, new_timestamp, new_headers = (
-                            self.on_watermark(None, None, timestamp, ())
+                        watermark_func = cast(TransformCallback, self.on_watermark)
+                        new_value, new_key, new_timestamp, new_headers = watermark_func(
+                            None, None, timestamp, ()
                         )
                         child_executor(
                             new_value,

quixstreams/dataframe/dataframe.py

@@ -1708,7 +1708,7 @@ def _sink_callback(
                 headers=headers,
                 partition=ctx.partition,
                 topic=ctx.topic,
-                offset=ctx.offset,
+                offset=ctx.offset or 0,
             )
 
         # uses apply without returning to make this operation terminal

quixstreams/dataframe/windows/base.py

@@ -33,6 +33,8 @@
 WindowResult: TypeAlias = dict[str, Any]
 WindowKeyResult: TypeAlias = tuple[Any, WindowResult]
 Message: TypeAlias = tuple[WindowResult, Any, int, Any]
+WindowBeforeUpdateCallback: TypeAlias = Callable[[Any, Any, Any, int, Any], bool]
+WindowAfterUpdateCallback: TypeAlias = Callable[[Any, Any, Any, int, Any], bool]
 
 WindowAggregateFunc = Callable[[Any, Any], Any]
 
@@ -58,6 +60,25 @@ def __init__(
     def name(self) -> str:
         return self._name
 
+    @abstractmethod
+    def process_window(
+        self,
+        value: Any,
+        key: Any,
+        timestamp_ms: int,
+        headers: Any,
+        transaction: WindowedPartitionTransaction,
+    ) -> tuple[Iterable[WindowKeyResult], Iterable[WindowKeyResult]]:
+        """
+        Process a window update for the given value and key.
+
+        Returns:
+            A tuple of (updated_windows, triggered_windows) where:
+            - updated_windows: Windows that were updated but not expired
+            - triggered_windows: Windows that were expired early due to before_update/after_update callbacks
+        """
+        pass
+
     def register_store(self) -> None:
         TopicManager.ensure_topics_copartitioned(*self._dataframe.topics)
         # Create a config for the changelog topic based on the underlying SDF topics
@@ -126,6 +147,7 @@ def final(self) -> "StreamingDataFrame":
         If some message keys appear irregularly in the stream, the latest windows
         can remain unprocessed until the message the same key is received.
         """
+        ...
 
     @abstractmethod
     def current(self) -> "StreamingDataFrame":

quixstreams/dataframe/windows/count_based.py

@@ -90,6 +90,7 @@ def window_callback(
                 value=value,
                 key=key,
                 timestamp_ms=timestamp_ms,
+                headers=_headers,
                 transaction=transaction,
             )
             # Use window start timestamp as a new record timestamp
@@ -135,6 +136,7 @@ def window_callback(
                 value=value,
                 key=key,
                 timestamp_ms=timestamp_ms,
+                headers=_headers,
                 transaction=transaction,
             )
 

quixstreams/dataframe/windows/sliding.py

@@ -18,7 +18,7 @@ def process_window(
         timestamp_ms: int,
         headers: Any,
         transaction: WindowedPartitionTransaction,
-    ) -> Iterable[WindowKeyResult]:
+    ) -> tuple[Iterable[WindowKeyResult], Iterable[WindowKeyResult]]:
         """
         The algorithm is based on the concept that each message
         is associated with a left and a right window.
@@ -87,7 +87,7 @@ def process_window(
                 timestamp_ms=timestamp_ms,
                 late_by_ms=max_expired_window_end - timestamp_ms,
             )
-            return []
+            return [], []
 
         right_start = timestamp_ms + 1
         right_end = right_start + duration
@@ -256,7 +256,9 @@ def process_window(
         if collect:
             state.add_to_collection(value=self._collect_value(value), id=timestamp_ms)
 
-        return reversed(updated_windows)
+        # Sliding windows don't support before_update/after_update callbacks yet,
+        # so triggered_windows is always empty
+        return reversed(updated_windows), []
 
     def expire_by_partition(
         self,