Repository: narwhals-dev/narwhals
Pull request #507 (2024-07-13 11:24:21)

feat: implement inner join for arrow
<!--
# 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

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] 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 768fcb32..b3ef89fc 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
+from typing import Literal
 from typing import Sequence
 from typing import overload
 
@@ -183,6 +184,41 @@ class ArrowDataFrame:
         df = self._native_dataframe.__class__.from_arrays(to_concat, names=output_names)
         return self._from_native_dataframe(df)
 
+    def join(
+        self,
+        other: Self,
+        *,
+        how: Literal["inner"] = "inner",
+        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":  # type: ignore[comparison-overlap]
+            raise NotImplementedError
+
+        if how == "anti":  # type: ignore[comparison-overlap]
+            raise NotImplementedError
+
+        if how == "semi":  # type: ignore[comparison-overlap]
+            raise NotImplementedError
+
+        if how == "left":  # type: ignore[comparison-overlap]
+            raise NotImplementedError
+
+        return self._from_native_dataframe(
+            self._native_dataframe.join(
+                other._native_dataframe,
+                keys=left_on,
+                right_keys=right_on,
+                join_type=how,
+                right_suffix="_right",
+            ),
+        )
+
     def drop(self, *columns: str | Iterable[str]) -> Self:
         return self._from_native_dataframe(
             self._native_dataframe.drop(list(flatten(columns)))
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index 2ca15a69..0cd6335b 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -18,7 +18,6 @@ MISSING = [
     "DataFrame.is_unique",
     "DataFrame.item",
     "DataFrame.iter_rows",
-    "DataFrame.join",
     "DataFrame.pipe",
     "DataFrame.rename",
     "DataFrame.unique",


Test patch:
diff --git a/tests/frame/join_test.py b/tests/frame/join_test.py
index c248a081..ebe7bacc 100644
--- a/tests/frame/join_test.py
+++ b/tests/frame/join_test.py
@@ -11,21 +11,20 @@ from narwhals._pandas_like.utils import Implementation
 from tests.utils import compare_dicts
 
 
-def test_inner_join_two_keys(constructor: Any) -> None:
+def test_inner_join_two_keys(constructor_with_pyarrow: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
-    df = nw.from_native(constructor(data), eager_only=True)
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
     df_right = df
     result = df.join(df_right, left_on=["a", "b"], right_on=["a", "b"], how="inner")
     expected = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9], "z_right": [7.0, 8, 9]}
     compare_dicts(result, expected)
 
 
-def test_inner_join_single_key(constructor: Any) -> None:
+def test_inner_join_single_key(constructor_with_pyarrow: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
-    df = nw.from_native(constructor(data), eager_only=True)
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
     df_right = df
     result = df.join(df_right, left_on="a", right_on="a", how="inner")
-    result_native = nw.to_native(result)
     expected = {
         "a": [1, 3, 2],
         "b": [4, 4, 6],
@@ -33,7 +32,7 @@ def test_inner_join_single_key(constructor: Any) -> None:
         "z": [7.0, 8, 9],
         "z_right": [7.0, 8, 9],
     }
-    compare_dicts(result_native, expected)
+    compare_dicts(result, expected)
 
 
 def test_cross_join(constructor: Any) -> None: