Repository: narwhals-dev/narwhals

Issue:
enh: support `nw.col('a').shift(n).over('b')`
We now support cumulative operations in `over` for pandas and Polars

We should do the same for `shift`, as pandas also supports that. One challenge will be figuring out how to pass the argument down, but it should be doable

cc @ClaudioSalvatoreArcidiacono in case this interests you

Discussion:
MarcoGorelli: nice, thanks @ClaudioSalvatoreArcidiacono !

I think @FBruzzesi 's work in https://github.com/narwhals-dev/narwhals/pull/1614 might make this easier, perhaps we should wait to get that in, then we can update here?
FBruzzesi: Thanks @ClaudioSalvatoreArcidiacono !

I was about to echo what @MarcoGorelli  just mentioned! Instead of an ad-hoc change, I think waiting for #1614 would make it easier. Possibly it will be ready by end of today.
ClaudioSalvatoreArcidiacono: Definetly! Nice PR @FBruzzesi 
MarcoGorelli: that PR's merged, if you fancy updating @ClaudioSalvatoreArcidiacono I think we'd love to get this in!
FBruzzesi: > that PR's merged, if you fancy updating @ClaudioSalvatoreArcidiacono I think we'd love to get this in!

Just a headsup that the PR was shipped with a bug, namely it does not track kwargs in expr composition. You would probably want to wait for #1645 which contains a fix

Patch:
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index e7e98d99..07543011 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
     from narwhals.utils import Implementation
     from narwhals.utils import Version
 
-CUMULATIVE_FUNCTIONS_TO_PANDAS_EQUIVALENT = {
+MANY_TO_MANY_AGG_FUNCTIONS_TO_PANDAS_EQUIVALENT = {
     "col->cum_sum": "cumsum",
     "col->cum_min": "cummin",
     "col->cum_max": "cummax",
@@ -33,6 +33,7 @@ CUMULATIVE_FUNCTIONS_TO_PANDAS_EQUIVALENT = {
     # Pandas cumcount counts nulls while Polars does not
     # So, instead of using "cumcount" we use "cumsum" on notna() to get the same result
     "col->cum_count": "cumsum",
+    "col->shift": "shift",
 }
 
 
@@ -417,7 +418,7 @@ class PandasLikeExpr(CompliantExpr[PandasLikeSeries]):
         )
 
     def over(self, keys: list[str]) -> Self:
-        if self._function_name in CUMULATIVE_FUNCTIONS_TO_PANDAS_EQUIVALENT:
+        if self._function_name in MANY_TO_MANY_AGG_FUNCTIONS_TO_PANDAS_EQUIVALENT:
 
             def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 if (
@@ -443,12 +444,18 @@ class PandasLikeExpr(CompliantExpr[PandasLikeSeries]):
                     plx = self.__narwhals_namespace__()
                     df = df.with_columns(~plx.col(*self._root_names).is_null())
 
+                if self._function_name == "col->shift":
+                    kwargs = {"periods": self._kwargs.get("n", 1)}
+                else:
+                    # Cumulative operation
+                    kwargs = {"skipna": True}
+
                 res_native = getattr(
                     df._native_frame.groupby(list(keys), as_index=False)[
                         self._root_names
                     ],
-                    CUMULATIVE_FUNCTIONS_TO_PANDAS_EQUIVALENT[self._function_name],
-                )(skipna=True)
+                    MANY_TO_MANY_AGG_FUNCTIONS_TO_PANDAS_EQUIVALENT[self._function_name],
+                )(**kwargs)
 
                 result_frame = df._from_native_frame(
                     rename(


Test patch:
diff --git a/tests/expr_and_series/over_test.py b/tests/expr_and_series/over_test.py
index d0d64acf..e824c6b7 100644
--- a/tests/expr_and_series/over_test.py
+++ b/tests/expr_and_series/over_test.py
@@ -165,6 +165,23 @@ def test_over_anonymous() -> None:
         nw.from_native(df).select(nw.all().cum_max().over("a"))
 
 
+def test_over_shift(request: pytest.FixtureRequest, constructor: Constructor) -> None:
+    if "pyarrow_table_constructor" in str(
+        constructor
+    ) or "dask_lazy_p2_constructor" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor(data))
+    expected = {
+        "a": ["a", "a", "b", "b", "b"],
+        "b": [1, 2, 3, 5, 3],
+        "c": [5, 4, 3, 2, 1],
+        "b_shift": [None, None, None, None, 3],
+    }
+    result = df.with_columns(b_shift=nw.col("b").shift(2).over("a"))
+    assert_equal_data(result, expected)
+
+
 def test_over_cum_reverse() -> None:
     df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})