Correct heartbeat handling in windowing

Author: gwaramadze

quixstreams/dataframe/windows/base.py

@@ -19,7 +19,6 @@
 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
 
@@ -93,13 +92,19 @@ def _apply_window(
             stream_id=self._dataframe.stream_id,
             processing_context=self._dataframe.processing_context,
             store_name=name,
-            heartbeat_func=heartbeat_func,
+        )
+        heartbeat_func = _as_heartbeat(
+            func=heartbeat_func,
+            stream_id=self._dataframe.stream_id,
+            processing_context=self._dataframe.processing_context,
+            store_name=name,
         )
         # Manually modify the Stream and clone the source StreamingDataFrame
         # to avoid adding "transform" API to it.
         # Transform callbacks can modify record key and timestamp,
         # and it's prone to misuse.
         stream = self._dataframe.stream.add_transform(func=windowed_func, expand=True)
+        stream = stream.add_heartbeat(func=heartbeat_func)
         return self._dataframe.__dataframe_clone__(stream=stream)
 
     def final(self) -> "StreamingDataFrame":
@@ -417,7 +422,6 @@ def _as_windowed(
     processing_context: "ProcessingContext",
     store_name: str,
     stream_id: str,
-    heartbeat_func,
 ) -> TransformExpandedCallback:
     @functools.wraps(func)
     def wrapper(
@@ -430,9 +434,6 @@ def wrapper(
                 stream_id=stream_id, partition=ctx.partition, store_name=store_name
             ),
         )
-        if is_heartbeat_message(key, value):
-            return heartbeat_func(timestamp, transaction)
-
         if key is None:
             logger.warning(
                 f"Skipping window processing for a message because the key is None, "
@@ -444,6 +445,26 @@ def wrapper(
     return wrapper
 
 
+def _as_heartbeat(
+    func,  # TODO: typing?
+    processing_context: "ProcessingContext",
+    store_name: str,
+    stream_id: str,
+):  # TODO: typing?
+    @functools.wraps(func)
+    def wrapper(timestamp: int) -> Iterable[Message]:
+        ctx = message_context()
+        transaction = cast(
+            WindowedPartitionTransaction,
+            processing_context.checkpoint.get_store_transaction(
+                stream_id=stream_id, partition=ctx.partition, store_name=store_name
+            ),
+        )
+        return func(timestamp, transaction)
+
+    return wrapper
+
+
 class WindowOnLateCallback(Protocol):
     def __call__(
         self,