Repository: narwhals-dev/narwhals
Pull request #488 (2024-07-11 11:32:42)

feat: Arrow `head`, `tail`, `null_count`
## 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.

Some more low hanging methods for pyarrow


Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 525402b8..fa03b625 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -239,6 +239,31 @@ class ArrowDataFrame:
         row_indices = pa.array(range(df.num_rows))
         return self._from_native_dataframe(df.append_column(name, row_indices))
 
+    def null_count(self) -> Self:
+        pa = get_pyarrow()
+        df = self._native_dataframe
+        names_and_values = zip(df.column_names, df.columns)
+
+        return self._from_native_dataframe(
+            pa.table({name: [col.null_count] for name, col in names_and_values})
+        )
+
+    def head(self, n: int) -> Self:
+        df = self._native_dataframe
+        if n >= 0:
+            return self._from_native_dataframe(df.slice(0, n))
+        else:
+            num_rows = df.num_rows
+            return self._from_native_dataframe(df.slice(0, max(0, num_rows + n)))
+
+    def tail(self, n: int) -> Self:
+        df = self._native_dataframe
+        if n >= 0:
+            num_rows = df.num_rows
+            return self._from_native_dataframe(df.slice(max(0, num_rows - n)))
+        else:
+            return self._from_native_dataframe(df.slice(abs(n)))
+
     def lazy(self) -> Self:
         return self
 
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 88bd32c3..cd0332ef 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -122,6 +122,15 @@ class ArrowExpr:
             backend_version=self._backend_version,
         )
 
+    def null_count(self) -> Self:
+        return reuse_series_implementation(self, "null_count", returns_scalar=True)
+
+    def head(self, n: int) -> Self:
+        return reuse_series_implementation(self, "head", n)
+
+    def tail(self, n: int) -> Self:
+        return reuse_series_implementation(self, "tail", n)
+
     @property
     def dt(self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index 36af9f1f..14f73381 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -73,7 +73,7 @@ class ArrowNamespace:
     def _create_series_from_scalar(self, value: Any, series: ArrowSeries) -> ArrowSeries:
         from narwhals._arrow.series import ArrowSeries
 
-        if self._backend_version < (13,):  # pragma: no cover
+        if self._backend_version < (13,) and hasattr(value, "as_py"):  # pragma: no cover
             value = value.as_py()
         return ArrowSeries._from_iterable(
             [value],
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 286bb37d..1fffd249 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -144,6 +144,25 @@ class ArrowSeries:
         dtype = reverse_translate_dtype(dtype)
         return self._from_native_series(pc.cast(ser, dtype))
 
+    def null_count(self: Self) -> int:
+        return self._native_series.null_count  # type: ignore[no-any-return]
+
+    def head(self, n: int) -> Self:
+        ser = self._native_series
+        if n >= 0:
+            return self._from_native_series(ser.slice(0, n))
+        else:
+            num_rows = len(ser)
+            return self._from_native_series(ser.slice(0, max(0, num_rows + n)))
+
+    def tail(self, n: int) -> Self:
+        ser = self._native_series
+        if n >= 0:
+            num_rows = len(ser)
+            return self._from_native_series(ser.slice(max(0, num_rows - n)))
+        else:
+            return self._from_native_series(ser.slice(abs(n)))
+
     @property
     def shape(self) -> tuple[int]:
         return (len(self._native_series),)
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index abed808c..f1d0b184 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -14,24 +14,20 @@ from narwhals._arrow.dataframe import ArrowDataFrame
 MISSING = [
     "DataFrame.filter",
     "DataFrame.group_by",
-    "DataFrame.head",
     "DataFrame.is_duplicated",
     "DataFrame.is_empty",
     "DataFrame.is_unique",
     "DataFrame.item",
     "DataFrame.iter_rows",
     "DataFrame.join",
-    "DataFrame.null_count",
     "DataFrame.pipe",
     "DataFrame.rename",
-    "DataFrame.tail",
     "DataFrame.unique",
     "DataFrame.write_parquet",
     "Series.drop_nulls",
     "Series.fill_null",
     "Series.filter",
     "Series.from_iterable",
-    "Series.head",
     "Series.is_between",
     "Series.is_duplicated",
     "Series.is_first_distinct",
@@ -45,14 +41,12 @@ MISSING = [
     "Series.max",
     "Series.min",
     "Series.n_unique",
-    "Series.null_count",
     "Series.quantile",
     "Series.round",
     "Series.sample",
     "Series.shift",
     "Series.sort",
     "Series.sum",
-    "Series.tail",
     "Series.to_frame",
     "Series.to_pandas",
     "Series.unique",


Test patch:
diff --git a/tests/expr/null_count_test.py b/tests/expr/null_count_test.py
index 5461db86..e5a02e26 100644
--- a/tests/expr/null_count_test.py
+++ b/tests/expr/null_count_test.py
@@ -9,8 +9,8 @@ data = {
 }
 
 
-def test_null_count(constructor: Any) -> None:
-    df = nw.from_native(constructor(data), eager_only=True)
+def test_null_count(constructor_with_pyarrow: Any) -> None:
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
     result = df.select(nw.all().null_count())
     expected = {
         "a": [2],
diff --git a/tests/frame/head_test.py b/tests/frame/head_test.py
index ef61334b..9b87e579 100644
--- a/tests/frame/head_test.py
+++ b/tests/frame/head_test.py
@@ -6,16 +6,21 @@ import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
 
-def test_head(constructor_with_lazy: Any) -> None:
+def test_head(constructor_with_pyarrow: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
-    df_raw = constructor_with_lazy(data)
-    df = nw.from_native(df_raw).lazy()
-    result = nw.to_native(df.head(2))
     expected = {"a": [1, 3], "b": [4, 4], "z": [7.0, 8.0]}
+
+    df_raw = constructor_with_pyarrow(data)
+    df = nw.from_native(df_raw).lazy()
+
+    result = df.head(2)
     compare_dicts(result, expected)
-    result = nw.to_native(df.collect().head(2))
-    expected = {"a": [1, 3], "b": [4, 4], "z": [7.0, 8.0]}
+
+    result = df.collect().head(2)  # type: ignore[assignment]
     compare_dicts(result, expected)
-    result = nw.to_native(df.collect().select(nw.col("a").head(2)))
-    expected = {"a": [1, 3]}
+
+    result = df.collect().head(-1)  # type: ignore[assignment]
     compare_dicts(result, expected)
+
+    result = df.collect().select(nw.col("a").head(2))  # type: ignore[assignment]
+    compare_dicts(result, {"a": expected["a"]})
diff --git a/tests/frame/null_count_test.py b/tests/frame/null_count_test.py
index 4e12dc61..5fc66daf 100644
--- a/tests/frame/null_count_test.py
+++ b/tests/frame/null_count_test.py
@@ -6,10 +6,10 @@ import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
 
-def test_null_count(constructor: Any) -> None:
+def test_null_count(constructor_with_pyarrow: Any) -> None:
     data = {"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]}
-    df_raw = constructor(data)
+    df_raw = constructor_with_pyarrow(data)
     df = nw.from_native(df_raw, eager_only=True)
-    result = nw.to_native(df.null_count())
+    result = df.null_count()
     expected = {"a": [1], "b": [0], "z": [1]}
     compare_dicts(result, expected)
diff --git a/tests/frame/tail_test.py b/tests/frame/tail_test.py
index 768b1969..28f9edf3 100644
--- a/tests/frame/tail_test.py
+++ b/tests/frame/tail_test.py
@@ -6,16 +6,21 @@ import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
 
-def test_tail(constructor_with_lazy: Any) -> None:
+def test_tail(constructor_with_pyarrow: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
-    df_raw = constructor_with_lazy(data)
-    df = nw.from_native(df_raw).lazy()
-    result = nw.to_native(df.tail(2))
     expected = {"a": [3, 2], "b": [4, 6], "z": [8.0, 9]}
+
+    df_raw = constructor_with_pyarrow(data)
+    df = nw.from_native(df_raw).lazy()
+
+    result = df.tail(2)
     compare_dicts(result, expected)
-    result = nw.to_native(df.collect().tail(2))
-    expected = {"a": [3, 2], "b": [4, 6], "z": [8.0, 9]}
+
+    result = df.collect().tail(2)  # type: ignore[assignment]
     compare_dicts(result, expected)
-    result = nw.to_native(df.collect().select(nw.col("a").tail(2)))
-    expected = {"a": [3, 2]}
+
+    result = df.collect().tail(-1)  # type: ignore[assignment]
     compare_dicts(result, expected)
+
+    result = df.collect().select(nw.col("a").tail(2))  # type: ignore[assignment]
+    compare_dicts(result, {"a": expected["a"]})
diff --git a/tests/series/head_test.py b/tests/series/head_test.py
new file mode 100644
index 00000000..3418bd92
--- /dev/null
+++ b/tests/series/head_test.py
@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+import narwhals as nw
+
+
+@pytest.mark.parametrize("n", [2, -1])
+def test_head(constructor_series_with_pyarrow: Any, n: int) -> None:
+    s = nw.from_native(constructor_series_with_pyarrow([1, 2, 3]), series_only=True)
+
+    assert s.head(n).to_list() == [1, 2]
diff --git a/tests/series/null_count_test.py b/tests/series/null_count_test.py
new file mode 100644
index 00000000..b2bf84bd
--- /dev/null
+++ b/tests/series/null_count_test.py
@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Any
+
+import narwhals as nw
+
+
+def test_null_count(constructor_series_with_pyarrow: Any) -> None:
+    data = [1, 2, None]
+    series = nw.from_native(constructor_series_with_pyarrow(data), series_only=True)
+    result = series.null_count()
+    assert result == 1
diff --git a/tests/series/tail_test.py b/tests/series/tail_test.py
new file mode 100644
index 00000000..21759321
--- /dev/null
+++ b/tests/series/tail_test.py
@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+import narwhals as nw
+
+
+@pytest.mark.parametrize("n", [2, -1])
+def test_tail(constructor_series_with_pyarrow: Any, n: int) -> None:
+    s = nw.from_native(constructor_series_with_pyarrow([1, 2, 3]), series_only=True)
+
+    assert s.tail(n).to_list() == [2, 3]
diff --git a/tests/series/test_common.py b/tests/series/test_common.py
index c69f852e..2a573958 100644
--- a/tests/series/test_common.py
+++ b/tests/series/test_common.py
@@ -172,13 +172,6 @@ def test_is_unique(df_raw: Any) -> None:
     assert (result.to_numpy() == expected).all()
 
 
-@pytest.mark.parametrize("s_raw", [pd.Series([1, 2, None]), pl.Series([1, 2, None])])
-def test_null_count(s_raw: Any) -> None:
-    series = nw.from_native(s_raw, series_only=True)
-    result = series.null_count()
-    assert result == 1
-
-
 @pytest.mark.parametrize("df_raw", [df_pandas, df_polars])
 def test_is_first_distinct(df_raw: Any) -> None:
     series = nw.from_native(df_raw["b"], series_only=True)
@@ -299,21 +292,3 @@ def test_item(df_raw: Any, index: int, expected: int) -> None:
         match=re.escape("can only call '.item()' if the Series is of length 1,"),
     ):
         s.item(None)
-
-
-@pytest.mark.parametrize("df_raw", [df_pandas, df_polars])
-@pytest.mark.parametrize("n", [1, 2, 3, 10])
-def test_head(df_raw: Any, n: int) -> None:
-    s_raw = df_raw["z"]
-    s = nw.from_native(s_raw, allow_series=True)
-
-    assert s.head(n) == nw.from_native(s_raw.head(n), series_only=True)
-
-
-@pytest.mark.parametrize("df_raw", [df_pandas, df_polars])
-@pytest.mark.parametrize("n", [1, 2, 3, 10])
-def test_tail(df_raw: Any, n: int) -> None:
-    s_raw = df_raw["z"]
-    s = nw.from_native(s_raw, allow_series=True)
-
-    assert s.tail(n) == nw.from_native(s_raw.tail(n), series_only=True)