Repository: narwhals-dev/narwhals
Pull request #359 (2024-06-30 11:56:16)

feat: add more methods to the pyarrow backend
<!--
# 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 b11d93f8..9cf65e72 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -41,6 +41,17 @@ class ArrowDataFrame:
     def shape(self) -> tuple[int, int]:
         return self._dataframe.shape  # type: ignore[no-any-return]
 
+    def __len__(self) -> int:
+        return len(self._dataframe)
+
+    def rows(
+        self, *, named: bool = False
+    ) -> list[tuple[Any, ...]] | list[dict[str, Any]]:
+        if not named:
+            msg = "Unnamed rows are not yet supported on PyArrow tables"
+            raise NotImplementedError(msg)
+        return self._dataframe.to_pylist()  # type: ignore[no-any-return]
+
     @overload
     def __getitem__(self, item: str) -> ArrowSeries: ...
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index be97fa80..0044e3da 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -3,11 +3,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 from typing import Any
 
+from narwhals._arrow.utils import translate_dtype
 from narwhals.dependencies import get_pyarrow_compute
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from narwhals.dtypes import DType
+
 
 class ArrowSeries:
     def __init__(
@@ -51,6 +54,10 @@ class ArrowSeries:
             name=name,
         )
 
+    @property
+    def dtype(self) -> DType:
+        return translate_dtype(self._series.type)
+
     def cum_sum(self) -> Self:
         pc = get_pyarrow_compute()
         return self._from_series(pc.cumulative_sum(self._series))
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index dc0588ba..e8740c0b 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -122,7 +122,7 @@ class PandasDataFrame:
         if not named:
             return list(self._dataframe.itertuples(index=False, name=None))
 
-        return self._dataframe.to_dict("records")  # type: ignore[no-any-return]
+        return self._dataframe.to_dict(orient="records")  # type: ignore[no-any-return]
 
     def iter_rows(
         self,
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 48315f5c..554f96a3 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -640,27 +640,27 @@ class DataFrame(BaseFrame):
         self,
         *,
         named: Literal[False],
-    ) -> tuple[Any, ...]: ...
+    ) -> list[tuple[Any, ...]]: ...
 
     @overload
     def rows(
         self,
         *,
         named: Literal[True],
-    ) -> dict[str, Any]: ...
+    ) -> list[dict[str, Any]]: ...
 
     @overload
     def rows(
         self,
         *,
         named: bool,
-    ) -> tuple[Any, ...] | dict[str, Any]: ...
+    ) -> list[tuple[Any, ...]] | list[dict[str, Any]]: ...
 
     def rows(
         self,
         *,
         named: bool = False,
-    ) -> tuple[Any, ...] | dict[str, Any]:
+    ) -> list[tuple[Any, ...]] | list[dict[str, Any]]:
         """
         Returns all data in the DataFrame as a list of rows of python-native values.
 
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 6ea1ccea..37fd9f2b 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -19,9 +19,7 @@ if TYPE_CHECKING:
     from narwhals.series import Series
 
 
-def to_native(
-    narwhals_object: LazyFrame | DataFrame | Series, *, strict: bool = True
-) -> Any:
+def to_native(narwhals_object: Any, *, strict: bool = True) -> Any:
     """
     Convert Narwhals object to native one.
 


Test patch:
diff --git a/tests/frame/len_test.py b/tests/frame/len_test.py
index 45921ae5..2cd85269 100644
--- a/tests/frame/len_test.py
+++ b/tests/frame/len_test.py
@@ -2,6 +2,7 @@ from typing import Any
 
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 
 import narwhals as nw
@@ -12,7 +13,7 @@ data = {
 }
 
 
-@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_drop_nulls(constructor: Any) -> None:
     result = len(nw.from_native(constructor(data)))
     assert result == 4
diff --git a/tests/frame/rows_test.py b/tests/frame/rows_test.py
index 0a0089fa..0e6805da 100644
--- a/tests/frame/rows_test.py
+++ b/tests/frame/rows_test.py
@@ -4,13 +4,14 @@ from typing import Any
 
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 
 import narwhals as nw
 from narwhals.utils import parse_version
-from tests.utils import maybe_get_modin_df
 
 df_pandas = pd.DataFrame({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
+df_pa = pa.table({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
 if parse_version(pd.__version__) >= parse_version("1.5.0"):
     df_pandas_pyarrow = pd.DataFrame(
         {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
@@ -34,15 +35,13 @@ else:  # pragma: no cover
     df_pandas_pyarrow = df_pandas
     df_pandas_nullable = df_pandas
 df_polars = pl.DataFrame({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
-df_mpd = maybe_get_modin_df(df_pandas)
 
 df_pandas_na = pd.DataFrame({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]})
 df_polars_na = pl.DataFrame({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]})
 
 
-@pytest.mark.parametrize("method_name", ["iter_rows", "rows"])
 @pytest.mark.parametrize(
-    "df_raw", [df_pandas, df_pandas_nullable, df_pandas_pyarrow, df_mpd, df_polars]
+    "df_raw", [df_pandas, df_pandas_nullable, df_pandas_pyarrow, df_polars]
 )
 @pytest.mark.parametrize(
     ("named", "expected"),
@@ -58,20 +57,47 @@ df_polars_na = pl.DataFrame({"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None,
         ),
     ],
 )
-@pytest.mark.filterwarnings("ignore::FutureWarning")
-def test_rows(
-    method_name: str,
+def test_iter_rows(
     df_raw: Any,
     named: bool,  # noqa: FBT001
     expected: list[tuple[Any, ...]] | list[dict[str, Any]],
 ) -> None:
-    # GIVEN
     df = nw.DataFrame(df_raw)
+    result = list(df.iter_rows(named=named))
+    assert result == expected
 
-    # WHEN
-    result = list(getattr(df, method_name)(named=named))
 
-    # THEN
+@pytest.mark.parametrize(
+    "df_raw", [df_pandas, df_pandas_nullable, df_pandas_pyarrow, df_polars, df_pa]
+)
+@pytest.mark.parametrize(
+    ("named", "expected"),
+    [
+        (False, [(1, 4, 7.0), (3, 4, 8.0), (2, 6, 9.0)]),
+        (
+            True,
+            [
+                {"a": 1, "b": 4, "z": 7.0},
+                {"a": 3, "b": 4, "z": 8.0},
+                {"a": 2, "b": 6, "z": 9.0},
+            ],
+        ),
+    ],
+)
+def test_rows(
+    df_raw: Any,
+    named: bool,  # noqa: FBT001
+    expected: list[tuple[Any, ...]] | list[dict[str, Any]],
+) -> None:
+    df = nw.DataFrame(df_raw)
+    if isinstance(df_raw, pa.Table) and not named:
+        with pytest.raises(
+            NotImplementedError,
+            match="Unnamed rows are not yet supported on PyArrow tables",
+        ):
+            df.rows(named=named)
+        return
+    result = df.rows(named=named)
     assert result == expected
 
 
diff --git a/tests/frame/schema_test.py b/tests/frame/schema_test.py
index af2ee233..bf77b7e6 100644
--- a/tests/frame/schema_test.py
+++ b/tests/frame/schema_test.py
@@ -47,7 +47,7 @@ def test_actual_object(constructor: Any) -> None:
     parse_version(pd.__version__) < parse_version("2.0.0"), reason="too old"
 )
 def test_dtypes() -> None:
-    df = pl.DataFrame(
+    df_pl = pl.DataFrame(
         {
             "a": [1],
             "b": [1],
@@ -85,7 +85,8 @@ def test_dtypes() -> None:
             "p": pl.Categorical,
         },
     )
-    result = nw.DataFrame(df).schema
+    df = nw.DataFrame(df_pl)
+    result = df.schema
     expected = {
         "a": nw.Int64,
         "b": nw.Int32,
@@ -105,7 +106,14 @@ def test_dtypes() -> None:
         "p": nw.Categorical,
     }
     assert result == expected
-    result_pd = nw.DataFrame(df.to_pandas(use_pyarrow_extension_array=True)).schema
+    assert {name: df[name].dtype for name in df.columns} == expected
+    df_pd = df_pl.to_pandas(use_pyarrow_extension_array=True)
+    df = nw.DataFrame(df_pd)
+    result_pd = df.schema
     assert result_pd == expected
-    result_pa = nw.DataFrame(df.to_arrow()).schema
+    assert {name: df[name].dtype for name in df.columns} == expected
+    df_pa = df_pl.to_arrow()
+    df = nw.DataFrame(df_pa)
+    result_pa = df.schema
     assert result_pa == expected
+    assert {name: df[name].dtype for name in df.columns} == expected