feat: support cum_sum for lazy backends

[MarcoGorelli]

Demo of this work:

For eager backends, just like now, you can keep using cum_sum liberally, there are no new restrictions introduced by this PR:

In [16]: df
Out[16]: 
┌──────────────────┐
|Narwhals DataFrame|
|------------------|
|     a  b  c  i   |
|  0  a  1  5  0   |
|  1  a  2  4  1   |
|  2  b  3  3  2   |
|  3  b  5  2  4   |
|  4  b  3  1  3   |
└──────────────────┘

In [17]: df.with_columns(b_cum_sum = (nw.col('b')**2).cum_sum())
Out[17]: 
┌────────────────────────┐
|   Narwhals DataFrame   |
|------------------------|
|   a  b  c  i  b_cum_sum|
|0  a  1  5  0          1|
|1  a  2  4  1          5|
|2  b  3  3  2         14|
|3  b  5  2  4         39|
|4  b  3  1  3         48|
└────────────────────────┘

You can also optionally specify _order_by (for now private, whilst we build out the functionality):

In [18]: df.with_columns(b_cum_sum = (nw.col('b')**2).cum_sum().over(_order_by='i'))
Out[18]: 
┌────────────────────────┐
|   Narwhals DataFrame   |
|------------------------|
|   a  b  c  i  b_cum_sum|
|0  a  1  5  0          1|
|1  a  2  4  1          5|
|2  b  3  3  2         14|
|3  b  5  2  4         48|
|4  b  3  1  3         23|
└────────────────────────┘

You can also partition by a column, but for pandas there's the usual limitation that only elementary expressions are supported:

In [19]: df.with_columns(b_cum_sum = nw.col('b').cum_sum().over('a', _order_by='i'))
Out[19]: 
┌────────────────────────┐
|   Narwhals DataFrame   |
|------------------------|
|   a  b  c  i  b_cum_sum|
|0  a  1  5  0          1|
|1  a  2  4  1          3|
|2  b  3  3  2          3|
|3  b  5  2  4         11|
|4  b  3  1  3          6|
└────────────────────────┘

And now, for the new functionality which this unlocks. Here is where we start enabling new things. For lazy backends, the above is supported, but specifying _order_by is required. Example using SQLFrame:

In [29]: lf
Out[29]: 
┌────────────────────────────────────────────────────────────────────┐
|                         Narwhals LazyFrame                         |
|--------------------------------------------------------------------|
|<sqlframe.duckdb.dataframe.DuckDBDataFrame object at 0x7f4f12317260>|
└────────────────────────────────────────────────────────────────────┘

In [30]: lf.to_native().show()
+---+---+---+---+
| a | b | c | i |
+---+---+---+---+
| a | 1 | 5 | 0 |
| a | 2 | 4 | 1 |
| b | 3 | 3 | 2 |
| b | 5 | 2 | 4 |
| b | 3 | 1 | 3 |
+---+---+---+---+

In [31]: lf.with_columns(b_cum_sum = nw.col('b').cum_sum().over(_order_by='i')).to_native().show()
+---+---+---+---+-----------+
| a | b | c | i | b_cum_sum |
+---+---+---+---+-----------+
| a | 1 | 5 | 0 |     1     |
| a | 2 | 4 | 1 |     3     |
| b | 3 | 3 | 2 |     6     |
| b | 3 | 1 | 3 |     9     |
| b | 5 | 2 | 4 |     14    |
+---+---+---+---+-----------+

In [32]: lf.with_columns(b_cum_sum = nw.col('b').cum_sum().over('a', _order_by='i')).to_native().show()
+---+---+---+---+-----------+
| a | b | c | i | b_cum_sum |
+---+---+---+---+-----------+
| a | 1 | 5 | 0 |     1     |
| a | 2 | 4 | 1 |     3     |
| b | 3 | 3 | 2 |     3     |
| b | 3 | 1 | 3 |     6     |
| b | 5 | 2 | 4 |     11    |
+---+---+---+---+-----------+

The implementation for spark-like looks like this:

https://github.com/MarcoGorelli/narwhals/blob/b9d4529a5756ef178e0b89cf244e786c63d2a0c8/narwhals/_spark_like/expr.py#L536-L550

then over just applies that window function with the given window:

https://github.com/MarcoGorelli/narwhals/blob/b9d4529a5756ef178e0b89cf244e786c63d2a0c8/narwhals/_spark_like/expr.py#L496-L506

At the Narwhals level, we enforce that, for lazyframes, window functions like cum_sum must be immediately followed by over with _order_by specified

We should be able to adapt this fairly straightforwardly to also cover:

• rolling_*
• diff
• shift
• is_*_distinct
• cum_*
• rank

All of these should be supportable immediately for SQLFrame / PySpark / Polars Lazy. Missing lazy backends are:

• Dask: blocked by groupby-transform-cumsum returns results which look wrong dask/dask#11806
• DuckDB: they don't have WindowExpression in their Python API. But, SQLFrame backed by DuckDB should be fine 👍

[MarcoGorelli]

thanks @FBruzzesi , excellent review!

have fixed the partition_by='' case

for nulls in order_by, have fixed + tested. if there's duplicates in order_by, then results aren't stable (we don't provide stable sorting guarantees) so i'm not sure there's much we should do there

[dangotbanned]

@MarcoGorelli I feel like I'm missing something 🤔

What is the relation between ArrowExpr.over and cum_sum for lazy backends?

[MarcoGorelli]

sure, thanks for asking

in nw.LazyFrame, cum_sum must be followed by over - there's some examples here #2132 (comment)

[MarcoGorelli]

right, let's go ahead with this, thanks all for reviews!