Repository: narwhals-dev/narwhals
Pull request #924 (2024-09-07 23:31:29)

feat: implement `to_arrow` for cuDF
<!--
# 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 #  https://github.com/narwhals-dev/narwhals/issues/862
- Closes #

## Checklist

- [x] 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/docs/installation.md b/docs/installation.md
index 796cd870..b89aa4b6 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -11,6 +11,6 @@ Then, if you start the Python REPL and see the following:
 ```python
 >>> import narwhals
 >>> narwhals.__version__
-'1.6.2'
+'1.6.3'
 ```
 then installation worked correctly!
diff --git a/docs/why.md b/docs/why.md
index adf8f39b..4ec605d1 100644
--- a/docs/why.md
+++ b/docs/why.md
@@ -27,7 +27,7 @@ pl_df_right = pl.DataFrame({"a": [1, 2, 3], "c": [4, 5, 6]})
 pl_left_merge = pl_df_left.join(pl_df_right, left_on="b", right_on="c", how="left")
 
 print(pd_left_merge.columns)
-print(pl_df_right.columns)
+print(pl_left_merge.columns)
 ```
 
 There are several such subtle difference between the libraries. Writing dataframe-agnostic code is hard!
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 56c638e8..b26cf949 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -53,7 +53,7 @@ from narwhals.utils import maybe_convert_dtypes
 from narwhals.utils import maybe_get_index
 from narwhals.utils import maybe_set_index
 
-__version__ = "1.6.2"
+__version__ = "1.6.3"
 
 __all__ = [
     "dependencies",
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 960d833a..2750f8c0 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -136,9 +136,12 @@ class ArrowDataFrame:
             and len(item) == 2
             and isinstance(item[1], (list, tuple))
         ):
-            return self._from_native_frame(
-                self._native_frame.take(item[0]).select(item[1])
-            )
+            if item[0] == slice(None):
+                selected_rows = self._native_frame
+            else:
+                selected_rows = self._native_frame.take(item[0])
+
+            return self._from_native_frame(selected_rows.select(item[1]))
 
         elif isinstance(item, tuple) and len(item) == 2:
             if isinstance(item[1], slice):
@@ -275,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..ce0cae8a 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
@@ -677,8 +672,7 @@ class PandasLikeDataFrame:
 
     def to_arrow(self: Self) -> Any:
         if self._implementation is Implementation.CUDF:  # pragma: no cover
-            msg = "`to_arrow` is not implemented for CuDF backend."
-            raise NotImplementedError(msg)
+            return self._native_frame.to_arrow(preserve_index=False)
 
         import pyarrow as pa  # ignore-banned-import()
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index e94c95a8..b28a0408 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -635,8 +635,7 @@ class PandasLikeSeries:
 
     def to_arrow(self: Self) -> Any:
         if self._implementation is Implementation.CUDF:  # pragma: no cover
-            msg = "`to_arrow` is not implemented for CuDF backend."
-            raise NotImplementedError(msg)
+            return self._native_series.to_arrow()
 
         import pyarrow as pa  # ignore-banned-import()
 
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 165b6598..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."
+        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 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."
-            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]):
@@ -585,10 +596,14 @@ class DataFrame(BaseFrame[FrameT]):
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
     @overload
+    def __getitem__(self, item: tuple[slice, Sequence[int]]) -> Self: ...
+    @overload
     def __getitem__(self, item: tuple[Sequence[int], str]) -> Series: ...  # type: ignore[overload-overlap]
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[str]]) -> Self: ...
     @overload
+    def __getitem__(self, item: tuple[slice, Sequence[str]]) -> Self: ...
+    @overload
     def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ...  # type: ignore[overload-overlap]
 
     @overload
@@ -606,7 +621,7 @@ class DataFrame(BaseFrame[FrameT]):
         | slice
         | Sequence[int]
         | tuple[Sequence[int], str | int]
-        | tuple[Sequence[int], Sequence[int] | Sequence[str] | slice],
+        | tuple[slice | Sequence[int], Sequence[int] | Sequence[str] | slice],
     ) -> Series | Self:
         """
         Extract column or slice of DataFrame.
@@ -623,6 +638,10 @@ class DataFrame(BaseFrame[FrameT]):
                     a `Series`.
                 - `df[[0, 1], [0, 1, 2]]` extracts the first two rows and the first three columns
                     and returns a `DataFrame`
+                - `df[:, [0, 1, 2]]` extracts all rows from the first three columns and returns a
+                  `DataFrame`.
+                - `df[:, ['a', 'c']]` extracts all rows and columns `'a'` and `'c'` and returns a
+                  `DataFrame`.
                 - `df[0: 2, ['a', 'c']]` extracts the first two rows and columns `'a'` and `'c'` and
                     returns a `DataFrame`
                 - `df[:, 0: 2]` extracts all rows from the first two columns and returns a `DataFrame`
@@ -1835,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.
@@ -1853,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
 
@@ -1900,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,
@@ -3560,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.
@@ -3578,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
 
@@ -3625,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,
diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py
index 98ee12a7..1af7a26f 100644
--- a/narwhals/stable/v1.py
+++ b/narwhals/stable/v1.py
@@ -76,11 +76,15 @@ class DataFrame(NwDataFrame[IntoDataFrameT]):
     def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ...
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
+    @overload
+    def __getitem__(self, item: tuple[slice, Sequence[int]]) -> Self: ...
 
     @overload
     def __getitem__(self, item: tuple[Sequence[int], str]) -> Series: ...  # type: ignore[overload-overlap]
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[str]]) -> Self: ...
+    @overload
+    def __getitem__(self, item: tuple[slice, Sequence[str]]) -> Self: ...
 
     @overload
     def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ...  # type: ignore[overload-overlap]
diff --git a/pyproject.toml b/pyproject.toml
index c4a10603..a928ae0d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "narwhals"
-version = "1.6.2"
+version = "1.6.3"
 authors = [
   { name="Marco Gorelli", email="33491632+MarcoGorelli@users.noreply.github.com" },
 ]


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]
diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py
index eea94d44..18b05bf3 100644
--- a/tests/frame/slice_test.py
+++ b/tests/frame/slice_test.py
@@ -141,6 +141,12 @@ def test_slice_slice_columns(constructor_eager: Any) -> None:
     result = df[[0, 1], 1:]
     expected = {"b": [4, 5], "c": [7, 8], "d": [1, 4]}
     compare_dicts(result, expected)
+    result = df[:, ["b", "d"]]
+    expected = {"b": [4, 5, 6], "d": [1, 4, 2]}
+    compare_dicts(result, expected)
+    result = df[:, [0, 2]]
+    expected = {"a": [1, 2, 3], "c": [7, 8, 9]}
+    compare_dicts(result, expected)
 
 
 def test_slice_invalid(constructor_eager: Any) -> None: