chore: Simplify PandasLikeDataFrame|DaskLazyFrame.join method

narwhals/_compliant/dataframe.py

@@ -331,7 +331,7 @@ def join(
         self,
         other: Self,
         *,
-        how: Literal["left", "inner", "cross", "anti", "semi"],
+        how: JoinStrategy,
         left_on: Sequence[str] | None,
         right_on: Sequence[str] | None,
         suffix: str,

narwhals/_dask/dataframe.py

@@ -8,6 +8,7 @@
 from narwhals._dask.utils import add_row_index, evaluate_exprs
 from narwhals._expression_parsing import ExprKind
 from narwhals._pandas_like.utils import native_to_narwhals_dtype, select_columns_by_name
+from narwhals._typing_compat import assert_never
 from narwhals._utils import (
     Implementation,
     _remap_full_join_keys,
@@ -271,7 +272,123 @@ def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) ->
             self.native.sort_values(list(by), ascending=ascending, na_position=position)
         )
 
-    def join(  # noqa: C901
+    def _join_inner(
+        self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str
+    ) -> dd.DataFrame:
+        return self.native.merge(
+            other.native,
+            left_on=left_on,
+            right_on=right_on,
+            how="inner",
+            suffixes=("", suffix),
+        )
+
+    def _join_left(
+        self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str
+    ) -> dd.DataFrame:
+        result_native = self.native.merge(
+            other.native,
+            how="left",
+            left_on=left_on,
+            right_on=right_on,
+            suffixes=("", suffix),
+        )
+        extra = [
+            right_key if right_key not in self.columns else f"{right_key}{suffix}"
+            for left_key, right_key in zip(left_on, right_on)
+            if right_key != left_key
+        ]
+        return result_native.drop(columns=extra)
+
+    def _join_full(
+        self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str
+    ) -> dd.DataFrame:
+        # dask does not retain keys post-join
+        # we must append the suffix to each key before-hand
+
+        right_on_mapper = _remap_full_join_keys(left_on, right_on, suffix)
+        other_native = other.native.rename(columns=right_on_mapper)
+        check_column_names_are_unique(other_native.columns)
+        right_suffixed = list(right_on_mapper.values())
+        return self.native.merge(
+            other_native,
+            left_on=left_on,
+            right_on=right_suffixed,
+            how="outer",
+            suffixes=("", suffix),
+        )
+
+    def _join_cross(self, other: Self, *, suffix: str) -> dd.DataFrame:
+        key_token = generate_temporary_column_name(
+            n_bytes=8, columns=(*self.columns, *other.columns)
+        )
+        return (
+            self.native.assign(**{key_token: 0})
+            .merge(
+                other.native.assign(**{key_token: 0}),
+                how="inner",
+                left_on=key_token,
+                right_on=key_token,
+                suffixes=("", suffix),
+            )
+            .drop(columns=key_token)
+        )
+
+    def _join_semi(
+        self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str]
+    ) -> dd.DataFrame:
+        other_native = self._join_filter_rename(
+            other=other,
+            columns_to_select=list(right_on),
+            columns_mapping=dict(zip(right_on, left_on)),
+        )
+        return self.native.merge(
+            other_native, how="inner", left_on=left_on, right_on=left_on
+        )
+
+    def _join_anti(
+        self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str]
+    ) -> dd.DataFrame:
+        indicator_token = generate_temporary_column_name(
+            n_bytes=8, columns=(*self.columns, *other.columns)
+        )
+        other_native = self._join_filter_rename(
+            other=other,
+            columns_to_select=list(right_on),
+            columns_mapping=dict(zip(right_on, left_on)),
+        )
+        df = self.native.merge(
+            other_native,
+            how="left",
+            indicator=indicator_token,  # pyright: ignore[reportArgumentType]
+            left_on=left_on,
+            right_on=left_on,
+        )
+        return df[df[indicator_token] == "left_only"].drop(columns=[indicator_token])
+
+    def _join_filter_rename(
+        self, other: Self, columns_to_select: list[str], columns_mapping: dict[str, str]
+    ) -> dd.DataFrame:
+        """Helper function to avoid creating extra columns and row duplication.
+
+        Used in `"anti"` and `"semi`" join's.
+
+        Notice that a native object is returned.
+        """
+        other_native: Incomplete = other.native
+        return (
+            select_columns_by_name(
+                other_native,
+                column_names=columns_to_select,
+                backend_version=self._backend_version,
+                implementation=self._implementation,
+            )
+            # rename to avoid creating extra columns in join
+            .rename(columns=columns_mapping)
+            .drop_duplicates()
+        )
+
+    def join(
         self,
         other: Self,
         *,
@@ -281,123 +398,30 @@ def join(  # noqa: C901
         suffix: str,
     ) -> Self:
         if how == "cross":
-            key_token = generate_temporary_column_name(
-                n_bytes=8, columns=[*self.columns, *other.columns]
-            )
-
-            return self._with_native(
-                self.native.assign(**{key_token: 0})
-                .merge(
-                    other.native.assign(**{key_token: 0}),
-                    how="inner",
-                    left_on=key_token,
-                    right_on=key_token,
-                    suffixes=("", suffix),
-                )
-                .drop(columns=key_token)
-            )
-        other_native: Incomplete = other.native
+            result = self._join_cross(other=other, suffix=suffix)
 
-        if how == "anti":
-            indicator_token = generate_temporary_column_name(
-                n_bytes=8, columns=[*self.columns, *other.columns]
-            )
-
-            if right_on is None:  # pragma: no cover
-                msg = "`right_on` cannot be `None` in anti-join"
-                raise TypeError(msg)
-            other_native = (
-                select_columns_by_name(
-                    other_native,
-                    list(right_on),
-                    self._backend_version,
-                    self._implementation,
-                )
-                .rename(  # rename to avoid creating extra columns in join
-                    columns=dict(zip(right_on, left_on))  # type: ignore[arg-type]
-                )
-                .drop_duplicates()
-            )
-            df = self.native.merge(
-                other_native,
-                how="outer",
-                indicator=indicator_token,  # pyright: ignore[reportArgumentType]
-                left_on=left_on,
-                right_on=left_on,
-            )
-            return self._with_native(
-                df[df[indicator_token] == "left_only"].drop(columns=[indicator_token])
-            )
-
-        if how == "semi":
-            if right_on is None:  # pragma: no cover
-                msg = "`right_on` cannot be `None` in semi-join"
-                raise TypeError(msg)
-            other_native = (
-                select_columns_by_name(
-                    other_native,
-                    list(right_on),
-                    self._backend_version,
-                    self._implementation,
-                )
-                .rename(  # rename to avoid creating extra columns in join
-                    columns=dict(zip(right_on, left_on))  # type: ignore[arg-type]
-                )
-                .drop_duplicates()  # avoids potential rows duplication from inner join
-            )
-            return self._with_native(
-                self.native.merge(
-                    other_native, how="inner", left_on=left_on, right_on=left_on
-                )
-            )
+        elif left_on is None or right_on is None:  # pragma: no cover
+            raise ValueError(left_on, right_on)
 
-        if how == "left":
-            result_native = self.native.merge(
-                other.native,
-                how="left",
-                left_on=left_on,
-                right_on=right_on,
-                suffixes=("", suffix),
+        elif how == "inner":
+            result = self._join_inner(
+                other=other, left_on=left_on, right_on=right_on, suffix=suffix
             )
-            extra = []
-            for left_key, right_key in zip(left_on, right_on):  # type: ignore[arg-type]
-                if right_key != left_key and right_key not in self.columns:
-                    extra.append(right_key)
-                elif right_key != left_key:
-                    extra.append(f"{right_key}_right")
-            return self._with_native(result_native.drop(columns=extra))
-
-        if how == "full":
-            # dask does not retain keys post-join
-            # we must append the suffix to each key before-hand
-
-            # help mypy
-            assert left_on is not None  # noqa: S101
-            assert right_on is not None  # noqa: S101
-
-            right_on_mapper = _remap_full_join_keys(left_on, right_on, suffix)
-            other_native = other.native.rename(columns=right_on_mapper)
-            check_column_names_are_unique(other_native.columns)
-            right_on = list(right_on_mapper.values())  # we now have the suffixed keys
-            return self._with_native(
-                self.native.merge(
-                    other_native,
-                    left_on=left_on,
-                    right_on=right_on,
-                    how="outer",
-                    suffixes=("", suffix),
-                )
+        elif how == "anti":
+            result = self._join_anti(other=other, left_on=left_on, right_on=right_on)
+        elif how == "semi":
+            result = self._join_semi(other=other, left_on=left_on, right_on=right_on)
+        elif how == "left":
+            result = self._join_left(
+                other=other, left_on=left_on, right_on=right_on, suffix=suffix
             )
-
-        return self._with_native(
-            self.native.merge(
-                other.native,
-                left_on=left_on,
-                right_on=right_on,
-                how=how,
-                suffixes=("", suffix),
+        elif how == "full":
+            result = self._join_full(
+                other=other, left_on=left_on, right_on=right_on, suffix=suffix
             )
-        )
+        else:
+            assert_never(how)
+        return self._with_native(result)
 
     def join_asof(
         self,