Support shift over grouped windows in cudf-polars

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

@@ -62,6 +62,11 @@ class CumSumOp(UnaryOp):
     pass
 
 
+@dataclass(frozen=True)
+class ShiftOp(UnaryOp):
+    pass
+
+
 def to_request(
     value: expr.Expr, orderby: Column, df: DataFrame
 ) -> plc.rolling.RollingRequest:
@@ -359,7 +364,13 @@ def __init__(
                 or (
                     isinstance(named_expr.value, expr.UnaryFunction)
                     and named_expr.value.name
-                    in {"rank", "fill_null_with_strategy", "cum_sum"}
+                    in {
+                        "rank",
+                        "fill_null_with_strategy",
+                        "cum_sum",
+                        "shift",
+                        "shift_and_fill",
+                    }
                 )
             )
         ]
@@ -569,6 +580,64 @@ def _(  # type: ignore[no-untyped-def]
                 result_tables.append(filled)
         return out_names, out_dtypes, result_tables
 
+    @_apply_unary_op.register
+    def _(
+        self,
+        op: ShiftOp,
+        df: DataFrame,
+        _: plc.groupby.GroupBy,
+    ) -> tuple[list[str], list[DataType], list[plc.Table]]:
+        plc_cols: list[plc.Column] = []
+        offsets: list[int] = []
+        fill_scalars: list[plc.Scalar] = []
+        out_names: list[str] = []
+        out_dtypes: list[DataType] = []
+
+        for ne in op.named_exprs:
+            shift_expr = ne.value
+            assert isinstance(shift_expr, expr.UnaryFunction)
+            data_expr, offset_expr = shift_expr.children[:2]
+            assert isinstance(offset_expr, expr.Literal)
+            offset = offset_expr.value
+            assert isinstance(offset, int)
+
+            plc_col = data_expr.evaluate(df, context=ExecutionContext.FRAME).obj
+            plc_cols.append(plc_col)
+            offsets.append(offset)
+            out_names.append(ne.name)
+            out_dtypes.append(shift_expr.dtype)
+            if shift_expr.name == "shift":
+                fill_scalars.append(
+                    plc.Scalar.from_py(None, plc_col.type(), stream=df.stream)
+                )
+            else:
+                assert shift_expr.name == "shift_and_fill"
+                fill_expr = shift_expr.children[2]
+                assert isinstance(fill_expr, expr.Literal)
+                fill_scalars.append(
+                    plc.Scalar.from_py(
+                        fill_expr.value, plc_col.type(), stream=df.stream
+                    )
+                )
+
+        assert op.order_index is not None
+        val_cols = plc.copying.gather(
+            plc.Table(plc_cols),
+            op.order_index,
+            plc.copying.OutOfBoundsPolicy.NULLIFY,
+            stream=df.stream,
+        ).columns()
+
+        assert isinstance(op.local_grouper, plc.groupby.GroupBy)
+        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()],
+        )
+
     def _reorder_to_input(
         self,
         row_id: plc.Column,
@@ -627,6 +696,7 @@ def _split_named_expr(
             "rank": [],
             "fill_null_with_strategy": [],
             "cum_sum": [],
+            "shift": [],
         }
 
         for ne in self.named_aggs:
@@ -640,6 +710,8 @@ def _split_named_expr(
                 unary_window_ops["cum_sum"].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)
             else:
                 reductions.append(ne)
         return reductions, unary_window_ops
@@ -1119,6 +1191,48 @@ def do_evaluate(  # noqa: D102
                 )
             )
 
+        if shift_named := unary_window_ops["shift"]:
+            order_index, shift_by_cols_for_scan, local = (
+                self._grouped_window_scan_setup(
+                    by_cols,
+                    row_id=row_id,
+                    order_by_col=order_by_col
+                    if self._order_by_expr is not None
+                    else None,
+                    ob_desc=self.options[2]
+                    if self._order_by_expr is not None
+                    else False,
+                    ob_nulls_last=self.options[3]
+                    if self._order_by_expr is not None
+                    else False,
+                    grouper=grouper,
+                    stream=df.stream,
+                    require_sorted_groups=True,
+                )
+            )
+            names, dtypes, tables = self._apply_unary_op(
+                ShiftOp(
+                    named_exprs=shift_named,
+                    order_index=order_index,
+                    by_cols_for_scan=shift_by_cols_for_scan,
+                    local_grouper=local,
+                ),
+                df,
+                grouper,
+            )
+            broadcasted_cols.extend(
+                self._reorder_to_input(
+                    row_id,
+                    by_cols,
+                    df.num_rows,
+                    tables,
+                    names,
+                    dtypes,
+                    order_index=order_index,
+                    stream=df.stream,
+                )
+            )
+
         # Create a temporary DataFrame with the broadcasted columns named by their
         # placeholder names from agg decomposition, then evaluate the post-expression.
         df = DataFrame(broadcasted_cols, stream=df.stream)

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -129,7 +129,8 @@ def _unsupported_fill_over_window(value: expr.Expr) -> bool:
     windowed = [
         node
         for node in traversal([value])
-        if isinstance(node, expr.UnaryFunction) and node.name in {"rank", "cum_sum"}
+        if isinstance(node, expr.UnaryFunction)
+        and node.name in {"rank", "cum_sum", "shift", "shift_and_fill"}
     ]
     if not windowed:
         return False
@@ -1229,7 +1230,14 @@ def _(
             if isinstance(v, expr.Agg)
             or (
                 isinstance(v, expr.UnaryFunction)
-                and v.name in {"rank", "fill_null_with_strategy", "cum_sum"}
+                and v.name
+                in {
+                    "rank",
+                    "fill_null_with_strategy",
+                    "cum_sum",
+                    "shift",
+                    "shift_and_fill",
+                }
             )
         ]
         children = (*by_exprs, *((order_by_expr,) if has_order_by else ()), *child_deps)

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

@@ -98,11 +98,24 @@ def decompose_single_agg(
         "rank",
         "fill_null_with_strategy",
         "cum_sum",
+        "shift",
+        "shift_and_fill",
     }:
         if context != ExecutionContext.WINDOW:
             raise NotImplementedError(
                 f"{agg.name} is not supported in groupby or rolling context"
             )
+        if agg.name in {"shift", "shift_and_fill"}:
+            if not isinstance(agg.children[1], expr.Literal):
+                raise NotImplementedError(
+                    "shift 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 == "fill_null_with_strategy" and (
             strategy := agg.options[0]
         ) not in {"forward", "backward"}:

python/cudf_polars/cudf_polars/streaming/select.py

@@ -25,7 +25,7 @@
 from cudf_polars.streaming.over import _fuse_over_nodes
 from cudf_polars.streaming.repartition import Repartition
 from cudf_polars.streaming.utils import (
-    _contains_cum_sum_without_order_by,
+    _contains_input_order_window_without_order_by,
     _contains_unsupported_fill_strategy,
     _dynamic_planning_on,
     _lower_ir_fallback,
@@ -414,15 +414,16 @@ def _(
             ),
         )
 
-    if rec.state["nranks"] > 1 and _contains_cum_sum_without_order_by(
+    if rec.state["nranks"] > 1 and _contains_input_order_window_without_order_by(
         [e.value for e in ir.exprs]
     ):
         return _lower_ir_fallback(
             ir.reconstruct([child]),
             rec,
             msg=(
-                "cum_sum() over a window without order_by is not supported across "
-                "multiple ranks; falling back to a single partition."
+                "input-order-sensitive window expressions without order_by are "
+                "not supported across multiple ranks; falling back to a single "
+                "partition."
             ),
         )
 

python/cudf_polars/cudf_polars/streaming/utils.py

@@ -125,8 +125,11 @@ def _contains_unsupported_fill_strategy(exprs: Sequence[Expr]) -> bool:
     return False
 
 
-def _contains_cum_sum_without_order_by(exprs: Sequence[Expr]) -> bool:
-    # Returns True for cum_sum(...) or fill_null_with_strategy(cum_sum(...)).
+_INPUT_ORDER_WINDOW_OPS = frozenset({"cum_sum", "shift", "shift_and_fill"})
+
+
+def _contains_input_order_window_without_order_by(exprs: Sequence[Expr]) -> bool:
+    """Return True for implicit input-order-sensitive window expressions."""
     for e in traversal(exprs):
         if not (isinstance(e, GroupedWindow) and not e.options[1]):
             continue
@@ -138,6 +141,6 @@ def _contains_cum_sum_without_order_by(exprs: Sequence[Expr]) -> bool:
                 and isinstance(v.children[0], UnaryFunction)
             ):
                 v = v.children[0]
-            if isinstance(v, UnaryFunction) and v.name == "cum_sum":
+            if isinstance(v, UnaryFunction) and v.name in _INPUT_ORDER_WINDOW_OPS:
                 return True
     return False

python/cudf_polars/tests/expressions/test_rolling.py

@@ -441,6 +441,66 @@ def test_cum_sum_over(
     assert_gpu_result_equal(q, engine=engine)
 
 
+@pytest.mark.parametrize("n", [1, -1, 2])
+@pytest.mark.parametrize(
+    "expr,group_key",
+    [
+        (pl.col("x"), "g"),
+        (pl.when((pl.col("x") % 2) == 0).then(None).otherwise(pl.col("x")), "g"),
+        (pl.col("x"), "g_null"),
+    ],
+)
+@pytest.mark.parametrize("order_by", ["x2", ["g2", pl.col("x2") * 2]])
+def test_shift_over(
+    engine: pl.GPUEngine,
+    df: pl.LazyFrame,
+    n: int,
+    expr: pl.Expr,
+    group_key: str,
+    order_by: str | list[str | pl.Expr],
+) -> None:
+    q = df.select(expr.shift(n).over(group_key, order_by=order_by))
+    assert_gpu_result_equal(q, engine=engine)
+
+
+@pytest.mark.parametrize("n,fill_value", [(1, 0), (-1, 99)])
+def test_shift_over_fill_value(
+    engine: pl.GPUEngine,
+    df: pl.LazyFrame,
+    n: int,
+    fill_value: int,
+) -> None:
+    q = df.select(pl.col("x").shift(n, fill_value=fill_value).over("g", order_by="x2"))
+    assert_gpu_result_equal(q, engine=engine)
+
+
+@pytest.mark.parametrize(
+    "expr",
+    [
+        pl.col("x").shift(pl.col("x2").min()).over("g"),
+        pl.col("x").shift(1, fill_value=pl.col("x2").min()).over("g"),
+    ],
+    ids=["nonliteral_offset", "nonliteral_fill_value"],
+)
+def test_shift_over_nonliteral_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,
+    df: pl.LazyFrame,
+    n: int,
+) -> None:
+    q = df.select(pl.col("x").shift(n).over("g"))
+    assert_gpu_result_equal(q, engine=engine_raise_on_fail)
+
+
 @pytest.mark.parametrize(
     "expr",
     [