Correct after rebase

Author: gwaramadze

quixstreams/dataframe/dataframe.py

@@ -1666,18 +1666,13 @@ def join(
                 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"
-        retention_ms = ensure_milliseconds(retention_ms)
-
         if merger is None:
             if on_overlap == "keep-left":
 
@@ -1710,20 +1705,26 @@ def merger(left_value, right_value):
                         )
                     return {**left_value, **right_value}
 
+        self.ensure_topics_copartitioned(*self.topics, *right.topics)
+        right.register_store(store_type=TimestampedStore)
+        is_inner_join = how == "inner"
+        retention_ms = ensure_milliseconds(retention_ms)
+
         def left_func(value, key, timestamp, headers):
             right_tx = _get_transaction(right)
-            right_value = right_tx.get_last(
-                timestamp=timestamp,
-                prefix=key,
-                retention_ms=retention_ms,
-            )
+            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)
+            right_tx.set_for_timestamp(
+                timestamp=timestamp,
+                value=value,
+                prefix=key,
+                retention_ms=retention_ms,
+            )
 
         left = self.apply(left_func, metadata=True).filter(
             lambda value: value is not DISCARDED

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -4,6 +4,7 @@
 import warnings
 from collections import namedtuple
 from datetime import timedelta
+from functools import partial
 from typing import Any, get_args
 from unittest import mock
 
@@ -2873,47 +2874,33 @@ def merger(left, right):
         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)
+        joined_sdf = left_sdf.join(right_sdf, retention_ms=10)
         assign_partition(right_sdf)
 
-        publish(
-            joined_sdf,
-            right_topic,
-            value={"right": 1},
-            key=b"key",
-            timestamp=right_timestamp,
-        )
-        joined_value = publish(
+        # min eligible timestamp is 15 - 10 = 5
+        publish(joined_sdf, right_topic, value={"right": 1}, key=b"key", timestamp=15)
+
+        # min eligible timestamp is still 5
+        publish(joined_sdf, right_topic, value={"right": 3}, key=b"key", timestamp=4)
+        publish(joined_sdf, right_topic, value={"right": 2}, key=b"key", timestamp=5)
+
+        publish_left = partial(
+            publish,
             joined_sdf,
             left_topic,
-            value={"left": 2},
+            value={"left": 4},
             key=b"key",
-            timestamp=left_timestamp,
         )
 
-        assert bool(joined_value) == joined
+        assert publish_left(timestamp=4) == []
+        assert publish_left(timestamp=5) == [({"left": 4, "right": 2}, b"key", 5, None)]