Repository: narwhals-dev/narwhals
Pull request #404 (2024-07-04 14:52:20)

feat: some pyarrow dataframe methods
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## 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.

Introduces the following methods for pyarrow dataframe:

- collect
- drop
- drop_nulls
- lazy
- sort
- to_pandas


Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index c3f2c0fb..6c9c390d 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Iterable
 from typing import Sequence
 from typing import overload
 
@@ -9,6 +10,7 @@ from narwhals._arrow.utils import translate_dtype
 from narwhals._pandas_like.utils import evaluate_into_exprs
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pyarrow
+from narwhals.utils import flatten
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -36,6 +38,9 @@ class ArrowDataFrame:
     def __narwhals_dataframe__(self) -> Self:
         return self
 
+    def __narwhals_lazyframe__(self) -> Self:
+        return self
+
     def _from_dataframe(self, df: Any) -> Self:
         return self.__class__(df)
 
@@ -112,3 +117,37 @@ class ArrowDataFrame:
         pa = get_pyarrow()
         df = pa.Table.from_arrays([s._series for s in new_series], names=names)
         return self._from_dataframe(df)
+
+    def drop(self, *columns: str | Iterable[str]) -> Self:
+        return self._from_dataframe(self._dataframe.drop(list(flatten(columns))))
+
+    def drop_nulls(self) -> Self:
+        return self._from_dataframe(self._dataframe.drop_null())
+
+    def sort(
+        self,
+        by: str | Iterable[str],
+        *more_by: str,
+        descending: bool | Sequence[bool] = False,
+    ) -> Self:
+        flat_keys = flatten([*flatten([by]), *more_by])
+        df = self._dataframe
+
+        if isinstance(descending, bool):
+            order = "descending" if descending else "ascending"
+            sorting = [(key, order) for key in flat_keys]
+        else:
+            sorting = [
+                (key, "descending" if is_descending else "ascending")
+                for key, is_descending in zip(flat_keys, descending)
+            ]
+        return self._from_dataframe(df.sort_by(sorting=sorting))
+
+    def to_pandas(self) -> Any:
+        return self._dataframe.to_pandas()
+
+    def lazy(self) -> Self:
+        return self
+
+    def collect(self) -> ArrowDataFrame:
+        return ArrowDataFrame(self._dataframe)
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 66a3aa1b..35641247 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -62,6 +62,9 @@ class ArrowSeries:
     def to_list(self) -> Any:
         return self._series.to_pylist()
 
+    def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any:
+        return self._series.__array__(dtype=dtype, copy=copy)
+
     def to_numpy(self) -> Any:
         return self._series.to_numpy()
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 001c899a..abfd49fb 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -431,7 +431,17 @@ class PandasSeries:
         ser = self._series
         return self._from_series(self._rename(ser, name))
 
-    def to_numpy(self) -> Any:
+    def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any:
+        # pandas used to always return object dtype for nullable dtypes.
+        # So, we intercept __array__ and pass to `to_numpy` ourselves to make
+        # sure an appropriate numpy dtype is returned.
+        return self.to_numpy(dtype=dtype, copy=copy)
+
+    def to_numpy(self, dtype: Any = None, copy: bool | None = None) -> Any:
+        # the default is meant to be None, but pandas doesn't allow it?
+        # https://numpy.org/doc/stable/reference/generated/numpy.ndarray.__array__.html
+        copy = copy or False
+
         has_missing = self._series.isna().any()
         if has_missing and str(self._series.dtype) in PANDAS_TO_NUMPY_DTYPE_MISSING:
             if self._implementation == "pandas" and parse_version(
@@ -441,16 +451,19 @@ class PandasSeries:
             else:
                 kwargs = {"na_value": float("nan")}
             return self._series.to_numpy(
-                dtype=PANDAS_TO_NUMPY_DTYPE_MISSING[str(self._series.dtype)], **kwargs
+                dtype=dtype or PANDAS_TO_NUMPY_DTYPE_MISSING[str(self._series.dtype)],
+                copy=copy,
+                **kwargs,
             )
         if (
             not has_missing
             and str(self._series.dtype) in PANDAS_TO_NUMPY_DTYPE_NO_MISSING
         ):
             return self._series.to_numpy(
-                dtype=PANDAS_TO_NUMPY_DTYPE_NO_MISSING[str(self._series.dtype)]
+                dtype=dtype or PANDAS_TO_NUMPY_DTYPE_NO_MISSING[str(self._series.dtype)],
+                copy=copy,
             )
-        return self._series.to_numpy()
+        return self._series.to_numpy(dtype=dtype, copy=copy)
 
     def to_pandas(self) -> Any:
         if self._implementation == "pandas":
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 6ba38b4f..239fcd2c 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -245,8 +245,8 @@ class DataFrame(BaseFrame):
             msg = f"Expected pandas-like dataframe, Polars dataframe, or Polars lazyframe, got: {type(df)}"
             raise TypeError(msg)
 
-    def __array__(self, *args: Any, **kwargs: Any) -> np.ndarray:
-        return self._dataframe.to_numpy(*args, **kwargs)
+    def __array__(self) -> np.ndarray:
+        return self._dataframe.to_numpy()
 
     def __repr__(self) -> str:  # pragma: no cover
         header = " Narwhals DataFrame                            "
diff --git a/narwhals/series.py b/narwhals/series.py
index cc0941b3..a5411bf7 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -75,8 +75,8 @@ class Series:
         )
         raise TypeError(msg)  # pragma: no cover
 
-    def __array__(self, *args: Any, **kwargs: Any) -> np.ndarray:
-        return self._series.to_numpy(*args, **kwargs)
+    def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray:
+        return self._series.__array__(dtype=dtype, copy=copy)
 
     def __getitem__(self, idx: int | slice) -> Any:
         if isinstance(idx, int):
diff --git a/noxfile.py b/noxfile.py
index 7d958fd0..9478f0a9 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -32,7 +32,15 @@ def pytest_coverage(session: Session) -> None:
 
 @nox.session(python=PYTHON_VERSIONS[0])  # type: ignore[misc]
 def minimum_versions(session: Session) -> None:
-    session.install("pandas==1.1.5", "polars==0.20.3", "numpy<=1.21")
+    session.install(
+        "pandas==0.25.3",
+        "polars==0.20.3",
+        "numpy==1.17.5",
+        "pyarrow==11.0.0",
+        "scipy==1.5.0",
+        "scikit-learn==1.1.0",
+        "tzdata",
+    )
     run_common(session, coverage_threshold=50)
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 90c9a24e..11139298 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,6 +17,14 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 
+[tool.hatch.build]
+exclude = [
+  "/.*",
+  "/docs",
+  "/tests",
+  "/tpch",
+]
+
 [project.optional-dependencies]
 pandas = ["pandas>=0.25.3"]
 polars = ["polars>=0.20.3"]
@@ -92,6 +100,7 @@ exclude_also = [
   "if sys.version_info() <",
   "if implementation == \"modin\"",
   "if implementation == \"cudf\"",
+  'request.applymarker\(pytest.mark.xfail\)'
 ]
 
 [tool.mypy]
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index dfb8853c..d387682a 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -12,9 +12,6 @@ import narwhals as nw
 from narwhals._arrow.dataframe import ArrowDataFrame
 
 MISSING = [
-    "DataFrame.collect",
-    "DataFrame.drop",
-    "DataFrame.drop_nulls",
     "DataFrame.filter",
     "DataFrame.group_by",
     "DataFrame.head",
@@ -24,15 +21,12 @@ MISSING = [
     "DataFrame.item",
     "DataFrame.iter_rows",
     "DataFrame.join",
-    "DataFrame.lazy",
     "DataFrame.null_count",
     "DataFrame.pipe",
     "DataFrame.rename",
-    "DataFrame.sort",
     "DataFrame.tail",
     "DataFrame.to_dict",
     "DataFrame.to_numpy",
-    "DataFrame.to_pandas",
     "DataFrame.unique",
     "DataFrame.with_columns",
     "DataFrame.with_row_index",


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 00b9274c..9b7fc942 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -112,8 +112,12 @@ def constructor_series(request: Any) -> Callable[[Any], Any]:
     return request.param  # type: ignore[no-any-return]
 
 
+def pyarrow_chunked_array_constructor(obj: Any) -> Any:
+    return pa.chunked_array([obj])
+
+
 # TODO: once pyarrow has complete coverage, we can remove this one,
 # and just put `pa.table` into `constructor`
-@pytest.fixture(params=[*params_series, lambda x: pa.chunked_array([x])])
+@pytest.fixture(params=[*params_series, pyarrow_chunked_array_constructor])
 def constructor_series_with_pyarrow(request: Any) -> Callable[[Any], Any]:
     return request.param  # type: ignore[no-any-return]
diff --git a/tests/frame/drop_nulls_test.py b/tests/frame/drop_nulls_test.py
index c63e0029..0646c803 100644
--- a/tests/frame/drop_nulls_test.py
+++ b/tests/frame/drop_nulls_test.py
@@ -9,12 +9,12 @@ data = {
 }
 
 
-def test_drop_nulls(constructor: Any) -> None:
-    result = nw.from_native(constructor(data)).drop_nulls()
+def test_drop_nulls(constructor_with_pyarrow: Any) -> None:
+    result = nw.from_native(constructor_with_pyarrow(data)).drop_nulls()
     expected = {
         "a": [2.0, 4.0],
         "b": [3.0, 5.0],
     }
     compare_dicts(result, expected)
-    result = nw.from_native(constructor(data)).lazy().drop_nulls()
+    result = nw.from_native(constructor_with_pyarrow(data)).lazy().drop_nulls()
     compare_dicts(result, expected)
diff --git a/tests/frame/sort_test.py b/tests/frame/sort_test.py
index 21cb508e..a5dd8e0b 100644
--- a/tests/frame/sort_test.py
+++ b/tests/frame/sort_test.py
@@ -4,22 +4,20 @@ import narwhals as nw
 from tests.utils import compare_dicts
 
 
-def test_sort(constructor: Any) -> None:
+def test_sort(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)
     result = df.sort("a", "b")
-    result_native = nw.to_native(result)
     expected = {
         "a": [1, 2, 3],
         "b": [4, 6, 4],
         "z": [7.0, 9.0, 8.0],
     }
-    compare_dicts(result_native, expected)
+    compare_dicts(result, expected)
     result = df.sort("a", "b", descending=[True, False])
-    result_native = nw.to_native(result)
     expected = {
         "a": [3, 2, 1],
         "b": [4, 6, 4],
         "z": [8.0, 9.0, 7.0],
     }
-    compare_dicts(result_native, expected)
+    compare_dicts(result, expected)
diff --git a/tests/frame/test_common.py b/tests/frame/test_common.py
index e555d8e4..ec73d7bd 100644
--- a/tests/frame/test_common.py
+++ b/tests/frame/test_common.py
@@ -55,6 +55,8 @@ df_lazy_na = pl.LazyFrame({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9
 df_right_pandas = pd.DataFrame({"c": [6, 12, -1], "d": [0, -4, 2]})
 df_right_lazy = pl.LazyFrame({"c": [6, 12, -1], "d": [0, -4, 2]})
 df_mpd = maybe_get_modin_df(df_pandas)
+df_pa = pa.table({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
+df_pa_na = pa.table({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]})
 
 
 @pytest.mark.parametrize(
@@ -418,7 +420,7 @@ def test_accepted_dataframes() -> None:
         nw.LazyFrame(array)
 
 
-@pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_mpd])
+@pytest.mark.parametrize("df_raw", [df_polars, df_pandas, df_mpd, df_pa])
 @pytest.mark.filterwarnings("ignore:.*Passing a BlockManager.*:DeprecationWarning")
 @pytest.mark.skipif(
     parse_version(pd.__version__) < parse_version("2.0.0"),
@@ -602,6 +604,7 @@ def test_drop_nulls(df_raw: Any) -> None:
         df_pandas,
         df_polars,
         df_mpd,
+        df_pa,
     ],
 )
 @pytest.mark.parametrize(
@@ -654,7 +657,7 @@ def test_concat_vertical(df_raw: Any, df_raw_right: Any) -> None:
         nw.concat([df_left, df_right.rename({"d": "i"})], how="vertical").collect()  # type: ignore[union-attr]
 
 
-@pytest.mark.parametrize("df_raw", [df_pandas, df_polars])
+@pytest.mark.parametrize("df_raw", [df_pandas, df_polars, df_pa])
 def test_lazy(df_raw: Any) -> None:
     df = nw.from_native(df_raw, eager_only=True)
     result = df.lazy()
diff --git a/tests/series/array_dunder_test.py b/tests/series/array_dunder_test.py
new file mode 100644
index 00000000..a9d3546a
--- /dev/null
+++ b/tests/series/array_dunder_test.py
@@ -0,0 +1,23 @@
+from typing import Any
+
+import numpy as np
+import polars as pl
+import pyarrow as pa
+import pytest
+
+import narwhals as nw
+from narwhals.utils import parse_version
+
+
+def test_array_dunder(constructor_series_with_pyarrow: Any, request: Any) -> None:
+    if "polars" in str(constructor_series_with_pyarrow) and parse_version(
+        pl.__version__
+    ) < parse_version("0.20.10"):
+        request.applymarker(pytest.mark.xfail)
+    if "chunked_array" in str(constructor_series_with_pyarrow) and parse_version(
+        pa.__version__
+    ) < parse_version("16.0.0"):
+        request.applymarker(pytest.mark.xfail)
+    s = nw.from_native(constructor_series_with_pyarrow([1, 2, 3]), series_only=True)
+    result = s.__array__(object)
+    np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype=object))