Refactor wall clock processing

Author: gwaramadze

quixstreams/app.py

@@ -1035,16 +1035,44 @@ def _process_wall_clock(self, wall_clock_executors):
 
         value, key, timestamp, headers = None, None, int(now * 1000), {}
 
-        for tp in self._consumer.assignment():
-            if executor := wall_clock_executors.get(tp.topic):
+        # Offsets processed in the current, open checkpoint (in-flight)
+        inflight_offsets = self._processing_context.checkpoint.inflight_offsets
+        assignment = self._consumer.assignment()
+
+        for topics, executor in wall_clock_executors.items():
+            topic_names = {t.name for t in topics}
+            seen_partitions: set[int] = set()
+            tps: list[tuple[str, int, int]] = []
+
+            # Gather assigned TPs that belong to the topics
+            for tp in assignment:
+                if tp.topic in topic_names and tp.partition not in seen_partitions:
+                    offset = inflight_offsets.get((tp.topic, tp.partition))
+                    if offset is None:
+                        committed_tp = self._consumer.committed([tp], timeout=30)[0]
+                        if committed_tp.error:
+                            raise RuntimeError(
+                                "Failed to get committed offsets for "
+                                f'"{committed_tp.topic}[{committed_tp.partition}]" '
+                                f"from the broker: {committed_tp.error}"
+                            )
+                        if committed_tp.offset >= 0:
+                            offset = committed_tp.offset - 1
+
+                    if offset is not None:
+                        seen_partitions.add(tp.partition)
+                        tps.append((tp.topic, tp.partition, offset))
+
+            # Execute callback for each selected topic-partition with its offset
+            for topic, partition, offset in tps:
                 row = Row(
                     value=value,
                     key=key,
                     timestamp=timestamp,
                     context=MessageContext(
-                        topic=tp.topic,
-                        partition=tp.partition,
-                        offset=-1,  # TODO: get correct offsets
+                        topic=topic,
+                        partition=partition,
+                        offset=offset,
                         size=-1,
                     ),
                     headers=headers,

quixstreams/checkpointing/checkpoint.py

@@ -1,7 +1,8 @@
 import logging
 import time
 from abc import abstractmethod
-from typing import Dict, Tuple
+from types import MappingProxyType
+from typing import Dict, Mapping, Tuple
 
 from confluent_kafka import KafkaException, TopicPartition
 
@@ -55,6 +56,15 @@ def __init__(
         self._commit_every = commit_every
         self._total_offsets_processed = 0
 
+    @property
+    def inflight_offsets(self) -> Mapping[Tuple[str, int], int]:
+        """
+        Read-only view of processed (but not yet committed) offsets in the current checkpoint.
+
+        :return: a read-only mapping {(topic, partition): last_processed_offset}
+        """
+        return MappingProxyType(self._tp_offsets)
+
     def expired(self) -> bool:
         """
         Returns `True` if checkpoint deadline has expired OR

quixstreams/dataframe/registry.py

@@ -22,7 +22,7 @@ class DataFrameRegistry:
 
     def __init__(self) -> None:
         self._registry: dict[str, Stream] = {}
-        self._wall_clock_registry: dict[str, Stream] = {}
+        self._wall_clock_registry: dict[str, tuple[tuple[Topic, ...], Stream]] = {}
         self._topics: list[Topic] = []
         self._repartition_origins: set[str] = set()
         self._topics_to_stream_ids: dict[str, set[str]] = {}
@@ -74,15 +74,11 @@ def register_wall_clock(
         self, dataframe: "StreamingDataFrame", stream: Stream
     ) -> None:
         """
-        Register a wall clock stream for the given topic.
+        Register a wall clock stream root for the given dataframe.
+        Stores the Stream itself to be composed later with an optional sink.
         """
-        topics = dataframe.topics
-        if len(topics) > 1:
-            raise ValueError(
-                f"Expected a StreamingDataFrame with one topic, got {len(topics)}"
-            )
-        topic = topics[0]
-        self._wall_clock_registry[topic.name] = stream
+        # TODO: What if there are more wall clock streams for the same stream_id?
+        self._wall_clock_registry[dataframe.stream_id] = (dataframe.topics, stream)
 
     def register_groupby(
         self,
@@ -128,28 +124,27 @@ def compose_all(
         :param sink: callable to accumulate the results of the execution, optional.
         :return: a {topic_name: composed} dict, where composed is a callable
         """
-        return self._compose(registry=self._registry, sink=sink)
-
-    def compose_wall_clock(self) -> dict[str, VoidExecutor]:
-        """
-        Composes all the wall clock streams and returns a dict of format {<topic>: <VoidExecutor>}
-        :return: a {topic_name: composed} dict, where composed is a callable
-        """
-        return self._compose(registry=self._wall_clock_registry)
-
-    def _compose(
-        self, registry: dict[str, Stream], sink: Optional[VoidExecutor] = None
-    ) -> dict[str, VoidExecutor]:
         executors = {}
         # Go over the registered topics with root Streams and compose them
-        for topic, root_stream in registry.items():
+        for topic, root_stream in self._registry.items():
             # If a root stream is connected to other roots, ".compose()" will
             # return them all.
             # Use the registered root Stream to filter them out.
             root_executors = root_stream.compose(sink=sink)
             executors[topic] = root_executors[root_stream]
         return executors
 
+    def compose_wall_clock(self) -> dict[tuple[Topic, ...], VoidExecutor]:
+        """
+        Compose all wall clock Streams and return executors keyed by stream_id.
+        Returns mapping: {stream_id: (topics, executor)}
+        """
+        executors = {}
+        for _, (topics, root_stream) in self._wall_clock_registry.items():
+            root_executors = root_stream.compose()
+            executors[topics] = root_executors[root_stream]
+        return executors
+
     def register_stream_id(self, stream_id: str, topic_names: list[str]):
         """
         Register a mapping between the stream_id and topic names.