Heartbeat basics

Author: gwaramadze

quixstreams/app.py

@@ -7,6 +7,7 @@
 import uuid
 import warnings
 from collections import defaultdict
+from datetime import datetime
 from pathlib import Path
 from typing import Callable, List, Literal, Optional, Protocol, Tuple, Type, Union
 
@@ -30,6 +31,8 @@
 from .logging import LogLevel, configure_logging
 from .models import (
     DeserializerType,
+    MessageContext,
+    Row,
     SerializerType,
     TimestampExtractor,
     Topic,
@@ -153,6 +156,7 @@ def __init__(
         topic_create_timeout: float = 60,
         processing_guarantee: ProcessingGuarantee = "at-least-once",
         max_partition_buffer_size: int = 10000,
+        heartbeat_interval: float = 0.0,
     ):
         """
         :param broker_address: Connection settings for Kafka.
@@ -222,6 +226,12 @@ def __init__(
             It is a soft limit, and the actual number of buffered messages can be up to x2 higher.
             Lower value decreases the memory use, but increases the latency.
             Default - `10000`.
+        :param heartbeat_interval: the interval (seconds) at which to send heartbeat messages.
+            The heartbeat timing starts counting from application start.
+            TODO: Save and respect last heartbeat timestamp.
+            The heartbeat is sent for every partition of every topic with registered heartbeat streams.
+            If the value is 0, no heartbeat messages will be sent.
+            Default - `0.0`.
 
         <br><br>***Error Handlers***<br>
         To handle errors, `Application` accepts callbacks triggered when
@@ -370,6 +380,10 @@ def __init__(
             recovery_manager=recovery_manager,
         )
 
+        self._heartbeat_active = heartbeat_interval > 0
+        self._heartbeat_interval = heartbeat_interval
+        self._heartbeat_last_sent = datetime.now().timestamp()
+
         self._source_manager = SourceManager()
         self._sink_manager = SinkManager()
         self._dataframe_registry = DataFrameRegistry()
@@ -899,6 +913,7 @@ def _run_dataframe(self, sink: Optional[VoidExecutor] = None):
         processing_context = self._processing_context
         source_manager = self._source_manager
         process_message = self._process_message
+        process_heartbeat = self._process_heartbeat
         printer = self._processing_context.printer
         run_tracker = self._run_tracker
         consumer = self._consumer
@@ -911,6 +926,9 @@ def _run_dataframe(self, sink: Optional[VoidExecutor] = None):
         )
 
         dataframes_composed = self._dataframe_registry.compose_all(sink=sink)
+        heartbeats_composed = self._dataframe_registry.compose_heartbeats()
+        if not heartbeats_composed:
+            self._heartbeat_active = False
 
         processing_context.init_checkpoint()
         run_tracker.set_as_running()
@@ -922,6 +940,7 @@ def _run_dataframe(self, sink: Optional[VoidExecutor] = None):
                 run_tracker.timeout_refresh()
             else:
                 process_message(dataframes_composed)
+                process_heartbeat(heartbeats_composed)
                 processing_context.commit_checkpoint()
                 consumer.resume_backpressured()
                 source_manager.raise_for_error()
@@ -1005,6 +1024,41 @@ def _process_message(self, dataframe_composed):
         if self._on_message_processed is not None:
             self._on_message_processed(topic_name, partition, offset)
 
+    def _process_heartbeat(self, heartbeats_composed):
+        if not self._heartbeat_active:
+            return
+
+        now = datetime.now().timestamp()
+        if self._heartbeat_last_sent > now - self._heartbeat_interval:
+            return
+
+        value, key, timestamp, headers = None, None, int(now * 1000), {}
+
+        for tp in self._consumer.assignment():
+            if executor := heartbeats_composed.get(tp.topic):
+                row = Row(
+                    value=value,
+                    key=key,
+                    timestamp=timestamp,
+                    context=MessageContext(
+                        topic=tp.topic,
+                        partition=tp.partition,
+                        offset=-1,  # TODO: get correct offsets
+                        size=-1,
+                    ),
+                    headers=headers,
+                )
+                context = copy_context()
+                context.run(set_message_context, row.context)
+                try:
+                    context.run(executor, value, key, timestamp, headers)
+                except Exception as exc:
+                    to_suppress = self._on_processing_error(exc, row, logger)
+                    if not to_suppress:
+                        raise
+
+        self._heartbeat_last_sent = now
+
     def _on_assign(self, _, topic_partitions: List[TopicPartition]):
         """
         Assign new topic partitions to consumer and state.

quixstreams/core/stream/functions/transform.py

@@ -1,7 +1,13 @@
 from typing import Any, Literal, Union, cast, overload
 
 from .base import StreamFunction
-from .types import TransformCallback, TransformExpandedCallback, VoidExecutor
+from .types import (
+    TransformCallback,
+    TransformExpandedCallback,
+    TransformHeartbeatCallback,
+    TransformHeartbeatExpandedCallback,
+    VoidExecutor,
+)
 
 __all__ = ("TransformFunction",)
 
@@ -23,28 +29,81 @@ class TransformFunction(StreamFunction):
 
     @overload
     def __init__(
-        self, func: TransformCallback, expand: Literal[False] = False
+        self,
+        func: TransformCallback,
+        expand: Literal[False] = False,
+        heartbeat: Literal[False] = False,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        func: TransformExpandedCallback,
+        expand: Literal[True],
+        heartbeat: Literal[False] = False,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        func: TransformHeartbeatCallback,
+        expand: Literal[False] = False,
+        heartbeat: Literal[True] = True,
     ) -> None: ...
 
     @overload
     def __init__(
-        self, func: TransformExpandedCallback, expand: Literal[True]
+        self,
+        func: TransformHeartbeatExpandedCallback,
+        expand: Literal[True],
+        heartbeat: Literal[True],
     ) -> None: ...
 
     def __init__(
         self,
-        func: Union[TransformCallback, TransformExpandedCallback],
+        func: Union[
+            TransformCallback,
+            TransformExpandedCallback,
+            TransformHeartbeatCallback,
+            TransformHeartbeatExpandedCallback,
+        ],
         expand: bool = False,
+        heartbeat: bool = False,
     ):
         super().__init__(func)
 
-        self.func: Union[TransformCallback, TransformExpandedCallback]
+        self.func: Union[
+            TransformCallback,
+            TransformExpandedCallback,
+            TransformHeartbeatCallback,
+            TransformHeartbeatExpandedCallback,
+        ]
         self.expand = expand
+        self.heartbeat = heartbeat
 
     def get_executor(self, *child_executors: VoidExecutor) -> VoidExecutor:
         child_executor = self._resolve_branching(*child_executors)
 
-        if self.expand:
+        if self.expand and self.heartbeat:
+            heartbeat_expanded_func = cast(
+                TransformHeartbeatExpandedCallback, self.func
+            )
+
+            def wrapper(
+                value: Any,
+                key: Any,
+                timestamp: int,
+                headers: Any,
+            ):
+                for (
+                    new_value,
+                    new_key,
+                    new_timestamp,
+                    new_headers,
+                ) in heartbeat_expanded_func(timestamp):
+                    child_executor(new_value, new_key, new_timestamp, new_headers)
+
+        elif self.expand:
             expanded_func = cast(TransformExpandedCallback, self.func)
 
             def wrapper(
@@ -57,8 +116,22 @@ def wrapper(
                 for new_value, new_key, new_timestamp, new_headers in result:
                     child_executor(new_value, new_key, new_timestamp, new_headers)
 
+        elif self.heartbeat:
+            heartbeat_func = cast(TransformHeartbeatCallback, self.func)
+
+            def wrapper(
+                value: Any,
+                key: Any,
+                timestamp: int,
+                headers: Any,
+            ):
+                new_value, new_key, new_timestamp, new_headers = heartbeat_func(
+                    timestamp
+                )
+                child_executor(new_value, new_key, new_timestamp, new_headers)
+
         else:
-            func = cast(TransformCallback, self.func)
+            regular_func = cast(TransformCallback, self.func)
 
             def wrapper(
                 value: Any,
@@ -67,7 +140,7 @@ def wrapper(
                 headers: Any,
             ):
                 # Execute a function on a single value and return its result
-                new_value, new_key, new_timestamp, new_headers = func(
+                new_value, new_key, new_timestamp, new_headers = regular_func(
                     value, key, timestamp, headers
                 )
                 child_executor(new_value, new_key, new_timestamp, new_headers)

quixstreams/core/stream/functions/types.py

@@ -14,6 +14,8 @@
     "FilterWithMetadataCallback",
     "TransformCallback",
     "TransformExpandedCallback",
+    "TransformHeartbeatCallback",
+    "TransformHeartbeatExpandedCallback",
 )
 
 
@@ -35,6 +37,10 @@ def __bool__(self) -> bool: ...
 TransformExpandedCallback = Callable[
     [Any, Any, int, Any], Iterable[Tuple[Any, Any, int, Any]]
 ]
+TransformHeartbeatCallback = Callable[[int], Tuple[Any, Any, int, Any]]
+TransformHeartbeatExpandedCallback = Callable[
+    [int], Iterable[Tuple[Any, Any, int, Any]]
+]
 
 StreamCallback = Union[
     ApplyCallback,

quixstreams/core/stream/stream.py

@@ -30,6 +30,8 @@
     TransformCallback,
     TransformExpandedCallback,
     TransformFunction,
+    TransformHeartbeatCallback,
+    TransformHeartbeatExpandedCallback,
     UpdateCallback,
     UpdateFunction,
     UpdateWithMetadataCallback,
@@ -249,18 +251,56 @@ def add_update(
         return self._add(update_func)
 
     @overload
-    def add_transform(self, func: TransformCallback, *, expand: Literal[False] = False):
+    def add_transform(
+        self,
+        func: TransformCallback,
+        *,
+        expand: Literal[False] = False,
+        heartbeat: Literal[False] = False,
+    ):
         pass
 
     @overload
-    def add_transform(self, func: TransformExpandedCallback, *, expand: Literal[True]):
+    def add_transform(
+        self,
+        func: TransformExpandedCallback,
+        *,
+        expand: Literal[True],
+        heartbeat: Literal[False] = False,
+    ):
         pass
 
+    @overload
     def add_transform(
         self,
-        func: Union[TransformCallback, TransformExpandedCallback],
+        func: TransformHeartbeatCallback,
+        *,
+        expand: Literal[False] = False,
+        heartbeat: Literal[True],
+    ):
+        pass
+
+    @overload
+    def add_transform(
+        self,
+        func: TransformHeartbeatExpandedCallback,
+        *,
+        expand: Literal[True],
+        heartbeat: Literal[True],
+    ):
+        pass
+
+    def add_transform(
+        self,
+        func: Union[
+            TransformCallback,
+            TransformExpandedCallback,
+            TransformHeartbeatCallback,
+            TransformHeartbeatExpandedCallback,
+        ],
         *,
         expand: bool = False,
+        heartbeat: bool = False,
     ) -> "Stream":
         """
         Add a "transform" function to the Stream, that will mutate the input value.
@@ -276,9 +316,11 @@ def add_transform(
         :param expand: if True, expand the returned iterable into individual items
             downstream. If returned value is not iterable, `TypeError` will be raised.
             Default - `False`.
+        :param heartbeat: if True, the callback is expected to accept timestamp only.
+            Default - `False`.
         :return: a new Stream derived from the current one
         """
-        return self._add(TransformFunction(func, expand=expand))  # type: ignore[call-overload]
+        return self._add(TransformFunction(func, expand=expand, heartbeat=heartbeat))  # type: ignore[call-overload]
 
     def merge(self, other: "Stream") -> "Stream":
         """

quixstreams/dataframe/dataframe.py

@@ -1696,6 +1696,10 @@ def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
                 stream=merged_stream, stream_id=merged_stream_id
             )
 
+    def concat_heartbeat(self, stream: Stream) -> "StreamingDataFrame":
+        self._registry.register_heartbeat(self, stream)
+        return self.__dataframe_clone__(stream=self.stream.merge(stream))
+
     def join_asof(
         self,
         right: "StreamingDataFrame",