Support diff() inside over() in cudf-polars

python/cudf_polars/cudf_polars/dsl/expressions/rolling.py

@@ -374,6 +374,7 @@ def __init__(
                         "rank",
                         "fill_null_with_strategy",
                         "cum_sum",
+                        "diff",
                         "shift",
                         "shift_and_fill",
                     }
@@ -612,7 +613,7 @@ def _(
             offsets.append(offset)
             out_names.append(ne.name)
             out_dtypes.append(shift_expr.dtype)
-            if shift_expr.name == "shift":
+            if shift_expr.name in {"diff", "shift"}:
                 fill_scalars.append(
                     plc.Scalar.from_py(None, plc_col.type(), stream=df.stream)
                 )
@@ -638,11 +639,29 @@ def _(
         shifted_tbl = op.local_grouper.shift(
             plc.Table(val_cols), offsets, fill_scalars, stream=df.stream
         )[1]
-        return (
-            out_names,
-            out_dtypes,
-            [plc.Table([column]) for column in shifted_tbl.columns()],
-        )
+        result_tables: list[plc.Table] = []
+        for val_col, shifted_col, ne in zip(
+            val_cols, shifted_tbl.columns(), op.named_exprs, strict=True
+        ):
+            shift_expr = ne.value
+            assert isinstance(shift_expr, expr.UnaryFunction)
+            if shift_expr.name == "diff":
+                result_tables.append(
+                    plc.Table(
+                        [
+                            plc.binaryop.binary_operation(
+                                val_col,
+                                shifted_col,
+                                plc.binaryop.BinaryOperator.SUB,
+                                shift_expr.dtype.plc_type,
+                                stream=df.stream,
+                            )
+                        ]
+                    )
+                )
+            else:
+                result_tables.append(plc.Table([shifted_col]))
+        return out_names, out_dtypes, result_tables
 
     @_apply_unary_op.register
     def _(
@@ -804,10 +823,13 @@ def _split_named_expr(
                 and v.children[0].name == "cum_sum"
             ):
                 unary_window_ops["cum_sum"].append(ne)
+            elif isinstance(v, expr.UnaryFunction) and v.name in {
+                "diff",
+                "shift_and_fill",
+            }:
+                unary_window_ops["shift"].append(ne)
             elif isinstance(v, expr.UnaryFunction) and v.name in unary_window_ops:
                 unary_window_ops[v.name].append(ne)
-            elif isinstance(v, expr.UnaryFunction) and v.name == "shift_and_fill":
-                unary_window_ops["shift"].append(ne)
             elif isinstance(v, FixedSizeRollingWindow):
                 unary_window_ops["fixed_size_rolling"].append(ne)
             else:

python/cudf_polars/cudf_polars/dsl/utils/aggregations.py

@@ -28,7 +28,7 @@
 
 
 _WINDOW_ONLY_UNARY_FUNCTIONS = frozenset(
-    {"rank", "fill_null_with_strategy", "cum_sum", "shift", "shift_and_fill"}
+    {"rank", "fill_null_with_strategy", "cum_sum", "diff", "shift", "shift_and_fill"}
 )
 
 
@@ -123,17 +123,21 @@ def decompose_single_agg(
             raise NotImplementedError(
                 f"{agg.name} over a window does not support nested fixed-size rolling"
             )
-        if agg.name in {"shift", "shift_and_fill"}:
+        if agg.name in {"diff", "shift", "shift_and_fill"}:
             if not isinstance(agg.children[1], expr.Literal):
                 raise NotImplementedError(
-                    "shift over a window only supports a literal offset"
+                    f"{agg.name} over a window only supports a literal offset"
                 )
             if agg.name == "shift_and_fill" and not isinstance(
                 agg.children[2], expr.Literal
             ):
                 raise NotImplementedError(
                     "shift over a window only supports a literal fill_value"
                 )
+        if agg.name == "diff" and agg.options[0] != "ignore":
+            raise NotImplementedError(
+                "diff over a window only supports null_behavior='ignore'"
+            )
         if agg.name == "fill_null_with_strategy" and (
             strategy := agg.options[0]
         ) not in {"forward", "backward"}:

python/cudf_polars/cudf_polars/streaming/utils.py

@@ -130,7 +130,7 @@ def _contains_unsupported_fill_strategy(exprs: Sequence[Expr]) -> bool:
     return False
 
 
-_INPUT_ORDER_WINDOW_OPS = frozenset({"cum_sum", "shift", "shift_and_fill"})
+_INPUT_ORDER_WINDOW_OPS = frozenset({"cum_sum", "diff", "shift", "shift_and_fill"})
 
 
 def _contains_input_order_window_without_order_by(exprs: Sequence[Expr]) -> bool:

python/cudf_polars/tests/expressions/test_rolling.py

@@ -303,7 +303,7 @@ def test_rank_over(
     method: RankMethod,
     *,
     descending: bool,
-    order_by: None | list[str | pl.Expr],
+    order_by: list[str | pl.Expr] | None,
 ) -> None:
     q = df.select(
         pl.col("x")
@@ -322,7 +322,7 @@ def test_rank_over_with_ties(
     method: RankMethod,
     *,
     descending: bool,
-    order_by: None | list[str | pl.Expr],
+    order_by: list[str | pl.Expr] | None,
 ) -> None:
     q = df.select(
         pl.when(pl.col("g") == 2)
@@ -343,7 +343,7 @@ def test_rank_over_with_null_values(
     method: RankMethod,
     *,
     descending: bool,
-    order_by: None | list[str | pl.Expr],
+    order_by: list[str | pl.Expr] | None,
 ) -> None:
     q = df.select(
         pl.when((pl.col("x") % 2) == 0)
@@ -364,7 +364,7 @@ def test_rank_over_with_null_group_keys(
     method: RankMethod,
     *,
     descending: bool,
-    order_by: None | list[str | pl.Expr],
+    order_by: list[str | pl.Expr] | None,
 ) -> None:
     q = df.select(
         pl.col("x")
@@ -395,7 +395,7 @@ def test_fill_over(
     engine: pl.GPUEngine,
     df: pl.LazyFrame,
     strategy: str,
-    order_by: None | list[str | pl.Expr],
+    order_by: list[str | pl.Expr] | None,
     group_key: str,
     expr: pl.Expr,
 ) -> None:
@@ -435,7 +435,7 @@ def test_cum_sum_over(
     *,
     expr: pl.Expr,
     group_key: str,
-    order_by: None | list[str | pl.Expr],
+    order_by: list[str | pl.Expr] | None,
 ) -> None:
     q = df.select(expr.cum_sum().over(group_key, order_by=order_by))
     assert_gpu_result_equal(q, engine=engine)
@@ -474,6 +474,21 @@ def test_shift_over_fill_value(
     assert_gpu_result_equal(q, engine=engine)
 
 
+@pytest.mark.parametrize("n", [1, -1, 2])
+@pytest.mark.parametrize("order_by", ["x2", None])
+def test_diff_over(
+    engine: pl.GPUEngine,
+    df: pl.LazyFrame,
+    n: int,
+    order_by: str | None,
+) -> None:
+    expr = pl.col("x").diff(n=n).over("g")
+    if order_by is not None:
+        expr = pl.col("x").diff(n=n).over("g", order_by=order_by)
+    q = df.select(expr)
+    assert_gpu_result_equal(q, engine=engine)
+
+
 @pytest.mark.parametrize(
     "expr",
     [
@@ -491,6 +506,23 @@ def test_shift_over_nonliteral_args_raises(
     assert_ir_translation_raises(q, engine, NotImplementedError)
 
 
+@pytest.mark.parametrize(
+    "expr",
+    [
+        pl.col("x").diff(n=pl.col("x2").min()).over("g"),
+        pl.col("x").diff(null_behavior="drop").over("g"),
+    ],
+    ids=["nonliteral_offset", "drop_null_behavior"],
+)
+def test_diff_over_unsupported_args_raises(
+    engine: pl.GPUEngine,
+    df: pl.LazyFrame,
+    expr: pl.Expr,
+) -> None:
+    q = df.select(expr)
+    assert_ir_translation_raises(q, engine, NotImplementedError)
+
+
 @pytest.mark.parametrize("n", [1, -1])
 def test_shift_over_without_order_by(
     engine_raise_on_fail: pl.GPUEngine,

python/cudf_polars/tests/streaming/test_rolling.py

@@ -63,6 +63,7 @@ def test_rolling_datetime(engine):
         pl.col("x").rank(method="dense", descending=True).over("g"),
         pl.col("x").rank(method="min").over("g", "g2"),
         pl.col("x").cum_sum().over("g", order_by="s"),
+        pl.col("x").diff().over("g", order_by="s"),
         pl.col("x").shift(1).over("g", order_by="s"),
         pl.col("x").shift(-1, fill_value=0).over("g", order_by="s"),
         pl.when((pl.col("x") % 2) == 0)
@@ -80,6 +81,7 @@ def test_rolling_datetime(engine):
         "rank_dense",
         "rank_min_multi_key",
         "cum_sum_order_by",
+        "diff_order_by",
         "shift_order_by",
         "shift_fill_order_by",
         "fill_null_forward",

python/cudf_polars/tests/streaming/test_spmd.py

@@ -500,6 +500,7 @@ def test_quent_context_default(spmd_engine: SPMDEngine) -> None:
     [
         (pl.col("x").sum().over("g").alias("result"), "sum"),
         (pl.col("x").rank(method="dense").over("g").alias("result"), "rank"),
+        (pl.col("x").diff().over("g", order_by="x").alias("result"), "diff"),
         (pl.col("x").shift(1).over("g", order_by="x").alias("result"), "shift"),
         pytest.param(
             pl.col("x")
@@ -513,7 +514,13 @@ def test_quent_context_default(spmd_engine: SPMDEngine) -> None:
             ),
         ),
     ],
-    ids=["scalar_sum", "nonscalar_rank", "nonscalar_shift", "nonscalar_rolling"],
+    ids=[
+        "scalar_sum",
+        "nonscalar_rank",
+        "nonscalar_diff",
+        "nonscalar_shift",
+        "nonscalar_rolling",
+    ],
 )
 @pytest.mark.parametrize(
     "cross_rank",
@@ -572,6 +579,8 @@ def test_over_multirank(
                 assert grp["result"].to_list() == [sum(expected_xs)] * 3
             elif expected == "rank":
                 assert grp["result"].to_list() == [1, 2, 3]
+            elif expected == "diff":
+                assert grp["result"].to_list() == [None, 1, 1]
             elif expected == "shift":
                 assert grp["result"].to_list() == [None, *expected_xs[:-1]]
             else:
@@ -585,6 +594,9 @@ def test_over_multirank(
     "expr,expected",
     [
         (pl.col("x").shift(1).over("g").alias("result"), "shift"),
+        (pl.col("x").diff().over("g").alias("result"), "diff"),
+        (pl.col("x").diff(n=2).over("g").alias("result"), "diff_n2"),
+        (pl.col("x").diff(n=-1).over("g").alias("result"), "diff_nneg1"),
         (pl.col("x").cum_sum().over("g").alias("result"), "cum_sum"),
         pytest.param(
             pl.col("x").rolling_mean(window_size=2).over("g").alias("result"),
@@ -606,7 +618,15 @@ def test_over_multirank(
             ),
         ),
     ],
-    ids=["shift", "cum_sum", "fixed_rolling", "fixed_rolling_ordered"],
+    ids=[
+        "shift",
+        "diff",
+        "diff_n2",
+        "diff_nneg1",
+        "cum_sum",
+        "fixed_rolling",
+        "fixed_rolling_ordered",
+    ],
 )
 def test_over_shared_group_ordering_multirank(
     comm: Communicator,
@@ -648,6 +668,12 @@ def test_over_shared_group_ordering_multirank(
         expected_values: list[float | int | None]
         if expected == "shift":
             expected_values = [None, *xs[:-1]]
+        elif expected == "diff":
+            expected_values = [None, *([1] * (len(xs) - 1))]
+        elif expected == "diff_n2":
+            expected_values = [None, None, *([2] * (len(xs) - 2))]
+        elif expected == "diff_nneg1":
+            expected_values = [*([-1] * (len(xs) - 1)), None]
         elif expected == "cum_sum":
             total = 0
             expected_values = []