Repository: narwhals-dev/narwhals
Pull request #927 (2024-09-08 11:43:05)

feat: add `on` key to `join`
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 064903d7..2750f8c0 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -278,11 +278,6 @@ class ArrowDataFrame:
         left_on: str | list[str] | None,
         right_on: str | list[str] | None,
     ) -> Self:
-        if isinstance(left_on, str):
-            left_on = [left_on]
-        if isinstance(right_on, str):
-            right_on = [right_on]
-
         how_to_join_map = {
             "anti": "left anti",
             "semi": "left semi",
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
index 5ef8c5a9..1a40d7a6 100644
--- a/narwhals/_dask/dataframe.py
+++ b/narwhals/_dask/dataframe.py
@@ -209,11 +209,6 @@ class DaskLazyFrame:
         left_on: str | list[str] | None,
         right_on: str | list[str] | None,
     ) -> Self:
-        if isinstance(left_on, str):
-            left_on = [left_on]
-        if isinstance(right_on, str):
-            right_on = [right_on]
-
         if how == "cross":
             key_token = generate_unique_token(
                 n_bytes=8, columns=[*self.columns, *other.columns]
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 3040adda..880e2d14 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -404,11 +404,6 @@ class PandasLikeDataFrame:
         left_on: str | list[str] | None,
         right_on: str | list[str] | None,
     ) -> Self:
-        if isinstance(left_on, str):
-            left_on = [left_on]
-        if isinstance(right_on, str):
-            right_on = [right_on]
-
         if how == "cross":
             if (
                 self._implementation is Implementation.MODIN
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index f6fccb27..ffd7ce36 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -186,6 +186,7 @@ class BaseFrame(Generic[FrameT]):
         how: Literal["inner", "left", "cross", "semi", "anti"] = "inner",
         left_on: str | list[str] | None = None,
         right_on: str | list[str] | None = None,
+        on: str | list[str] | None = None,
     ) -> Self:
         _supported_joins = ("inner", "left", "cross", "anti", "semi")
 
@@ -193,10 +194,25 @@ class BaseFrame(Generic[FrameT]):
             msg = f"Only the following join strategies are supported: {_supported_joins}; found '{how}'."
             raise NotImplementedError(msg)
 
-        if how == "cross" and (left_on or right_on):
-            msg = "Can not pass left_on, right_on for cross join"
+        if how == "cross" and (
+            left_on is not None or right_on is not None or on is not None
+        ):
+            msg = "Can not pass `left_on`, `right_on` or `on` keys for cross join"
+            raise ValueError(msg)
+
+        if how != "cross" and (on is None and (left_on is None or right_on is None)):
+            msg = f"Either (`left_on` and `right_on`) or `on` keys should be specified for {how}."
             raise ValueError(msg)
 
+        if how != "cross" and (
+            on is not None and (left_on is not None or right_on is not None)
+        ):
+            msg = f"If `on` is specified, `left_on` and `right_on` should be None for {how}."
+            raise ValueError(msg)
+
+        if on is not None:
+            left_on = right_on = on
+
         return self._from_compliant_dataframe(
             self._compliant_frame.join(
                 self._extract_compliant(other),
@@ -232,37 +248,24 @@ class BaseFrame(Generic[FrameT]):
             msg = f"Only the following strategies are supported: {_supported_strategies}; found '{strategy}'."
             raise NotImplementedError(msg)
 
-        if left_on is not None and right_on is not None and on is not None:
+        if (on is None) and (left_on is None or right_on is None):
             msg = "Either (`left_on` and `right_on`) or `on` keys should be specified."
             raise ValueError(msg)
-        if by_left is not None and by_right is not None and by is not None:
-            msg = "Can not specify `by_left`, `by_right`, and `by` keys at the same time."
-            raise ValueError(msg)
-        if by_left is not None and by_right is None and by is None:
-            msg = "`by_right` can not be None if `by_left` is specified."
-            raise ValueError(msg)
-        if by_left is None and by_right is not None and by is None:
-            msg = "`by_left` can not be None if `by_right` is specified."
+        if (on is not None) and (left_on is not None or right_on is not None):
+            msg = "If `on` is specified, `left_on` and `right_on` should be None."
             raise ValueError(msg)
-        if (
+        if (by is None) and (
             (by_left is None and by_right is not None)
             or (by_left is not None and by_right is None)
-        ) and by is not None:
-            msg = "Either (`by_left` and `by_right_`) or `by` keys should be specified."
-            raise ValueError(msg)
-        if left_on is not None and right_on is not None:
-            return self._from_compliant_dataframe(
-                self._compliant_frame.join_asof(
-                    self._extract_compliant(other),
-                    left_on=left_on,
-                    right_on=right_on,
-                    by_left=by_left,
-                    by_right=by_right,
-                    by=by,
-                    strategy=strategy,
-                )
+        ):
+            msg = (
+                "Can not specify only `by_left` or `by_right`, you need to specify both."
             )
-        elif on is not None:
+            raise ValueError(msg)
+        if (by is not None) and (by_left is not None or by_right is not None):
+            msg = "If `by` is specified, `by_left` and `by_right` should be None."
+            raise ValueError(msg)
+        if on is not None:
             return self._from_compliant_dataframe(
                 self._compliant_frame.join_asof(
                     self._extract_compliant(other),
@@ -273,9 +276,17 @@ class BaseFrame(Generic[FrameT]):
                     strategy=strategy,
                 )
             )
-        else:
-            msg = "Either (`left_on` and `right_on`) or `on` keys should be specified."
-            raise ValueError(msg)
+        return self._from_compliant_dataframe(
+            self._compliant_frame.join_asof(
+                self._extract_compliant(other),
+                left_on=left_on,
+                right_on=right_on,
+                by_left=by_left,
+                by_right=by_right,
+                by=by,
+                strategy=strategy,
+            )
+        )
 
 
 class DataFrame(BaseFrame[FrameT]):
@@ -1843,6 +1854,7 @@ class DataFrame(BaseFrame[FrameT]):
         how: Literal["inner", "left", "cross", "semi", "anti"] = "inner",
         left_on: str | list[str] | None = None,
         right_on: str | list[str] | None = None,
+        on: str | list[str] | None = None,
     ) -> Self:
         r"""
         Join in SQL-like fashion.
@@ -1861,6 +1873,8 @@ class DataFrame(BaseFrame[FrameT]):
 
             right_on: Name(s) of the right join column(s).
 
+            on: Join column of both DataFrames. If set, left_on and right_on should be None.
+
         Returns:
             A new joined DataFrame
 
@@ -1908,7 +1922,7 @@ class DataFrame(BaseFrame[FrameT]):
             │ 2   ┆ 7.0 ┆ b   ┆ y     │
             └─────┴─────┴─────┴───────┘
         """
-        return super().join(other, how=how, left_on=left_on, right_on=right_on)
+        return super().join(other, how=how, left_on=left_on, right_on=right_on, on=on)
 
     def join_asof(
         self,
@@ -3568,6 +3582,7 @@ class LazyFrame(BaseFrame[FrameT]):
         how: Literal["inner", "left", "cross", "semi", "anti"] = "inner",
         left_on: str | list[str] | None = None,
         right_on: str | list[str] | None = None,
+        on: str | list[str] | None = None,
     ) -> Self:
         r"""
         Add a join operation to the Logical Plan.
@@ -3586,6 +3601,8 @@ class LazyFrame(BaseFrame[FrameT]):
 
             right_on: Join column of the right DataFrame.
 
+            on: Join column of both DataFrames. If set, left_on and right_on should be None.
+
         Returns:
             A new joined LazyFrame
 
@@ -3633,7 +3650,7 @@ class LazyFrame(BaseFrame[FrameT]):
             │ 2   ┆ 7.0 ┆ b   ┆ y     │
             └─────┴─────┴─────┴───────┘
         """
-        return super().join(other, how=how, left_on=left_on, right_on=right_on)
+        return super().join(other, how=how, left_on=left_on, right_on=right_on, on=on)
 
     def join_asof(
         self,


Test patch:
diff --git a/tests/frame/join_test.py b/tests/frame/join_test.py
index 34a5961e..6615d503 100644
--- a/tests/frame/join_test.py
+++ b/tests/frame/join_test.py
@@ -18,8 +18,9 @@ def test_inner_join_two_keys(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     df_right = df
     result = df.join(df_right, left_on=["a", "b"], right_on=["a", "b"], how="inner")  # type: ignore[arg-type]
-    result = result.sort("index")
-    result = result.drop("index_right")
+    result_on = df.join(df_right, on=["a", "b"], how="inner")  # type: ignore[arg-type]
+    result = result.sort("index").drop("index_right")
+    result_on = result_on.sort("index").drop("index_right")
     expected = {
         "a": [1, 3, 2],
         "b": [4, 4, 6],
@@ -28,6 +29,7 @@ def test_inner_join_two_keys(constructor: Any) -> None:
         "index": [0, 1, 2],
     }
     compare_dicts(result, expected)
+    compare_dicts(result_on, expected)
 
 
 def test_inner_join_single_key(constructor: Any) -> None:
@@ -35,7 +37,9 @@ def test_inner_join_single_key(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     df_right = df
     result = df.join(df_right, left_on="a", right_on="a", how="inner").sort("index")  # type: ignore[arg-type]
+    result_on = df.join(df_right, on="a", how="inner").sort("index")  # type: ignore[arg-type]
     result = result.drop("index_right")
+    result_on = result_on.drop("index_right")
     expected = {
         "a": [1, 3, 2],
         "b": [4, 4, 6],
@@ -45,6 +49,7 @@ def test_inner_join_single_key(constructor: Any) -> None:
         "index": [0, 1, 2],
     }
     compare_dicts(result, expected)
+    compare_dicts(result_on, expected)
 
 
 def test_cross_join(constructor: Any) -> None:
@@ -57,7 +62,9 @@ def test_cross_join(constructor: Any) -> None:
     }
     compare_dicts(result, expected)
 
-    with pytest.raises(ValueError, match="Can not pass left_on, right_on for cross join"):
+    with pytest.raises(
+        ValueError, match="Can not pass `left_on`, `right_on` or `on` keys for cross join"
+    ):
         df.join(df, how="cross", left_on="a")  # type: ignore[arg-type]
 
 
@@ -206,6 +213,33 @@ def test_left_join_overlapping_column(constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
+@pytest.mark.parametrize("how", ["inner", "left", "semi", "anti"])
+def test_join_keys_exceptions(constructor: Any, how: str) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+    df = nw.from_native(constructor(data))
+
+    with pytest.raises(
+        ValueError,
+        match=rf"Either \(`left_on` and `right_on`\) or `on` keys should be specified for {how}.",
+    ):
+        df.join(df, how=how)  # type: ignore[arg-type]
+    with pytest.raises(
+        ValueError,
+        match=rf"Either \(`left_on` and `right_on`\) or `on` keys should be specified for {how}.",
+    ):
+        df.join(df, how=how, left_on="a")  # type: ignore[arg-type]
+    with pytest.raises(
+        ValueError,
+        match=rf"Either \(`left_on` and `right_on`\) or `on` keys should be specified for {how}.",
+    ):
+        df.join(df, how=how, right_on="a")  # type: ignore[arg-type]
+    with pytest.raises(
+        ValueError,
+        match=f"If `on` is specified, `left_on` and `right_on` should be None for {how}.",
+    ):
+        df.join(df, how=how, on="a", right_on="a")  # type: ignore[arg-type]
+
+
 def test_joinasof_numeric(constructor: Any, request: Any) -> None:
     if "pyarrow_table" in str(constructor):
         request.applymarker(pytest.mark.xfail)
@@ -361,31 +395,40 @@ def test_joinasof_not_implemented(constructor: Any, strategy: str) -> None:
         df.join_asof(df, left_on="a", right_on="a", strategy=strategy)  # type: ignore[arg-type]
 
 
-def test_joinasof_no_keys(constructor: Any) -> None:
+def test_joinasof_keys_exceptions(constructor: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
     df = nw.from_native(constructor(data))
 
-    msg = r"Either \(`left_on` and `right_on`\) or `on` keys should be specified."
     with pytest.raises(
         ValueError,
-        match=msg,
+        match=r"Either \(`left_on` and `right_on`\) or `on` keys should be specified.",
     ):
         df.join_asof(df, left_on="a")  # type: ignore[arg-type]
     with pytest.raises(
         ValueError,
-        match=msg,
+        match=r"Either \(`left_on` and `right_on`\) or `on` keys should be specified.",
     ):
         df.join_asof(df, right_on="a")  # type: ignore[arg-type]
     with pytest.raises(
         ValueError,
-        match=msg,
+        match=r"Either \(`left_on` and `right_on`\) or `on` keys should be specified.",
     ):
         df.join_asof(df)  # type: ignore[arg-type]
     with pytest.raises(
         ValueError,
-        match=msg,
+        match="If `on` is specified, `left_on` and `right_on` should be None.",
     ):
         df.join_asof(df, left_on="a", right_on="a", on="a")  # type: ignore[arg-type]
+    with pytest.raises(
+        ValueError,
+        match="If `on` is specified, `left_on` and `right_on` should be None.",
+    ):
+        df.join_asof(df, left_on="a", on="a")  # type: ignore[arg-type]
+    with pytest.raises(
+        ValueError,
+        match="If `on` is specified, `left_on` and `right_on` should be None.",
+    ):
+        df.join_asof(df, right_on="a", on="a")  # type: ignore[arg-type]
 
 
 def test_joinasof_by_exceptions(constructor: Any) -> None:
@@ -393,30 +436,30 @@ def test_joinasof_by_exceptions(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     with pytest.raises(
         ValueError,
-        match=r"Can not specify `by_left`, `by_right`, and `by` keys at the same time.",
+        match="If `by` is specified, `by_left` and `by_right` should be None.",
     ):
         df.join_asof(df, on="a", by_left="b", by_right="b", by="b")  # type: ignore[arg-type]
 
     with pytest.raises(
         ValueError,
-        match=r"`by_right` can not be None if `by_left` is specified.",
+        match="Can not specify only `by_left` or `by_right`, you need to specify both.",
     ):
         df.join_asof(df, on="a", by_left="b")  # type: ignore[arg-type]
 
     with pytest.raises(
         ValueError,
-        match=r"`by_left` can not be None if `by_right` is specified.",
+        match="Can not specify only `by_left` or `by_right`, you need to specify both.",
     ):
         df.join_asof(df, on="a", by_right="b")  # type: ignore[arg-type]
 
     with pytest.raises(
         ValueError,
-        match=r"Either \(`by_left` and `by_right_`\) or `by` keys should be specified.",
+        match="If `by` is specified, `by_left` and `by_right` should be None.",
     ):
         df.join_asof(df, on="a", by_left="b", by="b")  # type: ignore[arg-type]
 
     with pytest.raises(
         ValueError,
-        match=r"Either \(`by_left` and `by_right_`\) or `by` keys should be specified.",
+        match="If `by` is specified, `by_left` and `by_right` should be None.",
     ):
         df.join_asof(df, on="a", by_right="b", by="b")  # type: ignore[arg-type]