Feat: Series.shift Pyarrow Backend Implementation

narwhals/_arrow/expr.py

@@ -194,6 +194,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

@@ -209,6 +209,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,32 @@
+from __future__ import annotations
+
+from typing import Any
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+
+def test_shift(constructor: 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(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)

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,10 +11,7 @@
 }
 
 
-def test_shift(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_shift(constructor: Any) -> None:
     df = nw.from_native(constructor(data), eager_only=True)
     result = df.with_columns(nw.col("a", "b", "c").shift(2)).filter(nw.col("i") > 1)
     expected = {

utils/check_backend_completeness.py

@@ -14,7 +14,6 @@
 MISSING = [
     "DataFrame.pipe",
     "Series.round",
-    "Series.shift",
 ]