Correct heartbeat ignore

Author: gwaramadze

quixstreams/core/stream/functions/__init__.py

@@ -2,7 +2,7 @@
 from .apply import *
 from .base import *
 from .filter import *
+from .heartbeat import *
 from .transform import *
 from .types import *
 from .update import *
-from .heartbeat import *

quixstreams/core/stream/functions/apply.py

@@ -1,7 +1,7 @@
 from typing import Any, Literal, Union, overload
 
 from .base import StreamFunction
-from .heartbeat import ignore_heartbeat
+from .heartbeat import is_heartbeat_message
 from .types import (
     ApplyCallback,
     ApplyExpandedCallback,
@@ -49,6 +49,10 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
             ) -> None:
+                # Pass heartbeat messages downstream
+                if is_heartbeat_message(key, value):
+                    child_executor(value, key, timestamp, headers)
+
                 # Execute a function on a single value and wrap results into a list
                 # to expand them downstream
                 result = func(value)
@@ -63,11 +67,13 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
             ) -> None:
-                # Execute a function on a single value and return its result
-                result = func(value)
-                child_executor(result, key, timestamp, headers)
+                # Pass heartbeat messages downstream or execute
+                # a function on a single value and return its result
+                if not is_heartbeat_message(key, value):
+                    value = func(value)
+                child_executor(value, key, timestamp, headers)
 
-        return ignore_heartbeat(wrapper)
+        return wrapper
 
 
 class ApplyWithMetadataFunction(StreamFunction):
@@ -111,6 +117,10 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
             ):
+                # Pass heartbeat messages downstream
+                if is_heartbeat_message(key, value):
+                    child_executor(value, key, timestamp, headers)
+
                 # Execute a function on a single value and wrap results into a list
                 # to expand them downstream
                 result = func(value, key, timestamp, headers)
@@ -125,8 +135,10 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
             ):
-                # Execute a function on a single value and return its result
-                result = func(value, key, timestamp, headers)
-                child_executor(result, key, timestamp, headers)
+                # Pass heartbeat messages downstream or execute
+                # a function on a single value and return its result
+                if not is_heartbeat_message(key, value):
+                    value = func(value, key, timestamp, headers)
+                child_executor(value, key, timestamp, headers)
 
-        return ignore_heartbeat(wrapper)
+        return wrapper

quixstreams/core/stream/functions/filter.py

@@ -1,5 +1,7 @@
 from typing import Any
 
+from quixstreams.core.stream.functions.heartbeat import is_heartbeat_message
+
 from .base import StreamFunction
 from .types import FilterCallback, FilterWithMetadataCallback, VoidExecutor
 
@@ -30,7 +32,7 @@ def wrapper(
             headers: Any,
         ):
             # Filter a single value
-            if func(value):
+            if is_heartbeat_message(key, value) or func(value):
                 child_executor(value, key, timestamp, headers)
 
         return wrapper
@@ -62,7 +64,7 @@ def wrapper(
             headers: Any,
         ):
             # Filter a single value
-            if func(value, key, timestamp, headers):
+            if is_heartbeat_message(key, value) or func(value, key, timestamp, headers):
                 child_executor(value, key, timestamp, headers)
 
         return wrapper

quixstreams/core/stream/functions/heartbeat.py

@@ -1,15 +1,18 @@
-from functools import wraps
 from typing import Any
 
 from quixstreams.context import message_context
 
 from .base import StreamFunction
-from .types import VoidExecutor
+from .types import HeartbeatCallback, VoidExecutor
 
-__all__ = ("HeartbeatFunction", "is_heartbeat_message", "ignore_heartbeat")
+__all__ = ("HeartbeatFunction", "is_heartbeat_message")
 
 
 class HeartbeatFunction(StreamFunction):
+    def __init__(self, func: HeartbeatCallback) -> None:
+        super().__init__(func)
+        self.func: HeartbeatCallback
+
     def get_executor(self, *child_executors: VoidExecutor) -> VoidExecutor:
         child_executor = self._resolve_branching(*child_executors)
 
@@ -21,32 +24,13 @@ def wrapper(
             timestamp: int,
             headers: Any,
         ):
-            if is_heartbeat_message(key, value) and (result := func(timestamp)):
-                for new_value, new_key, new_timestamp, new_headers in result:
-                    child_executor(new_value, new_key, new_timestamp, new_headers)
-
+            if is_heartbeat_message(key, value):
+                # TODO: Heartbeats may return values (like expired windows)
+                func(timestamp)
             child_executor(value, key, timestamp, headers)
 
         return wrapper
 
 
 def is_heartbeat_message(key: Any, value: Any) -> bool:
     return message_context().heartbeat and key is None and value is None
-
-
-def ignore_heartbeat(func):
-    """
-    Decorator that wraps a function to return early if the message is a heartbeat.
-
-    The decorated function should expect (value, key, timestamp, headers) parameters.
-    If is_heartbeat_message(key, value) returns True, the function returns early
-    without executing the wrapped function.
-    """
-
-    @wraps(func)
-    def wrapper(value: Any, key: Any, timestamp: int, headers: Any):
-        if is_heartbeat_message(key, value):
-            return
-        return func(value, key, timestamp, headers)
-
-    return wrapper

quixstreams/core/stream/functions/transform.py

@@ -1,5 +1,7 @@
 from typing import Any, Literal, Union, cast, overload
 
+from quixstreams.core.stream.functions.heartbeat import is_heartbeat_message
+
 from .base import StreamFunction
 from .types import TransformCallback, TransformExpandedCallback, VoidExecutor
 
@@ -53,6 +55,10 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
             ):
+                # Pass heartbeat messages downstream
+                if is_heartbeat_message(key, value):
+                    child_executor(value, key, timestamp, headers)
+
                 result = expanded_func(value, key, timestamp, headers)
                 for new_value, new_key, new_timestamp, new_headers in result:
                     child_executor(new_value, new_key, new_timestamp, new_headers)
@@ -66,10 +72,12 @@ def wrapper(
                 timestamp: int,
                 headers: Any,
             ):
-                # Execute a function on a single value and return its result
-                new_value, new_key, new_timestamp, new_headers = func(
-                    value, key, timestamp, headers
-                )
-                child_executor(new_value, new_key, new_timestamp, new_headers)
+                # Pass heartbeat messages downstream or execute
+                # a function on a single value and return its result
+                if is_heartbeat_message(key, value):
+                    value, key, timestamp, headers = func(
+                        value, key, timestamp, headers
+                    )
+                child_executor(value, key, timestamp, headers)
 
         return wrapper

quixstreams/core/stream/functions/types.py

@@ -36,6 +36,15 @@ def __bool__(self) -> bool: ...
     [Any, Any, int, Any], Iterable[Tuple[Any, Any, int, Any]]
 ]
 
+HeartbeatCallback = Callable[
+    [int],  # timestamp
+    Union[
+        None,
+        Tuple[Any, Any, int, Any],  # single value
+        Iterable[Tuple[Any, Any, int, Any]],  # expanded values
+    ],
+]
+
 StreamCallback = Union[
     ApplyCallback,
     ApplyExpandedCallback,

quixstreams/core/stream/functions/update.py

@@ -1,5 +1,7 @@
 from typing import Any
 
+from quixstreams.core.stream.functions.heartbeat import is_heartbeat_message
+
 from .base import StreamFunction
 from .types import UpdateCallback, UpdateWithMetadataCallback, VoidExecutor
 
@@ -28,7 +30,8 @@ def get_executor(self, *child_executors: VoidExecutor) -> VoidExecutor:
 
         def wrapper(value: Any, key: Any, timestamp: int, headers: Any):
             # Update a single value and forward it
-            func(value)
+            if not is_heartbeat_message(key, value):
+                func(value)
             child_executor(value, key, timestamp, headers)
 
         return wrapper
@@ -56,7 +59,8 @@ def get_executor(self, *child_executors: VoidExecutor) -> VoidExecutor:
 
         def wrapper(value: Any, key: Any, timestamp: int, headers: Any):
             # Update a single value and forward it
-            func(value, key, timestamp, headers)
+            if not is_heartbeat_message(key, value):
+                func(value, key, timestamp, headers)
             child_executor(value, key, timestamp, headers)
 
         return wrapper

quixstreams/dataframe/windows/base.py

@@ -19,6 +19,7 @@
 from quixstreams.context import message_context
 from quixstreams.core.stream import TransformExpandedCallback
 from quixstreams.core.stream.exceptions import InvalidOperation
+from quixstreams.core.stream.functions.heartbeat import is_heartbeat_message
 from quixstreams.models.topics.manager import TopicManager
 from quixstreams.state import WindowedPartitionTransaction
 
@@ -83,6 +84,7 @@ def _apply_window(
         self,
         func: TransformRecordCallbackExpandedWindowed,
         name: str,
+        heartbeat_func,
     ) -> "StreamingDataFrame":
         self.register_store()
 
@@ -91,6 +93,7 @@ def _apply_window(
             stream_id=self._dataframe.stream_id,
             processing_context=self._dataframe.processing_context,
             store_name=name,
+            heartbeat_func=heartbeat_func,
         )
         # Manually modify the Stream and clone the source StreamingDataFrame
         # to avoid adding "transform" API to it.
@@ -140,9 +143,13 @@ def window_callback(
             for key, window in expired_windows:
                 yield (window, key, window["start"], None)
 
+        def heartbeat_callback(timestamp: int) -> Iterable[Message]:
+            return []
+
         return self._apply_window(
             func=window_callback,
             name=self._name,
+            heartbeat_func=heartbeat_callback,
         )
 
     def current(self) -> "StreamingDataFrame":
@@ -188,7 +195,14 @@ def window_callback(
             for key, window in updated_windows:
                 yield (window, key, window["start"], None)
 
-        return self._apply_window(func=window_callback, name=self._name)
+        def heartbeat_callback(timestamp: int) -> Iterable[Message]:
+            return []
+
+        return self._apply_window(
+            func=window_callback,
+            name=self._name,
+            heartbeat_func=heartbeat_callback,
+        )
 
     # Implemented by SingleAggregationWindowMixin and MultiAggregationWindowMixin
     # Single aggregation and multi aggregation windows store aggregations and collections