Ensure both sides are copartitioned

Author: gwaramadze

quixstreams/dataframe/dataframe.py

@@ -1648,7 +1648,7 @@ def concat(self, other: "StreamingDataFrame") -> "StreamingDataFrame":
         )
 
     def join(self, right: "StreamingDataFrame") -> "StreamingDataFrame":
-        # TODO: ensure copartitioning of left and right?
+        self.ensure_topics_copartitioned(*self.topics, *right.topics)
         right.register_store(store_type=TimestampedStore)
 
         def left_func(value, key, timestamp, headers):
@@ -1664,12 +1664,13 @@ def right_func(value, key, timestamp, headers):
         right = right.update(right_func, metadata=True).filter(lambda value: False)
         return left.concat(right)
 
-    def ensure_topics_copartitioned(self):
-        partitions_counts = set(t.broker_config.num_partitions for t in self._topics)
+    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}"

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -2703,3 +2703,10 @@ def test_join(self, create_topic, create_sdf, assign_partition, publish):
             joined_sdf, left_topic, value={"left": 2}, key=b"key", timestamp=2
         )
         assert joined_value == [({"left": 2, "right": 1}, b"key", 2, None)]
+
+    def test_mismatching_partitions_fails(self, create_topic, create_sdf):
+        left_topic, right_topic = create_topic(), create_topic(num_partitions=2)
+        left_sdf, right_sdf = create_sdf(left_topic), create_sdf(right_topic)
+
+        with pytest.raises(TopicPartitionsMismatch):
+            left_sdf.join(right_sdf)