Add how param with inner and left options

Author: gwaramadze

quixstreams/dataframe/dataframe.py

@@ -85,6 +85,9 @@
 FilterCallbackStateful = Callable[[Any, State], bool]
 FilterWithMetadataCallbackStateful = Callable[[Any, Any, int, Any, State], bool]
 
+# Constants related to join
+DISCARDED = object()
+JoinHow = Literal["inner", "left"]
 JoinOnOverlap = Literal["keep-left", "keep-right", "raise"]
 
 
@@ -1653,17 +1656,26 @@ def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
     def join(
         self,
         right: "StreamingDataFrame",
+        how: JoinHow = "inner",
         on_overlap: JoinOnOverlap = "raise",
         merger: Optional[Callable[[Any, Any], Any]] = None,
     ) -> "StreamingDataFrame":
+        if how not in get_args(JoinHow):
+            raise ValueError(
+                f"Invalid how value: {how}. "
+                f"Valid values are: {', '.join(get_args(JoinHow))}."
+            )
         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)
 
+        is_inner_join = how == "inner"
+
         if merger is None:
             if on_overlap == "keep-left":
 
@@ -1699,13 +1711,17 @@ def merger(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)
+            if is_inner_join and not right_value:
+                return DISCARDED
             return merger(value, right_value)
 
         def right_func(value, key, timestamp, headers):
             right_tx = _get_transaction(right)
             right_tx.set(timestamp=timestamp, value=value, prefix=key)
 
-        left = self.apply(left_func, metadata=True)
+        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)
 

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -10,7 +10,7 @@
 import pytest
 
 from quixstreams import State
-from quixstreams.dataframe.dataframe import JoinOnOverlap
+from quixstreams.dataframe.dataframe import JoinHow, JoinOnOverlap
 from quixstreams.dataframe.exceptions import (
     GroupByDuplicate,
     GroupByNestingLimit,
@@ -2693,17 +2693,79 @@ def _publish(sdf, topic, value, key, timestamp):
 
         return _publish
 
-    def test_join(self, create_topic, create_sdf, assign_partition, publish):
+    @pytest.mark.parametrize(
+        "how, left, right, expected",
+        [
+            (
+                "inner",
+                {"left": 1},
+                {"right": 2},
+                [({"left": 1, "right": 2}, b"key", 2, None)],
+            ),
+            (
+                "inner",
+                {"left": 1},
+                None,
+                [],
+            ),
+            (
+                "inner",
+                {"left": 1},
+                {},
+                [],
+            ),
+            (
+                "left",
+                {"left": 1},
+                {"right": 2},
+                [({"left": 1, "right": 2}, b"key", 2, None)],
+            ),
+            (
+                "left",
+                {"left": 1},
+                None,
+                [({"left": 1}, b"key", 2, None)],
+            ),
+            (
+                "left",
+                {"left": 1},
+                {},
+                [({"left": 1}, b"key", 2, None)],
+            ),
+        ],
+    )
+    def test_how(
+        self,
+        create_topic,
+        create_sdf,
+        assign_partition,
+        publish,
+        how,
+        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)
+        joined_sdf = left_sdf.join(right_sdf, how=how)
         assign_partition(right_sdf)
 
-        publish(joined_sdf, right_topic, value={"right": 1}, key=b"key", timestamp=1)
+        publish(joined_sdf, right_topic, value=right, key=b"key", timestamp=1)
         joined_value = publish(
-            joined_sdf, left_topic, value={"left": 2}, key=b"key", timestamp=2
+            joined_sdf, left_topic, value=left, key=b"key", timestamp=2
         )
-        assert joined_value == [({"left": 2, "right": 1}, b"key", 2, None)]
+        assert joined_value == expected
+
+    def test_how_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 how value: invalid. "
+            f"Valid values are: {', '.join(get_args(JoinHow))}."
+        )
+        with pytest.raises(ValueError, match=match):
+            left_sdf.join(right_sdf, how="invalid")
 
     def test_mismatching_partitions_fails(self, create_topic, create_sdf):
         left_topic, right_topic = create_topic(), create_topic(num_partitions=2)
@@ -2772,7 +2834,7 @@ def test_on_overlap(
     ):
         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)
+        joined_sdf = left_sdf.join(right_sdf, how="left", on_overlap=on_overlap)
         assign_partition(right_sdf)
 
         publish(joined_sdf, right_topic, value=right, key=b"key", timestamp=1)