Integrate with before/after update triggers

Author: gwaramadze

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,

quixstreams/dataframe/windows/time_based.py

@@ -10,6 +10,8 @@
     MultiAggregationWindowMixin,
     SingleAggregationWindowMixin,
     Window,
+    WindowAfterUpdateCallback,
+    WindowBeforeUpdateCallback,
     WindowKeyResult,
     WindowOnLateCallback,
     get_window_ranges,
@@ -30,6 +32,8 @@ def __init__(
         dataframe: "StreamingDataFrame",
         step_ms: Optional[int] = None,
         on_late: Optional[WindowOnLateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ):
         super().__init__(
             name=name,
@@ -40,6 +44,8 @@ def __init__(
         self._grace_ms = grace_ms
         self._step_ms = step_ms
         self._on_late = on_late
+        self._before_update = before_update
+        self._after_update = after_update
 
     def final(self) -> "StreamingDataFrame":
         """
@@ -69,13 +75,17 @@ def on_update(
             _headers: Any,
             transaction: WindowedPartitionTransaction,
         ):
-            self.process_window(
+            # Process the window and get windows triggered from callbacks
+            _, triggered_windows = self.process_window(
                 value=value,
                 key=key,
                 timestamp_ms=timestamp_ms,
+                headers=_headers,
                 transaction=transaction,
             )
-            return []
+            # Yield triggered windows (from before_update/after_update callbacks)
+            for key, window in triggered_windows:
+                yield window, key, window["start"], None
 
         def on_watermark(
             _value: Any,
@@ -133,15 +143,20 @@ def on_update(
             _headers: Any,
             transaction: WindowedPartitionTransaction,
         ):
-            updated_windows = self.process_window(
+            # Process the window and get both updated and triggered windows
+            updated_windows, triggered_windows = self.process_window(
                 value=value,
                 key=key,
                 timestamp_ms=timestamp_ms,
+                headers=_headers,
                 transaction=transaction,
             )
             # Use window start timestamp as a new record timestamp
+            # Yield both updated and triggered windows
             for key, window in updated_windows:
                 yield window, key, window["start"], None
+            for key, window in triggered_windows:
+                yield window, key, window["start"], None
 
         def on_watermark(
             _value: Any,
@@ -169,11 +184,22 @@ def process_window(
         value: Any,
         key: Any,
         timestamp_ms: int,
+        headers: Any,
         transaction: WindowedPartitionTransaction,
-    ) -> Iterable[WindowKeyResult]:
+    ) -> 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
+        """
         state = transaction.as_state(prefix=key)
         duration_ms = self._duration_ms
         grace_ms = self._grace_ms
+        before_update = self._before_update
+        after_update = self._after_update
 
         collect = self.collect
         aggregate = self.aggregate
@@ -190,6 +216,7 @@ def process_window(
         max_expired_window_end = latest_timestamp - grace_ms
         max_expired_window_start = max_expired_window_end - duration_ms
         updated_windows: list[WindowKeyResult] = []
+        triggered_windows: list[WindowKeyResult] = []
         for start, end in ranges:
             if start <= max_expired_window_start:
                 late_by_ms = max_expired_window_end - timestamp_ms
@@ -207,18 +234,78 @@ def process_window(
             # since actual values are stored separately and combined into an array
             # during window expiration.
             aggregated = None
+
             if aggregate:
                 current_value = state.get_window(start, end)
                 if current_value is None:
                     current_value = self._initialize_value()
 
+                # Check before_update trigger
+                if before_update and before_update(
+                    current_value, value, key, timestamp_ms, headers
+                ):
+                    # Get collected values for the result
+                    # Do NOT include the current value - before_update means
+                    # we expire BEFORE adding the current value
+                    collected = state.get_from_collection(start, end) if collect else []
+
+                    result = self._results(current_value, collected, start, end)
+                    triggered_windows.append((key, result))
+                    transaction.delete_window(start, end, prefix=key)
+                    # Note: We don't delete from collection here - normal expiration
+                    # will handle cleanup for both tumbling and hopping windows
+                    continue
+
                 aggregated = self._aggregate_value(current_value, value, timestamp_ms)
-                updated_windows.append(
-                    (
-                        key,
-                        self._results(aggregated, [], start, end),
-                    )
-                )
+
+                # Check after_update trigger
+                if after_update and after_update(
+                    aggregated, value, key, timestamp_ms, headers
+                ):
+                    # Get collected values for the result
+                    collected = []
+                    if collect:
+                        collected = state.get_from_collection(start, end)
+                        # Add the current value that's being collected
+                        collected.append(self._collect_value(value))
+
+                    result = self._results(aggregated, collected, start, end)
+                    triggered_windows.append((key, result))
+                    transaction.delete_window(start, end, prefix=key)
+                    # Note: We don't delete from collection here - normal expiration
+                    # will handle cleanup for both tumbling and hopping windows
+                    continue
+
+                result = self._results(aggregated, [], start, end)
+                updated_windows.append((key, result))
+            elif collect and (before_update or after_update):
+                # For collect-only windows, get the old collected values
+                old_collected = state.get_from_collection(start, end)
+
+                # Check before_update trigger (before adding new value)
+                if before_update and before_update(
+                    old_collected, value, key, timestamp_ms, headers
+                ):
+                    # Expire with the current collection (WITHOUT the new value)
+                    result = self._results(None, old_collected, start, end)
+                    triggered_windows.append((key, result))
+                    transaction.delete_window(start, end, prefix=key)
+                    # Note: We don't delete from collection here - normal expiration
+                    # will handle cleanup for both tumbling and hopping windows
+                    continue
+
+                # Check after_update trigger (conceptually after adding new value)
+                # For collect, "after update" means after the value would be added
+                if after_update:
+                    new_collected = [*old_collected, self._collect_value(value)]
+                    if after_update(new_collected, value, key, timestamp_ms, headers):
+                        result = self._results(None, new_collected, start, end)
+                        triggered_windows.append((key, result))
+                        transaction.delete_window(start, end, prefix=key)
+                        # Note: We don't delete from collection here - normal expiration
+                        # will handle cleanup for both tumbling and hopping windows
+                        continue
+
             state.update_window(start, end, value=aggregated, timestamp_ms=timestamp_ms)
 
         if collect:
@@ -227,7 +314,7 @@ def process_window(
                 id=timestamp_ms,
             )
 
-        return updated_windows
+        return updated_windows, triggered_windows
 
     def expire_by_partition(
         self,