Split to before_update and after_update

gwaramadze · gwaramadze · commit 4becc2258c7b · 2025-10-17T14:02:34.000+02:00
diff --git a/docs/windowing.md b/docs/windowing.md
@@ -596,19 +596,20 @@ 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 the `on_update` callback to `.tumbling_window()` and `.hopping_window()` methods.
+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.
 
 This is useful when you want to emit results as soon as certain conditions are met, rather than waiting for the window to close naturally.
 
 **How it works**:
 
-- The `on_update` callback is invoked every time a window is updated with a new value.
-- It receives `old_value` and `new_value` - the raw aggregated values before and after the update.
-- For `collect()` operations, it receives lists of collected values.
-- If the callback returns `True`, the window is immediately expired and emitted downstream.
+- The `before_update` callback is invoked before the window aggregation is updated with a new value.
+- The `after_update` callback is invoked after the window aggregation has been updated with a new value.
+- Both callbacks receive: `aggregated` (current or updated aggregated value), `value` (incoming value), `key`, `timestamp`, and `headers`.
+- For `collect()` operations without aggregation, `aggregated` contains the list of collected values.
+- If either callback returns `True`, the window is immediately expired and emitted downstream.
 - The expired window is removed from state, but may be "resurrected" if new data arrives within its time range before natural expiration.
 
-**Example**:
+**Example with after_update**:
 
 ```python
 from typing import Any
@@ -620,16 +621,18 @@ app = Application(...)
 sdf = app.dataframe(...)
 
 
-def trigger_on_threshold(old_value: int, new_value: int) -> bool:
+def trigger_on_threshold(
+    aggregated: int, value: Any, key: Any, timestamp: int, headers: Any
+) -> bool:
     """
     Expire the window early when the sum exceeds 1000.
     """
-    return new_value > 1000
+    return aggregated > 1000
 
 
 # Define a 1-hour tumbling window with early expiration trigger
 sdf = (
-    sdf.tumbling_window(timedelta(hours=1), on_update=trigger_on_threshold)
+    sdf.tumbling_window(timedelta(hours=1), after_update=trigger_on_threshold)
     .sum()
     .final()
 )
@@ -640,6 +643,25 @@ if __name__ == '__main__':
 
 ```
 
+**Example with before_update**:
+
+```python
+def trigger_before_large_value(
+    aggregated: int, value: Any, key: Any, timestamp: int, headers: Any
+) -> bool:
+    """
+    Expire the window before adding a value if it would make the sum too large.
+    """
+    return (aggregated + value) > 1000
+
+
+sdf = (
+    sdf.tumbling_window(timedelta(hours=1), before_update=trigger_before_large_value)
+    .sum()
+    .final()
+)
+```
+
 
 ## Emitting results
 
diff --git a/quixstreams/dataframe/dataframe.py b/quixstreams/dataframe/dataframe.py
@@ -72,7 +72,11 @@
     TumblingCountWindowDefinition,
     TumblingTimeWindowDefinition,
 )
-from .windows.base import WindowOnLateCallback, WindowOnUpdateCallback
+from .windows.base import (
+    WindowAfterUpdateCallback,
+    WindowBeforeUpdateCallback,
+    WindowOnLateCallback,
+)
 
 if typing.TYPE_CHECKING:
     from quixstreams.processing import ProcessingContext
@@ -1085,7 +1089,8 @@ def tumbling_window(
         grace_ms: Union[int, timedelta] = 0,
         name: Optional[str] = None,
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ) -> TumblingTimeWindowDefinition:
         """
         Create a time-based tumbling window transformation on this StreamingDataFrame.
@@ -1152,12 +1157,18 @@ def tumbling_window(
             (default behavior).
             Otherwise, no message will be logged.
 
-        :param on_update: an optional callback to trigger early window expiration based
-            on custom conditions.
-            The callback receives `old_value` and `new_value` (the raw aggregated values
-            before and after the update). If it returns `True`, the window will be expired
-            immediately, even if it hasn't reached its natural expiration time.
-            For `collect()` operations, the callback receives lists of collected values.
+        :param before_update: an optional callback to trigger early window expiration
+            before the window is updated.
+            The callback receives `aggregated` (current aggregated value or default/None),
+            `value`, `key`, `timestamp`, and `headers`.
+            If it returns `True`, the window will be expired immediately.
+            Default - `None`.
+
+        :param after_update: an optional callback to trigger early window expiration
+            after the window is updated.
+            The callback receives `aggregated` (updated aggregated value), `value`, `key`,
+            `timestamp`, and `headers`.
+            If it returns `True`, the window will be expired immediately.
             Default - `None`.
 
         :return: `TumblingTimeWindowDefinition` instance representing the tumbling window
@@ -1175,7 +1186,8 @@ def tumbling_window(
             dataframe=self,
             name=name,
             on_late=on_late,
-            on_update=on_update,
+            before_update=before_update,
+            after_update=after_update,
         )
 
     def tumbling_count_window(
@@ -1235,7 +1247,8 @@ def hopping_window(
         grace_ms: Union[int, timedelta] = 0,
         name: Optional[str] = None,
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ) -> HoppingTimeWindowDefinition:
         """
         Create a time-based hopping window transformation on this StreamingDataFrame.
@@ -1313,12 +1326,18 @@ def hopping_window(
             (default behavior).
             Otherwise, no message will be logged.
 
-        :param on_update: an optional callback to trigger early window expiration based
-            on custom conditions.
-            The callback receives `old_value` and `new_value` (the raw aggregated values
-            before and after the update). If it returns `True`, the window will be expired
-            immediately, even if it hasn't reached its natural expiration time.
-            For `collect()` operations, the callback receives lists of collected values.
+        :param before_update: an optional callback to trigger early window expiration
+            before the window is updated.
+            The callback receives `aggregated` (current aggregated value or default/None),
+            `value`, `key`, `timestamp`, and `headers`.
+            If it returns `True`, the window will be expired immediately.
+            Default - `None`.
+
+        :param after_update: an optional callback to trigger early window expiration
+            after the window is updated.
+            The callback receives `aggregated` (updated aggregated value), `value`, `key`,
+            `timestamp`, and `headers`.
+            If it returns `True`, the window will be expired immediately.
             Default - `None`.
 
         :return: `HoppingTimeWindowDefinition` instance representing the hopping
@@ -1338,7 +1357,8 @@ def hopping_window(
             dataframe=self,
             name=name,
             on_late=on_late,
-            on_update=on_update,
+            before_update=before_update,
+            after_update=after_update,
         )
 
     def hopping_count_window(
diff --git a/quixstreams/dataframe/windows/base.py b/quixstreams/dataframe/windows/base.py
@@ -34,7 +34,8 @@
 WindowResult: TypeAlias = dict[str, Any]
 WindowKeyResult: TypeAlias = tuple[Any, WindowResult]
 Message: TypeAlias = tuple[WindowResult, Any, int, Any]
-WindowOnUpdateCallback: TypeAlias = Callable[[Any, Any], bool]
+WindowBeforeUpdateCallback: TypeAlias = Callable[[Any, Any, Any, int, Any], bool]
+WindowAfterUpdateCallback: TypeAlias = Callable[[Any, Any, Any, int, Any], bool]
 
 WindowAggregateFunc = Callable[[Any, Any], Any]
 
@@ -66,6 +67,7 @@ def process_window(
         value: Any,
         key: Any,
         timestamp_ms: int,
+        headers: Any,
         transaction: WindowedPartitionTransaction,
     ) -> tuple[Iterable[WindowKeyResult], Iterable[WindowKeyResult]]:
         pass
@@ -135,6 +137,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
@@ -177,7 +180,11 @@ def window_callback(
             transaction: WindowedPartitionTransaction,
         ) -> Iterable[Message]:
             updated_windows, expired_windows = self.process_window(
-                value=value, key=key, timestamp_ms=timestamp_ms, transaction=transaction
+                value=value,
+                key=key,
+                timestamp_ms=timestamp_ms,
+                headers=_headers,
+                transaction=transaction,
             )
 
             # loop over the expired_windows generator to ensure the windows
diff --git a/quixstreams/dataframe/windows/count_based.py b/quixstreams/dataframe/windows/count_based.py
@@ -58,6 +58,7 @@ def process_window(
         value: Any,
         key: Any,
         timestamp_ms: int,
+        headers: Any,
         transaction: WindowedPartitionTransaction[str, CountWindowsData],
     ) -> tuple[Iterable[WindowKeyResult], Iterable[WindowKeyResult]]:
         """
diff --git a/quixstreams/dataframe/windows/definitions.py b/quixstreams/dataframe/windows/definitions.py
@@ -15,8 +15,9 @@
 )
 from .base import (
     Window,
+    WindowAfterUpdateCallback,
+    WindowBeforeUpdateCallback,
     WindowOnLateCallback,
-    WindowOnUpdateCallback,
 )
 from .count_based import (
     CountWindow,
@@ -55,13 +56,15 @@ def __init__(
         name: Optional[str],
         dataframe: "StreamingDataFrame",
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ) -> None:
         super().__init__()
 
         self._name = name
         self._on_late = on_late
-        self._on_update = on_update
+        self._before_update = before_update
+        self._after_update = after_update
         self._dataframe = dataframe
 
     @abstractmethod
@@ -242,7 +245,8 @@ def __init__(
         name: Optional[str] = None,
         step_ms: Optional[int] = None,
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ):
         if not isinstance(duration_ms, int):
             raise TypeError("Window size must be an integer")
@@ -257,7 +261,7 @@ def __init__(
                 f"got {step_ms}ms"
             )
 
-        super().__init__(name, dataframe, on_late, on_update)
+        super().__init__(name, dataframe, on_late, before_update, after_update)
 
         self._duration_ms = duration_ms
         self._grace_ms = grace_ms
@@ -285,7 +289,8 @@ def __init__(
         dataframe: "StreamingDataFrame",
         name: Optional[str] = None,
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ):
         super().__init__(
             duration_ms=duration_ms,
@@ -294,7 +299,8 @@ def __init__(
             name=name,
             step_ms=step_ms,
             on_late=on_late,
-            on_update=on_update,
+            before_update=before_update,
+            after_update=after_update,
         )
 
     def _get_name(self, func_name: Optional[str]) -> str:
@@ -326,7 +332,8 @@ def _create_window(
             aggregators=aggregators or {},
             collectors=collectors or {},
             on_late=self._on_late,
-            on_update=self._on_update,
+            before_update=self._before_update,
+            after_update=self._after_update,
         )
 
 
@@ -338,16 +345,18 @@ def __init__(
         dataframe: "StreamingDataFrame",
         name: Optional[str] = None,
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ):
         super().__init__(
             duration_ms=duration_ms,
             grace_ms=grace_ms,
             dataframe=dataframe,
             name=name,
             on_late=on_late,
+            before_update=before_update,
+            after_update=after_update,
         )
-        self._on_update = on_update
 
     def _get_name(self, func_name: Optional[str]) -> str:
         prefix = f"{self._name}_tumbling_window" if self._name else "tumbling_window"
@@ -377,7 +386,8 @@ def _create_window(
             aggregators=aggregators or {},
             collectors=collectors or {},
             on_late=self._on_late,
-            on_update=self._on_update,
+            before_update=self._before_update,
+            after_update=self._after_update,
         )
 
 
@@ -389,11 +399,12 @@ def __init__(
         dataframe: "StreamingDataFrame",
         name: Optional[str] = None,
         on_late: Optional[WindowOnLateCallback] = None,
-        on_update: Optional[WindowOnUpdateCallback] = None,
+        before_update: Optional[WindowBeforeUpdateCallback] = None,
+        after_update: Optional[WindowAfterUpdateCallback] = None,
     ):
-        if on_update is not None:
+        if before_update is not None or after_update is not None:
             raise ValueError(
-                "Sliding windows do not support the 'on_update' trigger callback. "
+                "Sliding windows do not support trigger callbacks (before_update/after_update). "
                 "Use tumbling or hopping windows instead."
             )
         super().__init__(
@@ -402,7 +413,8 @@ def __init__(
             dataframe=dataframe,
             name=name,
             on_late=on_late,
-            on_update=on_update,
+            before_update=before_update,
+            after_update=after_update,
         )
 
     def _get_name(self, func_name: Optional[str]) -> str:
@@ -434,7 +446,8 @@ def _create_window(
             aggregators=aggregators or {},
             collectors=collectors or {},
             on_late=self._on_late,
-            on_update=self._on_update,
+            before_update=self._before_update,
+            after_update=self._after_update,
         )
 
 
diff --git a/quixstreams/dataframe/windows/sliding.py b/quixstreams/dataframe/windows/sliding.py
@@ -35,6 +35,7 @@ def process_window(
         value: Any,
         key: Any,
         timestamp_ms: int,
+        headers: Any,
         transaction: WindowedPartitionTransaction,
     ) -> tuple[Iterable[WindowKeyResult], Iterable[WindowKeyResult]]:
         """
diff --git a/quixstreams/dataframe/windows/time_based.py b/quixstreams/dataframe/windows/time_based.py
diff --git a/tests/test_quixstreams/test_dataframe/test_windows/test_hopping.py b/tests/test_quixstreams/test_dataframe/test_windows/test_hopping.py
diff --git a/tests/test_quixstreams/test_dataframe/test_windows/test_tumbling.py b/tests/test_quixstreams/test_dataframe/test_windows/test_tumbling.py
