Self joins

Author: gwaramadze

quixstreams/dataframe/dataframe.py

@@ -1623,7 +1623,11 @@ def _sink_callback(
         # uses apply without returning to make this operation terminal
         self.apply(_sink_callback, metadata=True)
 
-    def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
+    def concat(
+        self,
+        other: "StreamingDataFrame",
+        stream_id: Optional[str] = None,
+    ) -> "StreamingDataFrame":
         """
         Concatenate two StreamingDataFrames together and return a new one.
         The transformations applied on this new StreamingDataFrame will update data
@@ -1650,7 +1654,7 @@ def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
             self._registry.require_time_alignment()
 
         return self.__dataframe_clone__(
-            *self.topics, *other.topics, stream=merged_stream
+            *self.topics, *other.topics, stream=merged_stream, stream_id=stream_id
         )
 
     def join_latest(
@@ -1726,11 +1730,11 @@ def right_func(value, key, timestamp, headers):
                 retention_ms=retention_ms,
             )
 
+        right = right.update(right_func, metadata=True).filter(lambda value: False)
         left = self.apply(left_func, metadata=True).filter(
             lambda value: value is not DISCARDED
         )
-        right = right.update(right_func, metadata=True).filter(lambda value: False)
-        return left.concat(right)
+        return left.concat(right, stream_id=f"{self.stream_id}-{right.stream_id}")
 
     def ensure_topics_copartitioned(self, *topics: Topic):
         topics = topics or self._topics
@@ -1766,7 +1770,11 @@ def _add_update(
         self._stream = self._stream.add_update(func, metadata=metadata)  # type: ignore[call-overload]
         return self
 
-    def register_store(self, store_type: Optional[StoreTypes] = None):
+    def register_store(
+        self,
+        stream_id: Optional[str] = None,
+        store_type: Optional[StoreTypes] = None,
+    ):
         """
         Register the default store for the current stream_id in StateStoreManager.
         """
@@ -1776,7 +1784,7 @@ def register_store(self, store_type: Optional[StoreTypes] = None):
         changelog_topic_config = self._topic_manager.derive_topic_config(self._topics)
 
         self._processing_context.state_manager.register_store(
-            stream_id=self.stream_id,
+            stream_id=stream_id or self.stream_id,
             store_type=store_type,
             changelog_config=changelog_topic_config,
         )
@@ -1938,8 +1946,10 @@ def wrapper(value: Any, key: Any, timestamp: int, headers: Any) -> Any:
     return wrapper
 
 
-def _get_transaction(sdf: StreamingDataFrame) -> PartitionTransaction:
+def _get_transaction(
+    sdf: StreamingDataFrame, stream_id: Optional[str] = None
+) -> PartitionTransaction:
     return sdf.processing_context.checkpoint.get_store_transaction(
-        stream_id=sdf.stream_id,
+        stream_id=stream_id or sdf.stream_id,
         partition=message_context().partition,
     )

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -2904,3 +2904,19 @@ def test_retention_ms(
 
         assert publish_left(timestamp=4) == []
         assert publish_left(timestamp=5) == [({"left": 4, "right": 2}, b"key", 5, None)]
+
+    def test_self_join(self, create_topic, create_sdf, assign_partition, publish):
+        topic = create_topic()
+        sdf = create_sdf(topic)
+        merger = lambda left, right: {"left": left, "right": right}
+        joined_sdf = sdf.join_latest(sdf, merger=merger)
+        assign_partition(sdf)
+
+        _publish = partial(publish, joined_sdf, topic, key=b"key")
+
+        assert _publish(value=1, timestamp=1) == [
+            ({"left": 1, "right": 1}, b"key", 1, None)
+        ]
+        assert _publish(value=2, timestamp=2) == [
+            ({"left": 2, "right": 2}, b"key", 2, None)
+        ]