Repository: narwhals-dev/narwhals
Pull request #113 (2024-05-08 13:40:59)

add cum_sum


Patch:
diff --git a/docs/api-reference/expressions.md b/docs/api-reference/expressions.md
index c487ff0f..ba096720 100644
--- a/docs/api-reference/expressions.md
+++ b/docs/api-reference/expressions.md
@@ -8,6 +8,7 @@
         - all
         - any
         - cast
+        - cum_sum
         - diff
         - drop_nulls
         - fill_null
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index e1171c31..84bcf363 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -8,6 +8,7 @@
         - all
         - any
         - cast
+        - cum_sum
         - diff
         - drop_nulls
         - dtype
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 180fa4da..35730d8d 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -196,6 +196,9 @@ class PandasExpr:
     def n_unique(self) -> Self:
         return register_expression_call(self, "n_unique")
 
+    def cum_sum(self) -> Self:
+        return register_expression_call(self, "cum_sum")
+
     def unique(self) -> Self:
         return register_expression_call(self, "unique")
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 28747366..802e2106 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -357,13 +357,18 @@ class PandasSeries:
         ser = self._series
         return self._from_series(ser.sample(n=n, frac=fraction, replace=with_replacement))
 
-    def unique(self) -> PandasSeries:
+    def cum_sum(self) -> PandasSeries:
         return self._from_series(
             self._series.__class__(
-                self._series.unique(), dtype=self._series.dtype, name=self._series.name
+                self._series.cumsum(), dtype=self._series.dtype, name=self._series.name
             )
         )
 
+    def unique(self) -> PandasSeries:
+        return self._from_series(
+            self._series.__class__(self._series.unique(), name=self._series.name)
+        )
+
     def diff(self) -> PandasSeries:
         return self._from_series(
             self._series.__class__(self._series.diff(), name=self._series.name)
diff --git a/narwhals/expression.py b/narwhals/expression.py
index e7100d4e..2cf4afdd 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -490,7 +490,7 @@ class Expr:
 
     def unique(self) -> Expr:
         """
-         Returns unique values
+        Return unique values
 
         Examples:
             >>> import polars as pl
@@ -527,6 +527,49 @@ class Expr:
         """
         return self.__class__(lambda plx: self._call(plx).unique())
 
+    def cum_sum(self) -> Expr:
+        """
+        Return cumulative sum.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> df_pd = pd.DataFrame({'a': [1, 1, 3, 5, 5], 'b': [2, 4, 4, 6, 6]})
+            >>> df_pl = pl.DataFrame({'a': [1, 1, 3, 5, 5], 'b': [2, 4, 4, 6, 6]})
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...    df = nw.from_native(df_any)
+            ...    df = df.select(nw.col('a', 'b').cum_sum())
+            ...    return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                a   b
+            0   1   2
+            1   2   6
+            2   5  10
+            3  10  16
+            4  15  22
+            >>> func(df_pl)
+            shape: (5, 2)
+            ┌─────┬─────┐
+            │ a   ┆ b   │
+            │ --- ┆ --- │
+            │ i64 ┆ i64 │
+            ╞═════╪═════╡
+            │ 1   ┆ 2   │
+            │ 2   ┆ 6   │
+            │ 5   ┆ 10  │
+            │ 10  ┆ 16  │
+            │ 15  ┆ 22  │
+            └─────┴─────┘
+        """
+        return self.__class__(lambda plx: self._call(plx).cum_sum())
+
     def diff(self) -> Expr:
         """
         Returns the difference between each element and the previous one.
diff --git a/narwhals/series.py b/narwhals/series.py
index ef549c6f..7aca6465 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -125,6 +125,43 @@ class Series:
     def drop_nulls(self) -> Self:
         return self._from_series(self._series.drop_nulls())
 
+    def cum_sum(self) -> Self:
+        """
+        Calculate the cumulative sum.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s = [2, 4, 3]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+
+            We define a data-frame agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     s = s.cum_sum()
+            ...     return nw.to_native(s)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    2
+            1    6
+            2    9
+            dtype: int64
+            >>> func(s_pl)  # doctest:+SKIP
+            shape: (3,)
+            Series: '' [i64]
+            [
+               2
+               6
+               9
+            ]
+        """
+        return self._from_series(self._series.cum_sum())
+
     def unique(self) -> Self:
         return self._from_series(self._series.unique())
 


Test patch:
diff --git a/tests/expr/cum_sum_test.py b/tests/expr/cum_sum_test.py
new file mode 100644
index 00000000..1be553cc
--- /dev/null
+++ b/tests/expr/cum_sum_test.py
@@ -0,0 +1,32 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {
+    "a": [0, 1, 2, 3, 4],
+    "b": [1, 2, 3, 5, 3],
+    "c": [5, 4, 3, 2, 1],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_over_single(constructor: Any) -> None:
+    df = nw.from_native(constructor(data), eager_only=True)
+    result = df.with_columns(nw.all().cum_sum())
+    expected = {
+        "a": [0, 1, 3, 6, 10],
+        "b": [1, 3, 6, 11, 14],
+        "c": [5, 9, 12, 14, 15],
+    }
+    compare_dicts(result, expected)
+    result = df.select(
+        df["a"].cum_sum(),
+        df["b"].cum_sum(),
+        df["c"].cum_sum(),
+    )
+    compare_dicts(result, expected)