Feature: Join Latest

quixstreams/dataframe/dataframe.py

@@ -49,12 +49,14 @@
 from quixstreams.models.serializers import DeserializerType, SerializerType
 from quixstreams.sinks import BaseSink
 from quixstreams.state.base import State
+from quixstreams.state.manager import StoreTypes
 from quixstreams.utils.printing import (
     DEFAULT_COLUMN_NAME,
     DEFAULT_LIVE,
     DEFAULT_LIVE_SLOWDOWN,
 )
 
+from . import joins
 from .exceptions import InvalidOperation, TopicPartitionsMismatch
 from .registry import DataFrameRegistry
 from .series import StreamingSeries
@@ -275,7 +277,7 @@ def func(d: dict, state: State):
             Default - `False`.
         """
         if stateful:
-            self._register_store()
+            self.register_store()
             # Force the callback to accept metadata
             if metadata:
                 with_metadata_func = cast(ApplyWithMetadataCallbackStateful, func)
@@ -284,11 +286,7 @@ def func(d: dict, state: State):
                     cast(ApplyCallbackStateful, func)
                 )
 
-            stateful_func = _as_stateful(
-                func=with_metadata_func,
-                processing_context=self._processing_context,
-                stream_id=self.stream_id,
-            )
+            stateful_func = _as_stateful(with_metadata_func, self)
             stream = self.stream.add_apply(stateful_func, expand=expand, metadata=True)  # type: ignore[call-overload]
         else:
             stream = self.stream.add_apply(
@@ -384,7 +382,7 @@ def func(values: list, state: State):
         :return: the updated StreamingDataFrame instance (reassignment NOT required).
         """
         if stateful:
-            self._register_store()
+            self.register_store()
             # Force the callback to accept metadata
             if metadata:
                 with_metadata_func = cast(UpdateWithMetadataCallbackStateful, func)
@@ -393,11 +391,7 @@ def func(values: list, state: State):
                     cast(UpdateCallbackStateful, func)
                 )
 
-            stateful_func = _as_stateful(
-                func=with_metadata_func,
-                processing_context=self._processing_context,
-                stream_id=self.stream_id,
-            )
+            stateful_func = _as_stateful(with_metadata_func, self)
             return self._add_update(stateful_func, metadata=True)
         else:
             return self._add_update(
@@ -486,7 +480,7 @@ def func(d: dict, state: State):
         """
 
         if stateful:
-            self._register_store()
+            self.register_store()
             # Force the callback to accept metadata
             if metadata:
                 with_metadata_func = cast(FilterWithMetadataCallbackStateful, func)
@@ -495,11 +489,7 @@ def func(d: dict, state: State):
                     cast(FilterCallbackStateful, func)
                 )
 
-            stateful_func = _as_stateful(
-                func=with_metadata_func,
-                processing_context=self._processing_context,
-                stream_id=self.stream_id,
-            )
+            stateful_func = _as_stateful(with_metadata_func, self)
             stream = self.stream.add_filter(stateful_func, metadata=True)
         else:
             stream = self.stream.add_filter(  # type: ignore[call-overload]
@@ -1656,12 +1646,62 @@ def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
             *self.topics, *other.topics, stream=merged_stream
         )
 
-    def ensure_topics_copartitioned(self):
-        partitions_counts = set(t.broker_config.num_partitions for t in self._topics)
+    def join_latest(
+        self,
+        right: "StreamingDataFrame",
+        how: joins.How = "inner",
+        on_overlap: Union[joins.OnOverlap, Callable[[Any, Any], Any]] = "raise",
+        grace_ms: Union[int, timedelta] = timedelta(days=7),
+        # TODO: Allow passing the store name here?
+    ) -> "StreamingDataFrame":
+        """
+        Join the StreamingDataFrame with the latest effective values on the right side.
+        This join is built with enrichment use case in mind when the left side is a data stream and the right side
+        is metadata.
+
+        The underlying topics of the dataframes must have the same number of partitions
+        and use the same partitioner (keys should be distributed between partitions using the same method).
+
+        Joining dataframes belonging to the same topics (aka "self-join") is not supported as of now.
+
+        How the joining works:
+            - Records from the right side get written to the state store without emitting any updates downstream.
+            - Records on the left side query the right store for the values with the same **key** and the timestamp lower or equal to the record's timestamp.
+              Left side emits data downstream.
+            - If the match is found, the two records are merged together into a new one according to the `on_overlap` logic.
+            - The size of the right store is controlled by the "grace_ms":
+              a newly added "right" record expires other values with the same key with timestamps below "<current timestamp> - <grace_ms>".
+
+        :param right: a StreamingDataFrame to join with.
+
+        :param how: the join strategy. Can be one of:
+          - "inner" - emits the result when the match on the right side is found for the left record.
+          - "left" - emits the result for each left record even if there is no match on the right side.
+          Default - `"inner"`.
+
+        :param on_overlap: how to merge the matched records together assuming they are dictionaries:
+            - "raise" - fail with an error if the same keys are found in both dictionaries
+            - "keep-left" - prefer the keys from the left record.
+            - "keep-right" - prefer the keys from the right record
+            - callback - a callback in form "(<left>, <right>) -> <new record>" to merge the records manually.
+              Use it to customize the merging logic or when one of the records is not a dictionary.
+
+        :param grace_ms: how long to keep the right records in the store.
+           Can be specified as either an `int` representing milliseconds or as a `timedelta` object.
+           Default - 7 days.
+
+        """
+        return joins.JoinLatest(how=how, on_overlap=on_overlap, grace_ms=grace_ms).join(
+            self, right
+        )
+
+    def ensure_topics_copartitioned(self, *topics: Topic):
+        topics = topics or self._topics
+        partitions_counts = set(t.broker_config.num_partitions for t in topics)
         if len(partitions_counts) > 1:
             msg = ", ".join(
                 f'"{t.name}" ({t.broker_config.num_partitions} partitions)'
-                for t in self._topics
+                for t in topics
             )
             raise TopicPartitionsMismatch(
                 f"The underlying topics must have the same number of partitions to use State; got {msg}"
@@ -1689,7 +1729,7 @@ def _add_update(
         self._stream = self._stream.add_update(func, metadata=metadata)  # type: ignore[call-overload]
         return self
 
-    def _register_store(self):
+    def register_store(self, store_type: Optional[StoreTypes] = None) -> None:
         """
         Register the default store for the current stream_id in StateStoreManager.
         """
@@ -1699,7 +1739,9 @@ def _register_store(self):
         changelog_topic_config = self._topic_manager.derive_topic_config(self._topics)
 
         self._processing_context.state_manager.register_store(
-            stream_id=self.stream_id, changelog_config=changelog_topic_config
+            stream_id=self.stream_id,
+            store_type=store_type,
+            changelog_config=changelog_topic_config,
         )
 
     def _groupby_key(
@@ -1847,19 +1889,16 @@ def wrapper(
 
 def _as_stateful(
     func: Callable[[Any, Any, int, Any, State], T],
-    processing_context: ProcessingContext,
-    stream_id: str,
+    sdf: StreamingDataFrame,
 ) -> Callable[[Any, Any, int, Any], T]:
     @functools.wraps(func)
     def wrapper(value: Any, key: Any, timestamp: int, headers: Any) -> Any:
-        ctx = message_context()
-        transaction = processing_context.checkpoint.get_store_transaction(
-            stream_id=stream_id,
-            partition=ctx.partition,
-        )
         # Pass a State object with an interface limited to the key updates only
         # and prefix all the state keys by the message key
-        state = transaction.as_state(prefix=key)
+        state = sdf.processing_context.checkpoint.get_store_transaction(
+            stream_id=sdf.stream_id,
+            partition=message_context().partition,
+        ).as_state(prefix=key)
         return func(value, key, timestamp, headers, state)
 
     return wrapper

quixstreams/dataframe/joins/__init__.py

@@ -0,0 +1,3 @@
+from .join_latest import How as How
+from .join_latest import JoinLatest as JoinLatest
+from .join_latest import OnOverlap as OnOverlap

quixstreams/dataframe/joins/join_latest.py

@@ -0,0 +1,113 @@
+import typing
+from datetime import timedelta
+from typing import Any, Callable, Literal, Union, cast, get_args
+
+from quixstreams.context import message_context
+from quixstreams.dataframe.utils import ensure_milliseconds
+from quixstreams.models.topics.manager import TopicManager
+from quixstreams.state.rocksdb.timestamped import TimestampedPartitionTransaction
+
+from .utils import keep_left_merger, keep_right_merger, raise_merger
+
+if typing.TYPE_CHECKING:
+    from quixstreams.dataframe.dataframe import StreamingDataFrame
+
+DISCARDED = object()
+How = Literal["inner", "left"]
+How_choices = get_args(How)
+
+OnOverlap = Literal["keep-left", "keep-right", "raise"]
+OnOverlap_choices = get_args(OnOverlap)
+
+
+class JoinLatest:
+    def __init__(
+        self,
+        how: How,
+        on_overlap: Union[OnOverlap, Callable[[Any, Any], Any]],
+        grace_ms: Union[int, timedelta],
+        store_name: str = "join",
+    ):
+        if how not in How_choices:
+            raise ValueError(
+                f'Invalid "how" value: {how}. '
+                f"Valid choices are: {', '.join(How_choices)}."
+            )
+        self._how = how
+
+        if callable(on_overlap):
+            self._merger = on_overlap
+        elif on_overlap == "keep-left":
+            self._merger = keep_left_merger
+        elif on_overlap == "keep-right":
+            self._merger = keep_right_merger
+        elif on_overlap == "raise":
+            self._merger = raise_merger
+        else:
+            raise ValueError(
+                f'Invalid "on_overlap" value: {on_overlap}. '
+                f"Provide either one of {', '.join(OnOverlap_choices)} or "
+                f"a callable to merge records manually."
+            )
+
+        self._retention_ms = ensure_milliseconds(grace_ms)
+        self._store_name = store_name
+
+    def join(
+        self,
+        left: "StreamingDataFrame",
+        right: "StreamingDataFrame",
+    ) -> "StreamingDataFrame":
+        if left.stream_id == right.stream_id:
+            raise ValueError(
+                "Joining dataframes originating from "
+                "the same topic is not yet supported.",
+            )
+        left.ensure_topics_copartitioned(*left.topics, *right.topics)
+
+        for sdf in (left, right):
+            changelog_config = TopicManager.derive_topic_config(sdf.topics)
+            sdf.processing_context.state_manager.register_timestamped_store(
+                stream_id=sdf.stream_id,
+                store_name=self._store_name,
+                changelog_config=changelog_config,
+            )
+
+        is_inner_join = self._how == "inner"
+
+        def left_func(value, key, timestamp, headers):
+            tx = cast(
+                TimestampedPartitionTransaction,
+                right.processing_context.checkpoint.get_store_transaction(
+                    stream_id=right.stream_id,
+                    partition=message_context().partition,
+                    store_name=self._store_name,
+                ),
+            )
+
+            right_value = tx.get_last(timestamp=timestamp, prefix=key)
+            if is_inner_join and not right_value:
+                return DISCARDED
+            return self._merger(value, right_value)
+
+        def right_func(value, key, timestamp, headers):
+            tx = cast(
+                TimestampedPartitionTransaction,
+                right.processing_context.checkpoint.get_store_transaction(
+                    stream_id=right.stream_id,
+                    partition=message_context().partition,
+                    store_name=self._store_name,
+                ),
+            )
+            tx.set_for_timestamp(
+                timestamp=timestamp,
+                value=value,
+                prefix=key,
+                retention_ms=self._retention_ms,
+            )
+
+        right = right.update(right_func, metadata=True).filter(lambda value: False)
+        left = left.apply(left_func, metadata=True).filter(
+            lambda value: value is not DISCARDED
+        )
+        return left.concat(right)

quixstreams/dataframe/joins/utils.py

@@ -0,0 +1,36 @@
+from typing import Mapping, Optional
+
+
+def keep_left_merger(left: Optional[Mapping], right: Optional[Mapping]) -> dict:
+    """
+    Merge two dictionaries, preferring values from the left dictionary
+    """
+    left = left or {}
+    right = right or {}
+    return {**right, **left}
+
+
+def keep_right_merger(left: Optional[Mapping], right: Optional[Mapping]) -> dict:
+    """
+    Merge two dictionaries, preferring values from the right dictionary
+    """
+    left = left or {}
+    right = right or {}
+    # TODO: Add try-except everywhere and tell to pass a callback if one of the objects is not a mapping
+    return {**left, **right}
+
+
+def raise_merger(left: Optional[Mapping], right: Optional[Mapping]) -> dict:
+    """
+    Merge two dictionaries and raise an error if overlapping keys detected
+    """
+    left = left or {}
+    right = right or {}
+    if overlapping_columns := left.keys() & right.keys():
+        overlapping_columns_str = ", ".join(sorted(overlapping_columns))
+        raise ValueError(
+            f"Overlapping columns: {overlapping_columns_str}."
+            'You need to provide either an "on_overlap" value of '
+            "'keep-left' or 'keep-right' or a custom merger function."
+        )
+    return {**left, **right}

quixstreams/state/exceptions.py

@@ -10,9 +10,13 @@ class PartitionStoreIsUsed(QuixException): ...
 class StoreNotRegisteredError(QuixException): ...
 
 
+# TODO: Merge these two exceptions together?
 class WindowedStoreAlreadyRegisteredError(QuixException): ...
 
 
+class StoreAlreadyRegisteredError(QuixException): ...
+
+
 class InvalidStoreTransactionStateError(QuixException): ...