Add retention_ms param

Author: gwaramadze

quixstreams/dataframe/dataframe.py

@@ -1659,6 +1659,7 @@ def join(
         how: JoinHow = "inner",
         on_overlap: JoinOnOverlap = "raise",
         merger: Optional[Callable[[Any, Any], Any]] = None,
+        retention_ms: Union[int, timedelta] = timedelta(days=7),
     ) -> "StreamingDataFrame":
         if how not in get_args(JoinHow):
             raise ValueError(
@@ -1675,6 +1676,7 @@ def join(
         right.register_store(store_type=TimestampedStore)
 
         is_inner_join = how == "inner"
+        retention_ms = ensure_milliseconds(retention_ms)
 
         if merger is None:
             if on_overlap == "keep-left":
@@ -1710,7 +1712,11 @@ 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)
+            right_value = right_tx.get_last(
+                timestamp=timestamp,
+                prefix=key,
+                retention_ms=retention_ms,
+            )
             if is_inner_join and not right_value:
                 return DISCARDED
             return merger(value, right_value)

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -2872,3 +2872,48 @@ def merger(left, right):
         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)]
+
+    @pytest.mark.parametrize(
+        "retention_ms, right_timestamp, left_timestamp, joined",
+        [
+            # Retention strategy includes right values greater or equal to 5 - 3 = 2
+            (3, 2, 5, True),
+            (timedelta(milliseconds=3), 2, 5, True),
+            # Retention strategy ignores right values lower than 6 - 3 = 3
+            (3, 2, 6, False),
+            (timedelta(milliseconds=3), 2, 6, False),
+        ],
+    )
+    def test_retention_ms(
+        self,
+        create_topic,
+        create_sdf,
+        assign_partition,
+        publish,
+        retention_ms,
+        right_timestamp,
+        left_timestamp,
+        joined,
+    ):
+        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, retention_ms=retention_ms)
+        assign_partition(right_sdf)
+
+        publish(
+            joined_sdf,
+            right_topic,
+            value={"right": 1},
+            key=b"key",
+            timestamp=right_timestamp,
+        )
+        joined_value = publish(
+            joined_sdf,
+            left_topic,
+            value={"left": 2},
+            key=b"key",
+            timestamp=left_timestamp,
+        )
+
+        assert bool(joined_value) == joined