shift -> kernel, add fancy test

Author: dangotbanned

narwhals/_plan/arrow/expr.py

@@ -398,22 +398,8 @@ def rolling_expr(self, node: ir.RollingExpr, frame: Frame, name: str) -> Self:
         raise NotImplementedError
 
     def shift(self, node: ir.FunctionExpr[Shift], frame: Frame, name: str) -> Self:
-        n = node.function.n
         series = self._dispatch_expr(node.input[0], frame, name)
-        native = series.native
-        if n == 0:
-            return self._with_native(native, name)
-        # NOTE: This might be zero-copy and the original wasn't?
-        # here, the non-nulls stay in their own chunk
-        # on main, the two (or more) chunks are merged
-        if n > 0:
-            arrays = [
-                pa.nulls(n, native.type),
-                *native.slice(length=native.length() - n).chunks,
-            ]
-        else:
-            arrays = [*native.slice(offset=-n).chunks, pa.nulls(-n, native.type)]
-        return self._with_native(fn.chunked_array(arrays), name)
+        return self._with_native(fn.shift(series.native, node.function.n), name)
 
     def diff(self, node: ir.FunctionExpr[Diff], frame: Frame, name: str) -> Self:
         series = self._dispatch_expr(node.input[0], frame, name)

narwhals/_plan/arrow/functions.py

@@ -262,6 +262,17 @@ def diff(native: ChunkedOrArrayT) -> ChunkedOrArrayT:
     )
 
 
+def shift(native: ChunkedArrayAny, n: int) -> ChunkedArrayAny:
+    if n == 0:
+        return native
+    arr = native
+    if n > 0:
+        arrays = [nulls_like(n, arr), *arr.slice(length=arr.length() - n).chunks]
+    else:
+        arrays = [*arr.slice(offset=-n).chunks, nulls_like(-n, arr)]
+    return pa.chunked_array(arrays)
+
+
 def is_between(
     native: ChunkedOrScalar[ScalarT],
     lower: ChunkedOrScalar[ScalarT],
@@ -325,6 +336,14 @@ def int_range(
     return pa.chunked_array([pa.array(np.arange(start, end, step), dtype)])
 
 
+def nulls_like(n: int, native: ArrowAny) -> ArrayAny:
+    """Create a strongly-typed Array instance with all elements null.
+
+    Uses the type of `native`.
+    """
+    return pa.nulls(n, native.type)  # type: ignore[no-any-return]
+
+
 def lit(value: Any, dtype: DataType | None = None) -> NativeScalar:
     return pa.scalar(value) if dtype is None else pa.scalar(value, dtype)
 

tests/plan/over_test.py

@@ -27,6 +27,11 @@ def data() -> Data:
     }
 
 
+@pytest.fixture
+def data_alt() -> Data:
+    return {"a": [3, 5, 1, 2, None], "b": [0, 1, 3, 2, 1], "c": [9, 1, 2, 1, 1]}
+
+
 @pytest.mark.parametrize(
     "partition_by",
     [
@@ -177,3 +182,16 @@ def test_len_over_2369() -> None:
     result = df.with_columns(a_len_per_group=nwp.len().over("b")).sort("a")
     expected = {"a": [1, 2, 4], "b": ["x", "x", "y"], "a_len_per_group": [2, 2, 1]}
     assert_equal_data(result, expected)
+
+
+def test_shift_kitchen_sink(data_alt: Data) -> None:
+    result = dataframe(data_alt).select(
+        nwp.nth(1, 2)
+        .shift(-1)
+        .over(order_by=nwp.nth(0))
+        .sort(nulls_last=True)
+        .fill_null(100)
+        * 5
+    )
+    expected = {"b": [0, 5, 10, 15, 500], "c": [5, 5, 10, 45, 500]}
+    assert_equal_data(result, expected)