Repository: narwhals-dev/narwhals
Pull request #2141 (2025-03-04 09:47:28)

feat: support reversed cumulative expressions in `over` for pandas-like
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues

- Related issue #\<issue number\>
- Closes #\<issue number\>

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below



Patch:
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 3d595520..bb7c97d0 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -51,10 +51,6 @@ WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT = {
 def window_kwargs_to_pandas_equivalent(
     function_name: str, kwargs: dict[str, object]
 ) -> dict[str, object]:
-    unsupported_reverse_msg = (
-        "Cumulative operation with `reverse=True` is not supported in "
-        "over context for pandas-like backend."
-    )
     if function_name == "shift":
         pandas_kwargs: dict[str, object] = {"periods": kwargs["n"]}
     elif function_name == "rank":
@@ -66,8 +62,6 @@ def window_kwargs_to_pandas_equivalent(
             "pct": False,
         }
     elif function_name.startswith("cum_"):  # Cumulative operation
-        if kwargs["reverse"]:
-            raise NotImplementedError(unsupported_reverse_msg)
         pandas_kwargs = {"skipna": True}
     else:
         pandas_kwargs = {}
@@ -458,7 +452,7 @@ class PandasLikeExpr(CompliantExpr["PandasLikeDataFrame", PandasLikeSeries]):
             call_kwargs=self._call_kwargs,
         )
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
+    def over(self: Self, partition_by: list[str], kind: ExprKind) -> Self:
         if not is_elementary_expression(self):
             msg = (
                 "Only elementary expressions are supported for `.over` in pandas-like backends.\n\n"
@@ -489,8 +483,19 @@ class PandasLikeExpr(CompliantExpr["PandasLikeDataFrame", PandasLikeSeries]):
             if function_name == "cum_count":
                 plx = self.__narwhals_namespace__()
                 df = df.with_columns(~plx.col(*output_names).is_null())
-
-            res_native = df._native_frame.groupby(keys)[list(output_names)].transform(
+            if function_name.startswith("cum_"):
+                reverse = self._call_kwargs["reverse"]
+            else:
+                assert "reverse" not in self._call_kwargs  # debug assertion  # noqa: S101
+                reverse = False
+            if reverse:
+                # Only select the columns we need to avoid reversing columns
+                # unnecessarily
+                columns = list(set(partition_by).union(output_names))
+                native_frame = df[columns]._native_frame[::-1]
+            else:
+                native_frame = df._native_frame
+            res_native = native_frame.groupby(partition_by)[list(output_names)].transform(
                 pandas_function_name, **pandas_kwargs
             )
             result_frame = df._from_native_frame(
@@ -501,6 +506,8 @@ class PandasLikeExpr(CompliantExpr["PandasLikeDataFrame", PandasLikeSeries]):
                     backend_version=self._backend_version,
                 )
             )
+            if reverse:
+                return [result_frame[name][::-1] for name in aliases]
             return [result_frame[name] for name in aliases]
 
         return self.__class__(


Test patch:
diff --git a/tests/expr_and_series/over_test.py b/tests/expr_and_series/over_test.py
index 4e9b2fe3..0c43f957 100644
--- a/tests/expr_and_series/over_test.py
+++ b/tests/expr_and_series/over_test.py
@@ -296,14 +296,37 @@ def test_over_diff(
     assert_equal_data(result, expected)
 
 
-def test_over_cum_reverse() -> None:
-    df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})
-
-    with pytest.raises(
-        NotImplementedError,
-        match=r"Cumulative operation with `reverse=True` is not supported",
-    ):
-        nw.from_native(df).select(nw.col("b").cum_max(reverse=True).over("a"))
+@pytest.mark.parametrize(
+    ("attr", "expected_b"),
+    [
+        ("cum_max", [5, 5, 9, None, 9]),
+        ("cum_min", [4, 5, 7, None, 9]),
+        ("cum_sum", [9, 5, 16, None, 9]),
+        ("cum_count", [2, 1, 2, 1, 1]),
+        ("cum_prod", [20, 5, 63, None, 9]),
+    ],
+)
+def test_over_cum_reverse(
+    constructor_eager: ConstructorEager,
+    request: pytest.FixtureRequest,
+    attr: str,
+    expected_b: list[object],
+) -> None:
+    if "pyarrow_table" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+    if "pandas_nullable" in str(constructor_eager) and attr in {"cum_max", "cum_min"}:
+        # https://github.com/pandas-dev/pandas/issues/61031
+        request.applymarker(pytest.mark.xfail)
+    df = constructor_eager(
+        {
+            "a": [1, 1, 2, 2, 2],
+            "b": [4, 5, 7, None, 9],
+        }
+    )
+    expr = getattr(nw.col("b"), attr)(reverse=True)
+    result = nw.from_native(df).with_columns(expr.over("a"))
+    expected = {"a": [1, 1, 2, 2, 2], "b": expected_b}
+    assert_equal_data(result, expected)
 
 
 def test_over_raise_len_change(constructor: Constructor) -> None: