Add on_overlap and merger params

Author: gwaramadze

quixstreams/dataframe/dataframe.py

@@ -20,6 +20,7 @@
     TypeVar,
     Union,
     cast,
+    get_args,
     overload,
 )
 
@@ -84,6 +85,8 @@
 FilterCallbackStateful = Callable[[Any, State], bool]
 FilterWithMetadataCallbackStateful = Callable[[Any, Any, int, Any, State], bool]
 
+JoinOnOverlap = Literal["keep-left", "keep-right", "raise"]
+
 
 class StreamingDataFrame:
     """
@@ -1647,14 +1650,56 @@ def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
             *self.topics, *other.topics, stream=merged_stream
         )
 
-    def join(self, right: "StreamingDataFrame") -> "StreamingDataFrame":
+    def join(
+        self,
+        right: "StreamingDataFrame",
+        on_overlap: JoinOnOverlap = "raise",
+        merger: Optional[Callable[[Any, Any], Any]] = None,
+    ) -> "StreamingDataFrame":
+        if on_overlap not in get_args(JoinOnOverlap):
+            raise ValueError(
+                f"Invalid on_overlap value: {on_overlap}. "
+                f"Valid values are: {', '.join(get_args(JoinOnOverlap))}."
+            )
         self.ensure_topics_copartitioned(*self.topics, *right.topics)
         right.register_store(store_type=TimestampedStore)
 
+        if merger is None:
+            if on_overlap == "keep-left":
+
+                def merger(left_value, right_value):
+                    """
+                    Merge two dictionaries, preferring values from the left dictionary
+                    and preserving order with left columns coming first.
+                    """
+                    right_value = {
+                        key: value
+                        for key, value in (right_value or {}).items()
+                        if key not in left_value
+                    }
+                    return {**left_value, **right_value}
+
+            elif on_overlap == "keep-right":
+
+                def merger(left_value, right_value):
+                    return {**left_value, **(right_value or {})}
+            elif on_overlap == "raise":
+
+                def merger(left_value, right_value):
+                    right_value = right_value or {}
+                    if overlapping_columns := left_value.keys() & right_value.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_value, **right_value}
+
         def left_func(value, key, timestamp, headers):
             right_tx = _get_transaction(right)
             right_value = right_tx.get_last(timestamp=timestamp, prefix=key)
-            return {**value, **(right_value or {})}
+            return merger(value, right_value)
 
         def right_func(value, key, timestamp, headers):
             right_tx = _get_transaction(right)

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -4,12 +4,13 @@
 import warnings
 from collections import namedtuple
 from datetime import timedelta
-from typing import Any
+from typing import Any, get_args
 from unittest import mock
 
 import pytest
 
 from quixstreams import State
+from quixstreams.dataframe.dataframe import JoinOnOverlap
 from quixstreams.dataframe.exceptions import (
     GroupByDuplicate,
     GroupByNestingLimit,
@@ -2710,3 +2711,102 @@ def test_mismatching_partitions_fails(self, create_topic, create_sdf):
 
         with pytest.raises(TopicPartitionsMismatch):
             left_sdf.join(right_sdf)
+
+    @pytest.mark.parametrize(
+        "on_overlap, left, right, expected",
+        [
+            (
+                "keep-left",
+                {"A": 1},
+                None,
+                {"A": 1},
+            ),
+            (
+                "keep-left",
+                {"A": 1, "B": "left"},
+                {"B": "right", "C": 2},
+                {"A": 1, "B": "left", "C": 2},
+            ),
+            (
+                "keep-right",
+                {"A": 1},
+                None,
+                {"A": 1},
+            ),
+            (
+                "keep-right",
+                {"A": 1, "B": "left"},
+                {"B": "right", "C": 2},
+                {"A": 1, "B": "right", "C": 2},
+            ),
+            (
+                "raise",
+                {"A": 1},
+                None,
+                {"A": 1},
+            ),
+            (
+                "raise",
+                {"A": 1},
+                {"B": 2},
+                {"A": 1, "B": 2},
+            ),
+            (
+                "raise",
+                {"A": 1, "B": "left B", "C": "left C"},
+                {"B": "right B", "C": "right C"},
+                ValueError("Overlapping columns: B, C."),
+            ),
+        ],
+    )
+    def test_on_overlap(
+        self,
+        create_topic,
+        create_sdf,
+        assign_partition,
+        publish,
+        on_overlap,
+        left,
+        right,
+        expected,
+    ):
+        left_topic, right_topic = create_topic(), create_topic()
+        left_sdf, right_sdf = create_sdf(left_topic), create_sdf(right_topic)
+        joined_sdf = left_sdf.join(right_sdf, on_overlap=on_overlap)
+        assign_partition(right_sdf)
+
+        publish(joined_sdf, right_topic, value=right, key=b"key", timestamp=1)
+
+        if isinstance(expected, Exception):
+            with pytest.raises(expected.__class__, match=expected.args[0]):
+                publish(joined_sdf, left_topic, value=left, key=b"key", timestamp=2)
+        else:
+            joined_value = publish(
+                joined_sdf, left_topic, value=left, key=b"key", timestamp=2
+            )
+            assert joined_value == [(expected, b"key", 2, None)]
+
+    def test_on_overlap_invalid_value(self, create_topic, create_sdf):
+        left_topic, right_topic = create_topic(), create_topic()
+        left_sdf, right_sdf = create_sdf(left_topic), create_sdf(right_topic)
+
+        match = (
+            "Invalid on_overlap value: invalid. "
+            f"Valid values are: {', '.join(get_args(JoinOnOverlap))}."
+        )
+        with pytest.raises(ValueError, match=match):
+            left_sdf.join(right_sdf, on_overlap="invalid")
+
+    def test_custom_merger(self, create_topic, create_sdf, assign_partition, publish):
+        left_topic, right_topic = create_topic(), create_topic()
+        left_sdf, right_sdf = create_sdf(left_topic), create_sdf(right_topic)
+
+        def merger(left, right):
+            return {"left": left, "right": right}
+
+        joined_sdf = left_sdf.join(right_sdf, merger=merger)
+        assign_partition(right_sdf)
+
+        publish(joined_sdf, right_topic, value=1, key=b"key", timestamp=1)
+        joined_value = publish(joined_sdf, left_topic, value=2, key=b"key", timestamp=2)
+        assert joined_value == [({"left": 2, "right": 1}, b"key", 2, None)]