WIP: Implement expression-based rolling

Still need to do the grouped version but that will be easier.

Author: wence-

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

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 # TODO: remove need for this
 # ruff: noqa: D101
@@ -8,24 +8,108 @@
 
 from typing import TYPE_CHECKING, Any
 
-from cudf_polars.dsl.expressions.base import Expr
+import pylibcudf as plc
+
+from cudf_polars.containers import Column
+from cudf_polars.dsl import expr
+from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
+from cudf_polars.dsl.utils.windows import range_window_bounds
 
 if TYPE_CHECKING:
-    import pylibcudf as plc
+    import pyarrow as pa
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.typing import ClosedInterval
+
+__all__ = ["GroupedRollingWindow", "RollingWindow", "to_request"]
+
+
+def to_request(
+    value: expr.Expr, orderby: Column, df: DataFrame
+) -> plc.rolling.RollingRequest:
+    """
+    Produce a rolling request for evaluation with pylibcudf.
 
-__all__ = ["GroupedRollingWindow", "RollingWindow"]
+    Parameters
+    ----------
+    value
+        The expression to perform the rolling aggregation on.
+    orderby
+        Orderby column, used as input to the request when the aggregation is Len.
+    df
+        DataFrame used to evaluate the inputs to the aggregation.
+    """
+    min_periods = 1
+    if isinstance(value, expr.Len):
+        # A count aggregation, we need a column so use the orderby column
+        col = orderby
+    elif isinstance(value, expr.Agg):
+        child = value.children[0]
+        col = child.evaluate(df, context=ExecutionContext.ROLLING)
+        if value.name == "var":
+            # Polars variance produces null if nvalues <= ddof
+            # libcudf produces NaN. However, we can get the polars
+            # behaviour by setting the minimum window size to ddof +
+            # 1.
+            min_periods = value.options + 1
+    else:
+        col = value.evaluate(df, context=ExecutionContext.ROLLING)
+    return plc.rolling.RollingRequest(col.obj, min_periods, value.agg_request)
 
 
 class RollingWindow(Expr):
-    __slots__ = ("options",)
-    _non_child = ("dtype", "options")
+    __slots__ = ("closed_window", "following", "orderby", "preceding")
+    _non_child = ("dtype", "preceding", "following", "closed_window", "orderby")
 
-    def __init__(self, dtype: plc.DataType, options: Any, agg: Expr) -> None:
+    def __init__(
+        self,
+        dtype: plc.DataType,
+        preceding: pa.Scalar,
+        following: pa.Scalar,
+        closed_window: ClosedInterval,
+        orderby: str,
+        agg: Expr,
+    ) -> None:
         self.dtype = dtype
-        self.options = options
+        self.preceding = preceding
+        self.following = following
+        self.closed_window = closed_window
+        self.orderby = orderby
         self.children = (agg,)
         self.is_pointwise = False
-        raise NotImplementedError("Rolling window not implemented")
+        if not plc.rolling.is_valid_rolling_aggregation(agg.dtype, agg.agg_request):
+            raise NotImplementedError(f"Unsupported rolling aggregation {agg}")
+
+    def do_evaluate(  # noqa: D102
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        if context != ExecutionContext.FRAME:
+            raise RuntimeError("Rolling aggregation inside groupby/over/rolling")
+        (agg,) = self.children
+        orderby = df.column_map[self.orderby]
+        preceding, following = range_window_bounds(
+            self.preceding, self.following, self.closed_window
+        )
+        if orderby.obj.null_count() != 0:
+            raise RuntimeError(
+                f"Index column '{self.orderby}' in rolling may not contain nulls"
+            )
+        if not orderby.check_sorted(
+            order=plc.types.Order.ASCENDING, null_order=plc.types.NullOrder.AFTER
+        ):
+            raise RuntimeError(
+                f"Index column '{self.orderby}' in rolling is not sorted, please sort first"
+            )
+        (result,) = plc.rolling.grouped_range_rolling_window(
+            plc.Table([]),
+            orderby.obj,
+            plc.types.Order.ASCENDING,
+            plc.types.NullOrder.AFTER,
+            preceding,
+            following,
+            [to_request(agg, orderby, df)],
+        ).columns()
+        return Column(result)
 
 
 class GroupedRollingWindow(Expr):

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -30,6 +30,7 @@
 
 import cudf_polars.dsl.expr as expr
 from cudf_polars.containers import Column, DataFrame
+from cudf_polars.dsl.expressions import rolling
 from cudf_polars.dsl.nodebase import Node
 from cudf_polars.dsl.to_ast import to_ast, to_parquet_filter
 from cudf_polars.dsl.utils.windows import range_window_bounds
@@ -1016,22 +1017,7 @@ def do_evaluate(
         df: DataFrame,
     ) -> DataFrame:
         keys = broadcast(*(k.evaluate(df) for k in keys_in), target_length=df.num_rows)
-        requests: list[plc.rolling.RollingRequest] = []
-        names: list[str] = []
         orderby = index.evaluate(df)
-        for request in aggs:
-            name = request.name
-            value = request.value
-            if isinstance(value, expr.Len):
-                # A count aggregation, we need a column so use the orderby column
-                col = orderby.obj
-            elif isinstance(value, expr.Agg):
-                (child,) = value.children
-                col = child.evaluate(df).obj
-            else:
-                col = value.evaluate(df).obj
-            requests.append(plc.rolling.RollingRequest(col, 1, value.agg_request))
-            names.append(name)
         preceding_window, following_window = range_window_bounds(
             preceding, following, closed_window
         )
@@ -1048,19 +1034,23 @@ def do_evaluate(
         values = plc.rolling.grouped_range_rolling_window(
             plc.Table([k.obj for k in keys]),
             orderby.obj,
-            orderby.order,
+            plc.types.Order.ASCENDING,  # Polars requires ascending orderby.
             plc.types.NullOrder.AFTER,  # Doesn't matter, polars doesn't allow nulls
             preceding_window,
             following_window,
-            requests,
+            [rolling.to_request(request.value, orderby, df) for request in aggs],
         )
         return DataFrame(
             itertools.chain(
                 keys,
                 [orderby],
                 (
                     Column(col, name=name)
-                    for col, name in zip(values.columns(), names, strict=True)
+                    for col, name in zip(
+                        values.columns(),
+                        (request.name for request in aggs),
+                        strict=True,
+                    )
                 ),
             )
         ).slice(zlice)

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -22,8 +22,13 @@
 
 from cudf_polars.dsl import expr, ir
 from cudf_polars.dsl.to_ast import insert_colrefs
+from cudf_polars.dsl.traversal import traversal
+from cudf_polars.dsl.utils.aggregations import decompose_single_agg
 from cudf_polars.dsl.utils.groupby import rewrite_groupby
+from cudf_polars.dsl.utils.naming import unique_names
+from cudf_polars.dsl.utils.replace import replace
 from cudf_polars.dsl.utils.rolling import rewrite_rolling
+from cudf_polars.dsl.utils.windows import offsets_to_windows
 from cudf_polars.typing import Schema
 from cudf_polars.utils import config, dtypes, sorting
 
@@ -51,6 +56,7 @@ def __init__(self, visitor: NodeTraverser, engine: GPUEngine):
         self.visitor = visitor
         self.config_options = config.ConfigOptions.from_polars_engine(engine)
         self.errors: list[Exception] = []
+        self.schema_stack = []
 
     def translate_ir(self, *, n: int | None = None) -> ir.IR:
         """
@@ -106,10 +112,13 @@ def translate_ir(self, *, n: int | None = None) -> ir.IR:
                 self.errors.append(e)
                 return ir.ErrorNode(schema, str(e))
             try:
+                self.schema_stack.append(schema)
                 result = _translate_ir(node, self, schema)
             except Exception as e:
                 self.errors.append(e)
                 return ir.ErrorNode(schema, str(e))
+            finally:
+                self.schema_stack.pop()
             if any(
                 isinstance(dtype, pl.Null)
                 for dtype in pl.datatypes.unpack_dtypes(*polars_schema.values())
@@ -621,9 +630,44 @@ def _(node: pl_expr.Window, translator: Translator, dtype: plc.DataType) -> expr
     # TODO: raise in groupby?
     if isinstance(node.options, pl_expr.RollingGroupOptions):
         # pl.col("a").rolling(...)
-        return expr.RollingWindow(
-            dtype, node.options, translator.translate_expr(n=node.function)
+        agg = translator.translate_expr(n=node.function)
+        name_generator = unique_names(
+            e.name for e in traversal([agg]) if isinstance(e, expr.Col)
         )
+        named_aggs, named_post_agg, _ = decompose_single_agg(
+            expr.NamedExpr(next(name_generator), agg), name_generator, is_top=True
+        )
+        orderby = node.options.index_column
+        preceding, following = offsets_to_windows(
+            translator.schema_stack[-1][orderby],
+            node.options.offset,
+            node.options.period,
+        )
+        closed_window = node.options.closed_window
+        if isinstance(named_post_agg.value, expr.Col):
+            (named_agg,) = named_aggs
+            return expr.RollingWindow(
+                named_agg.value.dtype,
+                preceding,
+                following,
+                closed_window,
+                orderby,
+                named_agg.value,
+            )
+        return replace(
+            [named_post_agg.value],  # type: ignore[misc]
+            {
+                expr.Col(agg.value.dtype, agg.name): expr.RollingWindow(
+                    agg.value.dtype,
+                    preceding,
+                    following,
+                    closed_window,
+                    orderby,
+                    agg.value,
+                )
+                for agg in named_aggs
+            },
+        )[0]
     elif isinstance(node.options, pl_expr.WindowMapping):
         # pl.col("a").over(...)
         return expr.GroupedRollingWindow(

python/cudf_polars/cudf_polars/dsl/utils/replace.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for replacing nodes in a DAG."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from cudf_polars.typing import GenericTransformer, NodeT
+
+__all__ = ["replace"]
+
+
+def _replace(node: NodeT, fn: GenericTransformer[NodeT, NodeT]) -> NodeT:
+    try:
+        return fn.state["replacements"][node]
+    except KeyError:
+        return reuse_if_unchanged(node, fn)
+
+
+def replace(nodes: Sequence[NodeT], replacements: Mapping[NodeT, NodeT]) -> list[NodeT]:
+    """
+    Replace nodes in expressions.
+
+    Parameters
+    ----------
+    nodes
+        Sequence of nodes to perform replacements in.
+    replacements
+        Mapping from nodes to be replaced to their replacements.
+
+    Returns
+    -------
+    list
+        Of nodes with replacements performed.
+    """
+    mapper: GenericTransformer[NodeT, NodeT] = CachingVisitor(
+        _replace, state={"replacements": replacements}
+    )
+    return [mapper(node) for node in nodes]