feat: support cum_sum for lazy backends

narwhals/_arrow/expr.py

@@ -6,6 +6,8 @@
 from typing import Literal
 from typing import Sequence
 
+import pyarrow.compute as pc
+
 from narwhals._arrow.expr_cat import ArrowExprCatNamespace
 from narwhals._arrow.expr_dt import ArrowExprDateTimeNamespace
 from narwhals._arrow.expr_list import ArrowExprListNamespace
@@ -417,27 +419,54 @@ def clip(self: Self, lower_bound: Any | None, upper_bound: Any | None) -> Self:
             self, "clip", lower_bound=lower_bound, upper_bound=upper_bound
         )
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
-        if not is_scalar_like(kind):
-            msg = "Only aggregation or literal operations are supported in `over` context for PyArrow."
+    def over(
+        self: Self,
+        partition_by: list[str],
+        kind: ExprKind,
+        order_by: Sequence[str] | None,
+    ) -> Self:
+        if partition_by and not is_scalar_like(kind):
+            msg = "Only aggregation or literal operations are supported in grouped `over` context for PyArrow."
             raise NotImplementedError(msg)
 
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
-            if overlap := set(output_names).intersection(keys):
-                # E.g. `df.select(nw.all().sum().over('a'))`. This is well-defined,
-                # we just don't support it yet.
-                msg = (
-                    f"Column names {overlap} appear in both expression output names and in `over` keys.\n"
-                    "This is not yet supported."
-                )
-                raise NotImplementedError(msg)
+        if not partition_by:
+            assert order_by is not None  # help type checkers  # noqa: S101
 
-            tmp = df.group_by(*keys, drop_null_keys=False).agg(self)
-            tmp = df.simple_select(*keys).join(
-                tmp, how="left", left_on=keys, right_on=keys, suffix="_right"
-            )
-            return [tmp[alias] for alias in aliases]
+            # This is something like `nw.col('a').cum_sum().order_by(key)`
+            # which we can always easily support, as it doesn't require grouping.
+            def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]:
+                native_frame = df._native_frame
+                sorting_indices = pc.sort_indices(
+                    native_frame, sort_keys=[(key, "ascending") for key in order_by]
+                )
+                native_frame = pc.take(native_frame, sorting_indices)  # type: ignore[call-overload]
+                result = self(df._from_native_frame(native_frame))
+                return [
+                    ser.scatter(sorting_indices, ser._native_series)  # type: ignore[arg-type]
+                    for ser in result
+                ]
+        else:
+
+            def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]:
+                output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
+                if overlap := set(output_names).intersection(partition_by):
+                    # E.g. `df.select(nw.all().sum().over('a'))`. This is well-defined,
+                    # we just don't support it yet.
+                    msg = (
+                        f"Column names {overlap} appear in both expression output names and in `over` keys.\n"
+                        "This is not yet supported."
+                    )
+                    raise NotImplementedError(msg)
+
+                tmp = df.group_by(*partition_by, drop_null_keys=False).agg(self)
+                tmp = df.simple_select(*partition_by).join(
+                    tmp,
+                    how="left",
+                    left_on=partition_by,
+                    right_on=partition_by,
+                    suffix="_right",
+                )
+                return [tmp[alias] for alias in aliases]
 
         return self.__class__(
             func,

narwhals/_dask/expr.py

@@ -25,6 +25,8 @@
 from narwhals.utils import generate_temporary_column_name
 
 if TYPE_CHECKING:
+    from narwhals._expression_parsing import ExprKind
+
     try:
         import dask.dataframe.dask_expr as dx
     except ModuleNotFoundError:
@@ -342,6 +344,7 @@ def shift(self: Self, n: int) -> Self:
 
     def cum_sum(self: Self, *, reverse: bool) -> Self:
         if reverse:  # pragma: no cover
+            # https://github.com/dask/dask/issues/11802
             msg = "`cum_sum(reverse=True)` is not supported with Dask backend"
             raise NotImplementedError(msg)
 
@@ -538,7 +541,12 @@ def null_count(self: Self) -> Self:
             lambda _input: _input.isna().sum().to_series(), "null_count"
         )
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
+    def over(
+        self: Self,
+        partition_by: list[str],
+        kind: ExprKind,
+        order_by: Sequence[str] | None,
+    ) -> Self:
         # pandas is a required dependency of dask so it's safe to import this
         from narwhals._pandas_like.group_by import AGGREGATIONS_TO_PANDAS_EQUIVALENT
 
@@ -549,30 +557,42 @@ def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
                 "https://narwhals-dev.github.io/narwhals/pandas_like_concepts/improve_group_by_operation/"
             )
             raise NotImplementedError(msg)
-        function_name = re.sub(r"(\w+->)", "", self._function_name)
-        try:
-            dask_function_name = AGGREGATIONS_TO_PANDAS_EQUIVALENT[function_name]
-        except KeyError:
-            msg = (
-                f"Unsupported function: {function_name} in `over` context.\n\n."
-                f"Supported functions are {', '.join(AGGREGATIONS_TO_PANDAS_EQUIVALENT)}\n"
-            )
-            raise NotImplementedError(msg) from None
+        if not partition_by:
+            assert order_by is not None  # help type checkers  # noqa: S101
 
-        def func(df: DaskLazyFrame) -> list[dx.Series]:
-            output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
-
-            with warnings.catch_warnings():
-                warnings.filterwarnings(
-                    "ignore", message=".*`meta` is not specified", category=UserWarning
-                )
-                res_native = df._native_frame.groupby(keys)[list(output_names)].transform(
-                    dask_function_name, **self._call_kwargs
+            # This is something like `nw.col('a').cum_sum().order_by(key)`
+            # which we can always easily support, as it doesn't require grouping.
+            def func(df: DaskLazyFrame) -> Sequence[dx.Series]:
+                return self(df.sort(*order_by, descending=False, nulls_last=False))
+        else:
+            function_name = re.sub(r"(\w+->)", "", self._function_name)
+            try:
+                dask_function_name = AGGREGATIONS_TO_PANDAS_EQUIVALENT[function_name]
+            except KeyError:
+                # window functions are unsupported: https://github.com/dask/dask/issues/11806
+                msg = (
+                    f"Unsupported function: {function_name} in `over` context.\n\n."
+                    f"Supported functions are {', '.join(AGGREGATIONS_TO_PANDAS_EQUIVALENT)}\n"
                 )
-            result_frame = df._from_native_frame(
-                res_native.rename(columns=dict(zip(output_names, aliases)))
-            )._native_frame
-            return [result_frame[name] for name in aliases]
+                raise NotImplementedError(msg) from None
+
+            def func(df: DaskLazyFrame) -> Sequence[dx.Series]:
+                output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
+
+                with warnings.catch_warnings():
+                    # https://github.com/dask/dask/issues/11804
+                    warnings.filterwarnings(
+                        "ignore",
+                        message=".*`meta` is not specified",
+                        category=UserWarning,
+                    )
+                    res_native = df._native_frame.groupby(partition_by)[
+                        list(output_names)
+                    ].transform(dask_function_name, **self._call_kwargs)
+                result_frame = df._from_native_frame(
+                    res_native.rename(columns=dict(zip(output_names, aliases)))
+                )._native_frame
+                return [result_frame[name] for name in aliases]
 
         return self.__class__(
             func,

narwhals/_pandas_like/expr.py

@@ -452,63 +452,98 @@ def alias_output_names(names: Sequence[str]) -> Sequence[str]:
             call_kwargs=self._call_kwargs,
         )
 
-    def over(self: Self, partition_by: list[str], kind: ExprKind) -> Self:
-        if not is_elementary_expression(self):
+    def over(
+        self: Self,
+        partition_by: list[str],
+        kind: ExprKind,
+        order_by: Sequence[str] | None = None,
+    ) -> Self:
+        if not partition_by:
+            # This is something like `nw.col('a').cum_sum().order_by(key)`
+            # which we can always easily support, as it doesn't require grouping.
+            def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
+                native_frame = df._native_frame
+                pdx = df.__native_namespace__()
+                sorting_indices = pdx.MultiIndex.from_frame(
+                    native_frame[order_by]
+                ).argsort()
+                native_frame = native_frame.iloc[sorting_indices]
+                result = self(df._from_native_frame(native_frame))
+                return [ser.scatter(sorting_indices, ser) for ser in result]
+        elif not is_elementary_expression(self):
             msg = (
                 "Only elementary expressions are supported for `.over` in pandas-like backends.\n\n"
                 "Please see: "
                 "https://narwhals-dev.github.io/narwhals/pandas_like_concepts/improve_group_by_operation/"
             )
             raise NotImplementedError(msg)
-        function_name = re.sub(r"(\w+->)", "", self._function_name)
-        try:
-            pandas_function_name = WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT[function_name]
-        except KeyError:
+        else:
+            function_name = re.sub(r"(\w+->)", "", self._function_name)
             try:
-                pandas_function_name = AGGREGATIONS_TO_PANDAS_EQUIVALENT[function_name]
+                pandas_function_name = WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT[
+                    function_name
+                ]
             except KeyError:
-                msg = (
-                    f"Unsupported function: {function_name} in `over` context.\n\n"
-                    f"Supported functions are {', '.join(WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT)}\n"
-                    f"and {', '.join(AGGREGATIONS_TO_PANDAS_EQUIVALENT)}."
-                )
-                raise NotImplementedError(msg) from None
+                try:
+                    pandas_function_name = AGGREGATIONS_TO_PANDAS_EQUIVALENT[
+                        function_name
+                    ]
+                except KeyError:
+                    msg = (
+                        f"Unsupported function: {function_name} in `over` context.\n\n"
+                        f"Supported functions are {', '.join(WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT)}\n"
+                        f"and {', '.join(AGGREGATIONS_TO_PANDAS_EQUIVALENT)}."
+                    )
+                    raise NotImplementedError(msg) from None
 
-        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
-            pandas_kwargs = window_kwargs_to_pandas_equivalent(
-                function_name, self._call_kwargs
-            )
+            def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
+                output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
+                pandas_kwargs = window_kwargs_to_pandas_equivalent(
+                    function_name, self._call_kwargs
+                )
 
-            if function_name == "cum_count":
-                plx = self.__narwhals_namespace__()
-                df = df.with_columns(~plx.col(*output_names).is_null())
-            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(
-                rename(
-                    res_native,
-                    columns=dict(zip(output_names, aliases)),
-                    implementation=self._implementation,
-                    backend_version=self._backend_version,
+                if function_name == "cum_count":
+                    plx = self.__narwhals_namespace__()
+                    df = df.with_columns(~plx.col(*output_names).is_null())
+
+                if function_name.startswith("cum_"):
+                    reverse = self._call_kwargs["reverse"]
+                else:
+                    assert "reverse" not in self._call_kwargs  # noqa: S101
+                    reverse = False
+                if order_by or reverse:
+                    # Only select the columns we need to avoid reordering columns
+                    # unnecessarily
+                    columns = list(set(partition_by).union(output_names))
+                    if order_by:
+                        pdx = df.__native_namespace__()
+                        sorting_indices = pdx.MultiIndex.from_frame(
+                            df[order_by]._native_frame
+                        ).argsort()
+                        if reverse:
+                            sorting_indices = sorting_indices[::-1]
+                        native_frame = df[columns]._native_frame.iloc[sorting_indices]
+                    else:  # reverse
+                        native_frame = df._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(
+                    rename(
+                        res_native,
+                        columns=dict(zip(output_names, aliases)),
+                        implementation=self._implementation,
+                        backend_version=self._backend_version,
+                    )
                 )
-            )
-            if reverse:
-                return [result_frame[name][::-1] for name in aliases]
-            return [result_frame[name] for name in aliases]
+                results = (result_frame[name] for name in aliases)
+                if order_by:
+                    return [s.scatter(sorting_indices, s) for s in results]
+                elif reverse:
+                    return [s[::-1] for s in results]
+                return list(results)
 
         return self.__class__(
             func,

narwhals/_polars/expr.py

@@ -97,8 +97,22 @@ def is_nan(self: Self) -> Self:
             )
         return self._from_native_expr(self._native_expr.is_nan())
 
-    def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
-        return self._from_native_expr(self._native_expr.over(keys))
+    def over(
+        self: Self,
+        partition_by: list[str],
+        kind: ExprKind,
+        order_by: Sequence[str] | None,
+    ) -> Self:
+        if self._backend_version < (1, 9):
+            if order_by:
+                msg = "`order_by` in Polars requires version 1.10 or greater"
+                raise NotImplementedError(msg)
+            return self._from_native_expr(
+                self._native_expr.over(partition_by or pl.lit(1))
+            )
+        return self._from_native_expr(
+            self._native_expr.over(partition_by or pl.lit(1), order_by=order_by)
+        )
 
     def rolling_var(
         self: Self,