Repository: narwhals-dev/narwhals
Pull request #810 (2024-08-18 21:06:21)

feat: `DaskExpr.over` method
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #637 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [ ] Documented the changes




Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index b251b2a5..08ad0f7c 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -39,6 +39,7 @@
         - unique
         - with_columns
         - with_row_index
+        - write_csv
         - write_parquet
       show_source: false
       show_bases: false
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 865d1709..b5a14f5a 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -20,6 +20,7 @@ from narwhals.utils import generate_unique_token
 from narwhals.utils import parse_columns_to_drop
 
 if TYPE_CHECKING:
+    import numpy as np
     from typing_extensions import Self
 
     from narwhals._arrow.group_by import ArrowGroupBy
@@ -100,6 +101,9 @@ class ArrowDataFrame:
             backend_version=self._backend_version,
         )
 
+    def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray:
+        return self._native_frame.__array__(dtype, copy=copy)
+
     @overload
     def __getitem__(self, item: tuple[Sequence[int], str | int]) -> ArrowSeries: ...  # type: ignore[overload-overlap]
 
@@ -420,6 +424,17 @@ class ArrowDataFrame:
 
         pp.write_table(self._native_frame, file)
 
+    def write_csv(self, file: Any) -> Any:
+        import pyarrow as pa  # ignore-banned-import
+        import pyarrow.csv as pa_csv  # ignore-banned-import
+
+        pa_table = self._native_frame
+        if file is None:
+            csv_buffer = pa.BufferOutputStream()
+            pa_csv.write_csv(pa_table, csv_buffer)
+            return csv_buffer.getvalue().to_pybytes().decode()
+        return pa_csv.write_csv(pa_table, file)
+
     def is_duplicated(self: Self) -> ArrowSeries:
         import numpy as np  # ignore-banned-import
         import pyarrow as pa  # ignore-banned-import()
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
index faedb609..db0641a4 100644
--- a/narwhals/_dask/expr.py
+++ b/narwhals/_dask/expr.py
@@ -598,6 +598,33 @@ class DaskExpr:
         msg = "`Expr.gather_every` is not supported for the Dask backend. Please use `LazyFrame.gather_every` instead."
         raise NotImplementedError(msg)
 
+    def over(self: Self, keys: list[str]) -> Self:
+        def func(df: DaskLazyFrame) -> list[Any]:
+            if self._output_names is None:
+                msg = (
+                    "Anonymous expressions are not supported in over.\n"
+                    "Instead of `nw.all()`, try using a named expression, such as "
+                    "`nw.col('a', 'b')`\n"
+                )
+                raise ValueError(msg)
+            tmp = df.group_by(*keys).agg(self)
+            tmp = (
+                df.select(*keys)
+                .join(tmp, how="left", left_on=keys, right_on=keys)
+                ._native_frame
+            )
+            return [tmp[name] for name in self._output_names]
+
+        return self.__class__(
+            func,
+            depth=self._depth + 1,
+            function_name=self._function_name + "->over",
+            root_names=self._root_names,
+            output_names=self._output_names,
+            returns_scalar=False,
+            backend_version=self._backend_version,
+        )
+
     @property
     def str(self: Self) -> DaskExprStringNamespace:
         return DaskExprStringNamespace(self)
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 4faf61a0..dbe1999b 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -25,6 +25,7 @@ from narwhals.utils import generate_unique_token
 from narwhals.utils import parse_columns_to_drop
 
 if TYPE_CHECKING:
+    import numpy as np
     import pandas as pd
     from typing_extensions import Self
 
@@ -100,6 +101,9 @@ class PandasLikeDataFrame:
             backend_version=self._backend_version,
         )
 
+    def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray:
+        return self.to_numpy(dtype=dtype, copy=copy)
+
     @overload
     def __getitem__(self, item: tuple[Sequence[int], str | int]) -> PandasLikeSeries: ...  # type: ignore[overload-overlap]
 
@@ -520,19 +524,24 @@ class PandasLikeDataFrame:
             }
         return self._native_frame.to_dict(orient="list")  # type: ignore[no-any-return]
 
-    def to_numpy(self) -> Any:
+    def to_numpy(self, dtype: Any = None, copy: bool | None = None) -> Any:
         from narwhals._pandas_like.series import PANDAS_TO_NUMPY_DTYPE_MISSING
 
-        # pandas return `object` dtype for nullable dtypes, so we cast each
-        # Series to numpy and let numpy find a common dtype.
+        if dtype is not None:
+            return self._native_frame.to_numpy(dtype=dtype, copy=copy)
+
+        # pandas return `object` dtype for nullable dtypes if dtype=None,
+        # so we cast each Series to numpy and let numpy find a common dtype.
         # If there aren't any dtypes where `to_numpy()` is "broken" (i.e. it
         # returns Object) then we just call `to_numpy()` on the DataFrame.
-        for dtype in self._native_frame.dtypes:
-            if str(dtype) in PANDAS_TO_NUMPY_DTYPE_MISSING:
+        for col_dtype in self._native_frame.dtypes:
+            if str(col_dtype) in PANDAS_TO_NUMPY_DTYPE_MISSING:
                 import numpy as np  # ignore-banned-import
 
-                return np.hstack([self[col].to_numpy()[:, None] for col in self.columns])
-        return self._native_frame.to_numpy()
+                return np.hstack(
+                    [self[col].to_numpy(copy=copy)[:, None] for col in self.columns]
+                )
+        return self._native_frame.to_numpy(copy=copy)
 
     def to_pandas(self) -> Any:
         if self._implementation is Implementation.PANDAS:
@@ -544,6 +553,9 @@ class PandasLikeDataFrame:
     def write_parquet(self, file: Any) -> Any:
         self._native_frame.to_parquet(file)
 
+    def write_csv(self, file: Any = None) -> Any:
+        return self._native_frame.to_csv(file, index=False)
+
     # --- descriptive ---
     def is_duplicated(self: Self) -> PandasLikeSeries:
         from narwhals._pandas_like.series import PandasLikeSeries
diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py
index 2137f891..c2b771f9 100644
--- a/narwhals/_polars/dataframe.py
+++ b/narwhals/_polars/dataframe.py
@@ -11,6 +11,7 @@ from narwhals.utils import Implementation
 from narwhals.utils import parse_columns_to_drop
 
 if TYPE_CHECKING:
+    import numpy as np
     from typing_extensions import Self
 
 
@@ -58,6 +59,14 @@ class PolarsDataFrame:
 
         return func
 
+    def __array__(self, dtype: Any | None = None, copy: bool | None = None) -> np.ndarray:
+        if self._backend_version < (0, 20, 28) and copy is not None:  # pragma: no cover
+            msg = "`copy` in `__array__` is only supported for Polars>=0.20.28"
+            raise NotImplementedError(msg)
+        if self._backend_version < (0, 20, 28):  # pragma: no cover
+            return self._native_frame.__array__(dtype)
+        return self._native_frame.__array__(dtype)
+
     @property
     def schema(self) -> dict[str, Any]:
         schema = self._native_frame.schema
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 9aa7bb64..4fa7e913 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -233,8 +233,8 @@ class DataFrame(BaseFrame[FrameT]):
             msg = f"Expected an object which implements `__narwhals_dataframe__`, got: {type(df)}"
             raise AssertionError(msg)
 
-    def __array__(self) -> np.ndarray:
-        return self._compliant_frame.to_numpy()
+    def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray:
+        return self._compliant_frame.__array__(dtype, copy=copy)
 
     def __repr__(self) -> str:  # pragma: no cover
         header = " Narwhals DataFrame                            "
@@ -343,6 +343,38 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return self._compliant_frame.to_pandas()
 
+    def write_csv(self, file: str | Path | BytesIO | None = None) -> Any:
+        r"""
+        Write dataframe to parquet file.
+
+        Examples:
+            Construct pandas and Polars DataFrames:
+
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> df = {"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0], "ham": ["a", "b", "c"]}
+            >>> df_pd = pd.DataFrame(df)
+            >>> df_pl = pl.DataFrame(df)
+
+            We define a library agnostic function:
+
+            >>> def func(df):
+            ...     df = nw.from_native(df)
+            ...     return df.write_csv()
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)  # doctest: +SKIP
+            'foo,bar,ham\n1,6.0,a\n2,7.0,b\n3,8.0,c\n'
+            >>> func(df_pl)  # doctest: +SKIP
+            'foo,bar,ham\n1,6.0,a\n2,7.0,b\n3,8.0,c\n'
+
+            If we had passed a file name to `write_csv`, it would have been
+            written to that file.
+        """
+        return self._compliant_frame.write_csv(file)
+
     def write_parquet(self, file: str | Path | BytesIO) -> Any:
         """
         Write dataframe to parquet file.


Test patch:
diff --git a/tests/expr_and_series/over_test.py b/tests/expr_and_series/over_test.py
index d9f6d607..fb01a3cf 100644
--- a/tests/expr_and_series/over_test.py
+++ b/tests/expr_and_series/over_test.py
@@ -12,10 +12,7 @@ data = {
 }
 
 
-def test_over_single(request: Any, constructor: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_over_single(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     result = df.with_columns(c_max=nw.col("c").max().over("a"))
     expected = {
@@ -27,10 +24,7 @@ def test_over_single(request: Any, constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_over_multiple(request: Any, constructor: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_over_multiple(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     result = df.with_columns(c_min=nw.col("c").min().over("a", "b"))
     expected = {
@@ -43,7 +37,7 @@ def test_over_multiple(request: Any, constructor: Any) -> None:
 
 
 def test_over_invalid(request: Any, constructor: Any) -> None:
-    if "polars" in str(constructor) or "dask" in str(constructor):
+    if "polars" in str(constructor):
         request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(constructor(data))
diff --git a/tests/frame/array_dunder_test.py b/tests/frame/array_dunder_test.py
new file mode 100644
index 00000000..8f9cbd16
--- /dev/null
+++ b/tests/frame/array_dunder_test.py
@@ -0,0 +1,61 @@
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+import pytest
+
+import narwhals.stable.v1 as nw
+from narwhals.utils import parse_version
+from tests.utils import compare_dicts
+
+
+def test_array_dunder(request: Any, constructor_eager: Any) -> None:
+    if "pyarrow_table" in str(constructor_eager) and parse_version(
+        pa.__version__
+    ) < parse_version("16.0.0"):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)
+    result = df.__array__()
+    np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype="int64"))
+
+
+def test_array_dunder_with_dtype(request: Any, constructor_eager: Any) -> None:
+    if "pyarrow_table" in str(constructor_eager) and parse_version(
+        pa.__version__
+    ) < parse_version("16.0.0"):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)
+    result = df.__array__(object)
+    np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype=object))
+
+
+def test_array_dunder_with_copy(request: Any, constructor_eager: Any) -> None:
+    if "pyarrow_table" in str(constructor_eager) and parse_version(pa.__version__) < (
+        16,
+        0,
+        0,
+    ):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor_eager) and parse_version(pl.__version__) < (
+        0,
+        20,
+        28,
+    ):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)
+    result = df.__array__(copy=True)
+    np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype="int64"))
+    if "pandas_constructor" in str(constructor_eager) and parse_version(
+        pd.__version__
+    ) < (3,):
+        # If it's pandas, we know that `copy=False` definitely took effect.
+        # So, let's check it!
+        result = df.__array__(copy=False)
+        np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype="int64"))
+        result[0, 0] = 999
+        compare_dicts(df, {"a": [999, 2, 3]})
diff --git a/tests/frame/to_numpy_test.py b/tests/frame/to_numpy_test.py
index 4310f676..d573f432 100644
--- a/tests/frame/to_numpy_test.py
+++ b/tests/frame/to_numpy_test.py
@@ -7,7 +7,7 @@ import numpy as np
 import narwhals.stable.v1 as nw
 
 
-def test_convert_numpy(constructor_eager: Any) -> None:
+def test_to_numpy(constructor_eager: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.1, 8, 9]}
     df_raw = constructor_eager(data)
     result = nw.from_native(df_raw, eager_only=True).to_numpy()
@@ -15,7 +15,3 @@ def test_convert_numpy(constructor_eager: Any) -> None:
     expected = np.array([[1, 3, 2], [4, 4, 6], [7.1, 8, 9]]).T
     np.testing.assert_array_equal(result, expected)
     assert result.dtype == "float64"
-
-    result = nw.from_native(df_raw, eager_only=True).__array__()
-    np.testing.assert_array_equal(result, expected)
-    assert result.dtype == "float64"
diff --git a/tests/frame/write_csv_test.py b/tests/frame/write_csv_test.py
new file mode 100644
index 00000000..ed930360
--- /dev/null
+++ b/tests/frame/write_csv_test.py
@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Any
+
+import narwhals.stable.v1 as nw
+from tests.utils import is_windows
+
+if TYPE_CHECKING:
+    import pytest
+
+
+def test_write_csv(constructor_eager: Any, tmpdir: pytest.TempdirFactory) -> None:
+    data = {"a": [1, 2, 3]}
+    path = tmpdir / "foo.csv"  # type: ignore[operator]
+    result = nw.from_native(constructor_eager(data), eager_only=True).write_csv(str(path))
+    assert path.exists()
+    assert result is None
+    result = nw.from_native(constructor_eager(data), eager_only=True).write_csv()
+    if is_windows():  # pragma: no cover
+        result = result.replace("\r\n", "\n")
+    if "pyarrow_table" in str(constructor_eager):
+        assert result == '"a"\n1\n2\n3\n'
+    else:
+        assert result == "a\n1\n2\n3\n"
diff --git a/tests/series_only/array_dunder_test.py b/tests/series_only/array_dunder_test.py
index 199059ae..0449199e 100644
--- a/tests/series_only/array_dunder_test.py
+++ b/tests/series_only/array_dunder_test.py
@@ -1,11 +1,13 @@
 from typing import Any
 
 import numpy as np
+import pandas as pd
 import pyarrow as pa
 import pytest
 
 import narwhals.stable.v1 as nw
 from narwhals.utils import parse_version
+from tests.utils import compare_dicts
 
 
 def test_array_dunder(request: Any, constructor_eager: Any) -> None:
@@ -14,6 +16,37 @@ def test_array_dunder(request: Any, constructor_eager: Any) -> None:
     ) < parse_version("16.0.0"):  # pragma: no cover
         request.applymarker(pytest.mark.xfail)
 
+    s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"]
+    result = s.__array__()
+    np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype="int64"))
+
+
+def test_array_dunder_with_dtype(request: Any, constructor_eager: Any) -> None:
+    if "pyarrow_table" in str(constructor_eager) and parse_version(
+        pa.__version__
+    ) < parse_version("16.0.0"):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
+
     s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"]
     result = s.__array__(object)
     np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype=object))
+
+
+def test_array_dunder_with_copy(request: Any, constructor_eager: Any) -> None:
+    if "pyarrow_table" in str(constructor_eager) and parse_version(
+        pa.__version__
+    ) < parse_version("16.0.0"):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
+
+    s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"]
+    result = s.__array__(copy=True)
+    np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype="int64"))
+    if "pandas_constructor" in str(constructor_eager) and parse_version(
+        pd.__version__
+    ) < (3,):
+        # If it's pandas, we know that `copy=False` definitely took effect.
+        # So, let's check it!
+        result = s.__array__(copy=False)
+        np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype="int64"))
+        result[0] = 999
+        compare_dicts({"a": s}, {"a": [999, 2, 3]})