Repository: narwhals-dev/narwhals
Pull request #826 (2024-08-20 15:43:14)

feat: support dtype and copy in DataFrame.__array__
<!--
# 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
- [ ] ✨ 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 865d1709..e0d3d51a 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]
 
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 4faf61a0..23633564 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:
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..aae29107 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                            "


Test patch:
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/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]})