Feat: Series.shift Pyarrow Backend Implementation

narwhals/_arrow/expr.py

@@ -200,6 +200,9 @@ def sum(self) -> Self:
     def drop_nulls(self) -> Self:
         return reuse_series_implementation(self, "drop_nulls")
 
+    def shift(self, n: int) -> Self:
+        return reuse_series_implementation(self, "shift", n)
+
     def alias(self, name: str) -> Self:
         # Define this one manually, so that we can
         # override `output_names` and not increase depth

narwhals/_arrow/series.py

@@ -221,6 +221,30 @@ def drop_nulls(self) -> ArrowSeries:
         pc = get_pyarrow_compute()
         return self._from_native_series(pc.drop_null(self._native_series))
 
+    def shift(self, n: int) -> Self:
+        pa = get_pyarrow()
+        chunks = self._native_series.chunks
+        dtype = self._native_series.type
+
+        # This condition shifts the series by n positions.
+        if n > 0:
+            # If n is positive, prepend nulls and truncate the end
+            shifted_chunks = [pa.array([None] * n, type=dtype)] + [
+                chunk.slice(0, len(chunk) - n) for chunk in chunks
+            ]
+        elif n < 0:
+            # If n is negative, append nulls and truncate the start
+            shifted_chunks = [chunk.slice(-n, len(chunk) + n) for chunk in chunks] + [
+                pa.array([None] * -n, type=dtype)
+            ]
+        else:
+            # If n is 0, return the original series
+            shifted_chunks = chunks
+
+        # Flatten the list of chunks to create a single array
+        shifted_array = pa.concat_arrays(shifted_chunks)
+        return self._from_native_series(shifted_array)
+
     def std(self, ddof: int = 1) -> int:
         pc = get_pyarrow_compute()
         return pc.stddev(self._native_series, ddof=ddof)  # type: ignore[no-any-return]

tests/expr_and_series/series_shift_test.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pyarrow as pa
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+
+def test_shift(constructor_eager: Any) -> None:
+    data = {
+        "A": [1, 2, None, 4],
+        "B": [5, 6, 7, 8],
+        "C": [None, None, None, None],
+        "D": [9, 10, 11, 12],
+    }
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+
+    result_a = df.select(nw.col("A").shift(1))
+    result_b = df.select(nw.col("B").shift(-1))
+    result_c = df.select(nw.col("C").shift(1))
+    result_d = df.select(nw.col("D").shift(2))
+
+    expected_a = {"A": [float("nan"), 1.0, 2.0, float("nan")]}
+    expected_b = {"B": [6.0, 7.0, 8.0, float("nan")]}
+    expected_c = {"C": [float("nan"), float("nan"), float("nan"), float("nan")]}
+    expected_d = {"D": [float("nan"), float("nan"), 9, 10]}
+
+    compare_dicts(result_a, expected_a)
+    compare_dicts(result_b, expected_b)
+    compare_dicts(result_c, expected_c)
+    compare_dicts(result_d, expected_d)
+
+
+def test_shift_multi_chunk_pyarrow() -> None:
+    tbl = pa.table({"a": [1, 2, 3]})
+    tbl = pa.concat_tables([tbl, tbl, tbl])
+    tbl = tbl.combine_chunks()
+    df = nw.from_native(tbl, eager_only=True)
+    result_a = df.select(nw.col("a").shift(1))
+    result_b = df.select(nw.col("a").shift(-1))
+    expected_a = {"a": [None, 1, 2, 3, 1, 2, 3, 1, 2]}
+    expected_b = {"a": [2, 3, 1, 2, 3, 1, 2, 3, None]}
+
+    compare_dicts(result_a, expected_a)
+    compare_dicts(result_b, expected_b)

tests/expr_and_series/shift_test.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-import pytest
-
 import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
@@ -13,11 +11,8 @@
 }
 
 
-def test_shift(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
-    df = nw.from_native(constructor(data))
+def test_shift(constructor_eager: Any) -> None:
+    df = nw.from_native(constructor_eager(data), eager_only=True)
     result = df.with_columns(nw.col("a", "b", "c").shift(2)).filter(nw.col("i") > 1)
     expected = {
         "i": [2, 3, 4],
@@ -28,10 +23,7 @@ def test_shift(request: Any, constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_shift_series(request: Any, constructor_eager: Any) -> None:
-    if "pyarrow_table" in str(constructor_eager):
-        request.applymarker(pytest.mark.xfail)
-
+def test_shift_series(constructor_eager: Any) -> None:
     df = nw.from_native(constructor_eager(data), eager_only=True)
     expected = {
         "i": [2, 3, 4],

utils/check_backend_completeness.py

@@ -0,0 +1,95 @@
+"""
+Hopefully temporary script which tracks which methods we're missing
+for the PyArrow table backend.
+
+If you implement a method, please remove it from the `MISSING` list.
+"""
+
+# ruff: noqa
+import sys
+
+import narwhals as nw
+from narwhals._arrow.dataframe import ArrowDataFrame
+
+MISSING = [
+    "DataFrame.pipe",
+    "Series.round",
+]
+
+
+class MockDataFrame:
+    # Make a little mock object so we can instantiate
+    # PandasLikeDataFrame without having pandas installed
+    def __init__(self, dataframe): ...
+
+    def __narwhals_dataframe__(self):
+        return self
+
+    @property
+    def columns(self):
+        return []
+
+    @property
+    def loc(self):
+        return self
+
+    def __getitem__(self, *args):
+        return MockSeries(self)
+
+
+class MockSeries:
+    # Make a little mock object so we can instantiate
+    # nw.DataFrame without having dataframe libraries
+    # installed
+    def __init__(self, series): ...
+
+    def __narwhals_series__(self):
+        return self
+
+    @property
+    def name(self):
+        return "a"
+
+
+if __name__ == "__main__":
+    missing = []
+    no_longer_missing = []
+
+    df_pa = ArrowDataFrame(MockDataFrame({"a": [1, 2, 3]}), backend_version=(13, 0))
+    df_nw = nw.DataFrame(
+        MockDataFrame({"a": [1, 2, 3]}),
+        level="full",
+    )
+    pa_methods = [f"DataFrame.{x}" for x in df_pa.__dir__() if not x.startswith("_")]
+    nw_methods = [f"DataFrame.{x}" for x in df_nw.__dir__() if not x.startswith("_")]
+    missing.extend(
+        [
+            x
+            for x in nw_methods
+            if x not in pa_methods and x not in MISSING and x not in {"level"}
+        ]
+    )
+    no_longer_missing.extend([x for x in MISSING if x in pa_methods and x in nw_methods])
+
+    ser_pa = df_pa["a"]
+    ser_pd = df_nw["a"]
+    pa_methods = [f"Series.{x}" for x in ser_pa.__dir__() if not x.startswith("_")]
+    nw_methods = [f"Series.{x}" for x in ser_pd.__dir__() if not x.startswith("_")]
+    missing.extend([x for x in nw_methods if x not in pa_methods and x not in MISSING])
+    no_longer_missing.extend([x for x in MISSING if x in pa_methods and x in nw_methods])
+
+    if missing:
+        print(
+            "The following have not been implemented for the Arrow backend: ",
+            sorted(missing),
+        )
+        sys.exit(1)
+
+    if no_longer_missing:
+        print(
+            "Please remove the following from MISSING in utils/check_backend_completeness.py: ",
+            sorted(no_longer_missing),
+        )
+        sys.exit(1)
+
+    sys.exit(0)