diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index 4adcf611388b..bdcd6c737641 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/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: diff --git a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py index dc242d734125..4b66ebb1846a 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py +++ b/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,10 +123,10 @@ 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 @@ -134,6 +134,10 @@ def decompose_single_agg( 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"}: diff --git a/python/cudf_polars/cudf_polars/streaming/utils.py b/python/cudf_polars/cudf_polars/streaming/utils.py index 2355f4215e4f..98e775b582f2 100644 --- a/python/cudf_polars/cudf_polars/streaming/utils.py +++ b/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: diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index c579e22abde6..268cca2f1cb0 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/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, diff --git a/python/cudf_polars/tests/streaming/test_rolling.py b/python/cudf_polars/tests/streaming/test_rolling.py index fd6fd34c1e68..de1682c173eb 100644 --- a/python/cudf_polars/tests/streaming/test_rolling.py +++ b/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", diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index ea397c3addab..eb9ef7703b32 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/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 = []