From 6f9cc22d1862b12cf057f53349d612c91def33e6 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 11 Aug 2026 13:41:02 -0700 Subject: [PATCH 1/7] initial support for sum().rolling().over() --- .../cudf_polars/dsl/expressions/rolling.py | 190 +++++++++++++++++- .../cudf_polars/cudf_polars/dsl/translate.py | 22 +- .../cudf_polars/dsl/utils/aggregations.py | 12 ++ .../tests/expressions/test_rolling.py | 45 +++++ .../tests/streaming/test_rolling.py | 49 +++++ 5 files changed, 301 insertions(+), 17 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index bdcd6c737641..95213811cdba 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -72,8 +72,16 @@ class FixedSizeRollingOp(UnaryOp): pass +@dataclass(frozen=True) +class RollingWindowOp(UnaryOp): + pass + + def to_request( - value: expr.Expr, orderby: Column, df: DataFrame + value: expr.Expr, + orderby: Column, + df: DataFrame, + order_index: plc.Column | None = None, ) -> plc.rolling.RollingRequest: """ Produce a rolling request for evaluation with pylibcudf. @@ -86,14 +94,16 @@ def to_request( Orderby column, used as input to the request when the aggregation is Len. df DataFrame used to evaluate the inputs to the aggregation. + order_index + Optional row ordering to apply to the request column. """ min_periods = 1 if isinstance(value, expr.Len): # A count aggregation, we need a column so use the orderby column - col = orderby + col_obj = orderby.obj elif isinstance(value, expr.Agg): child = value.children[0] - col = child.evaluate(df, context=ExecutionContext.ROLLING) + col_obj = child.evaluate(df, context=ExecutionContext.ROLLING).obj if (POLARS_VERSION_LT_136 or not POLARS_VERSION_LT_139) and value.name == "var": # Polars variance produces null if nvalues <= ddof # libcudf produces NaN. However, we can get the polars @@ -104,11 +114,18 @@ def to_request( # See https://github.com/pola-rs/polars/pull/25117 min_periods = value.options + 1 else: - col = value.evaluate( + col_obj = value.evaluate( df, context=ExecutionContext.ROLLING - ) # pragma: no cover; raise before we get here because we + ).obj # pragma: no cover; raise before we get here because we # don't do correct handling of empty groups - return plc.rolling.RollingRequest(col.obj, min_periods, value.agg_request) + if order_index is not None: + col_obj = plc.copying.gather( + plc.Table([col_obj]), + order_index, + plc.copying.OutOfBoundsPolicy.NULLIFY, + stream=df.stream, + ).columns()[0] + return plc.rolling.RollingRequest(col_obj, min_periods, value.agg_request) class RollingWindow(Expr): # pragma: no cover; polars >1.36 uses AExpr::Rolling now @@ -365,7 +382,8 @@ def __init__( for named_expr in self.named_aggs if not ( isinstance( - named_expr.value, (expr.Len, expr.Agg, FixedSizeRollingWindow) + named_expr.value, + (expr.Len, expr.Agg, FixedSizeRollingWindow, RollingWindow), ) or ( isinstance(named_expr.value, expr.UnaryFunction) @@ -752,6 +770,116 @@ def get_window_key(ne: expr.NamedExpr) -> tuple[int, int]: return names, dtypes, tables + @_apply_unary_op.register + def _( + self, + op: RollingWindowOp, + df: DataFrame, + _: plc.groupby.GroupBy, + ) -> tuple[list[str], list[DataType], list[plc.Table]]: + assert op.order_index is not None + assert op.by_cols_for_scan is not None + + orderby_names = { + rolling_expr.orderby + for ne in op.named_exprs + for rolling_expr in (ne.value,) + if isinstance(rolling_expr, RollingWindow) + } + if len(orderby_names) != 1: + raise NotImplementedError( + "rolling(...).over(...) only supports one rolling index column" + ) + (orderby_name,) = orderby_names + orderby = df.column_map[orderby_name] + sorted_orderby_obj = plc.copying.gather( + plc.Table([orderby.obj]), + op.order_index, + plc.copying.OutOfBoundsPolicy.NULLIFY, + stream=df.stream, + ).columns()[0] + if ( + plc.traits.is_integral(sorted_orderby_obj.type()) + and sorted_orderby_obj.type().id() != plc.TypeId.INT64 + ): + sorted_orderby_obj = plc.unary.cast( + sorted_orderby_obj, plc.DataType(plc.TypeId.INT64), stream=df.stream + ) + if sorted_orderby_obj.null_count() != 0: + raise RuntimeError( + f"Index column '{orderby_name}' in rolling may not contain nulls" + ) + + request_groups: dict[tuple[int, int, Any], list[expr.NamedExpr]] = defaultdict( + list + ) + for ne in op.named_exprs: + rolling_expr = ne.value + assert isinstance(rolling_expr, RollingWindow) + request_groups[ + ( + rolling_expr.preceding_ordinal, + rolling_expr.following_ordinal, + rolling_expr.closed_window, + ) + ].append(ne) + + group_keys = plc.Table([c.obj for c in op.by_cols_for_scan]) + results_by_name: dict[str, tuple[DataType, plc.Table]] = {} + for window_exprs in request_groups.values(): + sample = window_exprs[0].value + assert isinstance(sample, RollingWindow) + preceding_scalar, following_scalar = offsets_to_windows( + sample.orderby_dtype, + sample.preceding_ordinal, + sample.following_ordinal, + stream=df.stream, + ) + preceding, following = range_window_bounds( + preceding_scalar, following_scalar, sample.closed_window + ) + requests: list[plc.rolling.RollingRequest] = [] + for ne in window_exprs: + rolling_expr = ne.value + assert isinstance(rolling_expr, RollingWindow) + requests.append( + to_request( + rolling_expr.children[0], + orderby, + df, + order_index=op.order_index, + ) + ) + + result_cols = plc.rolling.grouped_range_rolling_window( + group_keys, + sorted_orderby_obj, + plc.types.Order.ASCENDING, + plc.types.NullOrder.BEFORE, + preceding, + following, + requests, + stream=df.stream, + ).columns() + + for ne, result_col in zip(window_exprs, result_cols, strict=True): + rolling_expr = ne.value + assert isinstance(rolling_expr, RollingWindow) + output_col = result_col + if result_col.type() != rolling_expr.dtype.plc_type: + output_col = plc.unary.cast( + result_col, rolling_expr.dtype.plc_type, stream=df.stream + ) + results_by_name[ne.name] = ( + rolling_expr.dtype, + plc.Table([output_col]), + ) + + names = [ne.name for ne in op.named_exprs] + dtypes = [results_by_name[name][0] for name in names] + tables = [results_by_name[name][1] for name in names] + return names, dtypes, tables + def _reorder_to_input( self, row_id: plc.Column, @@ -812,6 +940,7 @@ def _split_named_expr( "cum_sum": [], "shift": [], "fixed_size_rolling": [], + "range_rolling": [], } for ne in self.named_aggs: @@ -832,6 +961,8 @@ def _split_named_expr( unary_window_ops[v.name].append(ne) elif isinstance(v, FixedSizeRollingWindow): unary_window_ops["fixed_size_rolling"].append(ne) + elif isinstance(v, RollingWindow): + unary_window_ops["range_rolling"].append(ne) else: reductions.append(ne) return reductions, unary_window_ops @@ -1353,6 +1484,51 @@ def do_evaluate( # noqa: D102 ) ) + if rolling_named := unary_window_ops["range_rolling"]: + orderby_names = { + rolling_expr.orderby + for ne in rolling_named + for rolling_expr in (ne.value,) + if isinstance(rolling_expr, RollingWindow) + } + if len(orderby_names) != 1: + raise NotImplementedError( + "rolling(...).over(...) only supports one rolling index column" + ) + (orderby_name,) = orderby_names + rolling_order_by_col = df.column_map[orderby_name] + order_index, rolling_by_cols_for_scan, _ = self._grouped_window_scan_setup( + by_cols, + row_id=row_id, + order_by_col=rolling_order_by_col, + ob_desc=False, + ob_nulls_last=False, + grouper=grouper, + stream=df.stream, + require_sorted_groups=True, + ) + names, dtypes, tables = self._apply_unary_op( + RollingWindowOp( + named_exprs=rolling_named, + order_index=order_index, + by_cols_for_scan=rolling_by_cols_for_scan, + ), + 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, + ) + ) + if fixed_rolling_named := unary_window_ops["fixed_size_rolling"]: order_index, rolling_by_cols_for_scan, _ = self._grouped_window_scan_setup( by_cols, diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 8316879b0c12..541edd4300d9 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -1216,30 +1216,32 @@ def _( translator.translate_expr(n=n, schema=schema) for n in node.partition_by ] - child_deps = [ - v.children[0].children[0] + child_deps: list[expr.Expr] = [] + for ne in named_aggs: + v = ne.value if ( isinstance(v, expr.UnaryFunction) and v.name == "fill_null_with_strategy" and isinstance(v.children[0], expr.UnaryFunction) and v.children[0].name == "cum_sum" - ) - else v.children[0] - for ne in named_aggs - for v in (ne.value,) - if isinstance(v, expr.Agg) - or ( + ): + child_deps.append(v.children[0].children[0]) + elif isinstance(v, expr.RollingWindow): + child_deps.append(v.children[0]) + child_deps.append(expr.Col(schema[v.orderby], v.orderby)) + elif isinstance(v, expr.Agg) or ( isinstance(v, expr.UnaryFunction) and v.name in { "rank", "fill_null_with_strategy", "cum_sum", + "diff", "shift", "shift_and_fill", } - ) - ] + ): + child_deps.append(v.children[0]) children = (*by_exprs, *((order_by_expr,) if has_order_by else ()), *child_deps) return expr.GroupedWindow( dtype, diff --git a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py index 4b66ebb1846a..f6c88112995b 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py @@ -163,6 +163,18 @@ def decompose_single_agg( "window-only unary expressions" ) return [(named_expr, True)], named_expr.reconstruct(expr.Col(agg.dtype, name)) + if isinstance(agg, expr.RollingWindow): + if context != ExecutionContext.WINDOW: + raise NotImplementedError( + "Range rolling is not supported in groupby or rolling context" + ) + if _contains_window_only_unary(agg.children[0]) or ( + _contains_fixed_size_rolling_window(agg.children[0]) + ): + raise NotImplementedError( + "Range rolling over a window does not support nested window expressions" + ) + return [(named_expr, True)], named_expr.reconstruct(expr.Col(agg.dtype, name)) if isinstance(agg, expr.UnaryFunction) and agg.name == "null_count": (child,) = agg.children diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index 268cca2f1cb0..f93047fce6d6 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -3,6 +3,7 @@ from __future__ import annotations +import datetime as dt from typing import TYPE_CHECKING, Literal, cast import pytest @@ -185,6 +186,50 @@ def test_rolling_sum_all_null_window_returns_null(engine: pl.GPUEngine): assert_gpu_result_equal(q, engine=engine) +@skip_rolling_expr_136_to_138 +def test_rolling_sum_over(engine: pl.GPUEngine) -> None: + df = ( + pl.LazyFrame( + { + "ric": ["A", "A", "A", "B", "B", "B"], + "ts": [ + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 1), + dt.datetime(2025, 1, 1, 9, 3), + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 2), + dt.datetime(2025, 1, 1, 9, 3), + ], + "price": [10.0, 11.0, 12.0, 20.0, 21.0, 22.0], + "volume": [100, 200, 300, 400, 500, 600], + } + ) + .with_columns(notional=pl.col("price") * pl.col("volume")) + .sort("ric", "ts") + ) + q = df.with_columns( + volume_before=pl.col("volume") + .sum() + .rolling("ts", period="2m", offset="-2m", closed="left") + .over("ric"), + notional_before=pl.col("notional") + .sum() + .rolling("ts", period="2m", offset="-2m", closed="left") + .over("ric"), + volume_after=pl.col("volume") + .sum() + .rolling("ts", period="2m", closed="right") + .over("ric"), + ).select( + "ric", + "ts", + "volume_before", + "notional_before", + "volume_after", + ) + assert_gpu_result_equal(q, engine=engine) + + @pytest.mark.parametrize( "expr", [ diff --git a/python/cudf_polars/tests/streaming/test_rolling.py b/python/cudf_polars/tests/streaming/test_rolling.py index de1682c173eb..030dfbf03154 100644 --- a/python/cudf_polars/tests/streaming/test_rolling.py +++ b/python/cudf_polars/tests/streaming/test_rolling.py @@ -3,6 +3,8 @@ from __future__ import annotations +import datetime as dt + import pytest import polars as pl @@ -100,6 +102,53 @@ def test_over_select(engine, expr): assert_gpu_result_equal(df.select(expr), engine=engine, check_row_order=True) +@pytest.mark.skipif( + not POLARS_VERSION_LT_136 and POLARS_VERSION_LT_139, + reason="Rolling window expressions are not accessible in polars 1.36-1.38", +) +def test_rolling_sum_over(engine): + df = ( + pl.LazyFrame( + { + "ric": ["A", "A", "A", "B", "B", "B"], + "ts": [ + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 1), + dt.datetime(2025, 1, 1, 9, 3), + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 2), + dt.datetime(2025, 1, 1, 9, 3), + ], + "price": [10.0, 11.0, 12.0, 20.0, 21.0, 22.0], + "volume": [100, 200, 300, 400, 500, 600], + } + ) + .with_columns(notional=pl.col("price") * pl.col("volume")) + .sort("ric", "ts") + ) + q = df.with_columns( + volume_before=pl.col("volume") + .sum() + .rolling("ts", period="2m", offset="-2m", closed="left") + .over("ric"), + notional_before=pl.col("notional") + .sum() + .rolling("ts", period="2m", offset="-2m", closed="left") + .over("ric"), + volume_after=pl.col("volume") + .sum() + .rolling("ts", period="2m", closed="right") + .over("ric"), + ).select( + "ric", + "ts", + "volume_before", + "notional_before", + "volume_after", + ) + assert_gpu_result_equal(q, engine=engine, check_row_order=True) + + @pytest.mark.parametrize("strategy", ["forward", "backward"]) def test_over_cum_sum_fill_null_per_partition(engine, strategy): df = pl.LazyFrame( From 8e298e05448a9be6f77d2c9d86247728c7892f9e Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 11 Aug 2026 13:49:41 -0700 Subject: [PATCH 2/7] extend coverage to all basic aggs --- .../tests/expressions/test_rolling.py | 20 ++++++++++++++++ .../tests/streaming/test_rolling.py | 23 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index f93047fce6d6..9645ef77847d 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -230,6 +230,26 @@ def test_rolling_sum_over(engine: pl.GPUEngine) -> None: assert_gpu_result_equal(q, engine=engine) +@skip_rolling_expr_136_to_138 +def test_rolling_common_aggs_over(engine: pl.GPUEngine) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "A", "B", "B", "B"], + "ts": [1, 2, 4, 1, 3, 4], + "x": [100, 200, 300, 400, 500, 600], + } + ).sort("g", "ts") + q = df.select( + pl.col("x").sum().rolling("ts", period="2i").over("g").alias("sum"), + pl.col("x").min().rolling("ts", period="2i").over("g").alias("min"), + pl.col("x").max().rolling("ts", period="2i").over("g").alias("max"), + pl.col("x").mean().rolling("ts", period="2i").over("g").alias("mean"), + pl.col("x").count().rolling("ts", period="2i").over("g").alias("count"), + pl.len().rolling("ts", period="2i").over("g").alias("len"), + ) + assert_gpu_result_equal(q, engine=engine) + + @pytest.mark.parametrize( "expr", [ diff --git a/python/cudf_polars/tests/streaming/test_rolling.py b/python/cudf_polars/tests/streaming/test_rolling.py index 030dfbf03154..810598ed6a71 100644 --- a/python/cudf_polars/tests/streaming/test_rolling.py +++ b/python/cudf_polars/tests/streaming/test_rolling.py @@ -149,6 +149,29 @@ def test_rolling_sum_over(engine): assert_gpu_result_equal(q, engine=engine, check_row_order=True) +@pytest.mark.skipif( + not POLARS_VERSION_LT_136 and POLARS_VERSION_LT_139, + reason="Rolling window expressions are not accessible in polars 1.36-1.38", +) +def test_rolling_common_aggs_over(engine): + df = pl.LazyFrame( + { + "g": ["A", "A", "A", "B", "B", "B"], + "ts": [1, 2, 4, 1, 3, 4], + "x": [100, 200, 300, 400, 500, 600], + } + ).sort("g", "ts") + q = df.select( + pl.col("x").sum().rolling("ts", period="2i").over("g").alias("sum"), + pl.col("x").min().rolling("ts", period="2i").over("g").alias("min"), + pl.col("x").max().rolling("ts", period="2i").over("g").alias("max"), + pl.col("x").mean().rolling("ts", period="2i").over("g").alias("mean"), + pl.col("x").count().rolling("ts", period="2i").over("g").alias("count"), + pl.len().rolling("ts", period="2i").over("g").alias("len"), + ) + assert_gpu_result_equal(q, engine=engine, check_row_order=True) + + @pytest.mark.parametrize("strategy", ["forward", "backward"]) def test_over_cum_sum_fill_null_per_partition(engine, strategy): df = pl.LazyFrame( From 55180fecd1d3579a73f409441f794ab9c8e12e71 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 12 Aug 2026 10:57:37 -0700 Subject: [PATCH 3/7] reduce some overlap --- .../cudf_polars/dsl/expressions/rolling.py | 252 +++++++----------- 1 file changed, 101 insertions(+), 151 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index 95213811cdba..5f314479a068 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -7,7 +7,7 @@ from __future__ import annotations from collections import Counter, defaultdict -from dataclasses import dataclass +from dataclasses import dataclass, replace from functools import singledispatchmethod from typing import TYPE_CHECKING, Any, ClassVar @@ -119,6 +119,8 @@ def to_request( ).obj # pragma: no cover; raise before we get here because we # don't do correct handling of empty groups if order_index is not None: + assert order_index.size() == df.num_rows + assert col_obj.size() == df.num_rows col_obj = plc.copying.gather( plc.Table([col_obj]), order_index, @@ -780,17 +782,7 @@ def _( assert op.order_index is not None assert op.by_cols_for_scan is not None - orderby_names = { - rolling_expr.orderby - for ne in op.named_exprs - for rolling_expr in (ne.value,) - if isinstance(rolling_expr, RollingWindow) - } - if len(orderby_names) != 1: - raise NotImplementedError( - "rolling(...).over(...) only supports one rolling index column" - ) - (orderby_name,) = orderby_names + orderby_name = self._rolling_orderby_name(op.named_exprs) orderby = df.column_map[orderby_name] sorted_orderby_obj = plc.copying.gather( plc.Table([orderby.obj]), @@ -967,6 +959,19 @@ def _split_named_expr( reductions.append(ne) return reductions, unary_window_ops + @staticmethod + def _rolling_orderby_name(named_exprs: Sequence[expr.NamedExpr]) -> str: + orderby_names: set[str] = set() + for ne in named_exprs: + rolling_expr = ne.value + assert isinstance(rolling_expr, RollingWindow) + orderby_names.add(rolling_expr.orderby) + if len(orderby_names) != 1: + raise NotImplementedError( + "rolling(...).over(...) only supports one rolling index column" + ) + return next(iter(orderby_names)) + def _build_window_order_index( self, by_cols: list[Column], @@ -1115,6 +1120,50 @@ def _broadcast_agg_results( for name, dtype, col in zip(names, dtypes, out_cols, strict=True) ] + def _apply_ordered_unary_op( + self, + op: UnaryOp, + df: DataFrame, + grouper: plc.groupby.GroupBy, + by_cols: list[Column], + row_id: plc.Column, + *, + order_by_col: Column | None, + ob_desc: bool = False, + ob_nulls_last: bool = False, + require_sorted_groups: bool = False, + ) -> list[Column]: + order_index, by_cols_for_scan, local = self._grouped_window_scan_setup( + by_cols, + row_id=row_id, + order_by_col=order_by_col, + ob_desc=ob_desc, + ob_nulls_last=ob_nulls_last, + grouper=grouper, + stream=df.stream, + require_sorted_groups=require_sorted_groups, + ) + names, dtypes, tables = self._apply_unary_op( + replace( + op, + order_index=order_index, + by_cols_for_scan=by_cols_for_scan, + local_grouper=local, + ), + df, + grouper, + ) + return self._reorder_to_input( + row_id, + by_cols, + df.num_rows, + tables, + names, + dtypes, + order_index=order_index, + stream=df.stream, + ) + def _build_groupby_requests( self, named_exprs: list[expr.NamedExpr], @@ -1274,6 +1323,11 @@ def do_evaluate( # noqa: D102 plc.Scalar.from_py(1, plc.types.SIZE_TYPE, stream=df.stream), stream=df.stream, ) + over_order_by_col = order_by_col if self._order_by_expr is not None else None + over_ob_desc = self.options[2] if self._order_by_expr is not None else False + over_ob_nulls_last = ( + self.options[3] if self._order_by_expr is not None else False + ) if rank_named := unary_window_ops["rank"]: if self._order_by_expr is not None: @@ -1401,166 +1455,62 @@ def do_evaluate( # noqa: D102 and ne.value.name == "fill_null_with_strategy" for ne in cum_named ) - order_index, cum_sum_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=has_fill, - ) - ) - names, dtypes, tables = self._apply_unary_op( - CumSumOp( - named_exprs=cum_named, - order_index=order_index, - by_cols_for_scan=cum_sum_by_cols_for_scan, - local_grouper=local, - ), - df, - grouper, - ) broadcasted_cols.extend( - self._reorder_to_input( - row_id, + self._apply_ordered_unary_op( + CumSumOp(named_exprs=cum_named), + df, + grouper, by_cols, - df.num_rows, - tables, - names, - dtypes, - order_index=order_index, - stream=df.stream, + row_id, + order_by_col=over_order_by_col, + ob_desc=over_ob_desc, + ob_nulls_last=over_ob_nulls_last, + require_sorted_groups=has_fill, ) ) 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, + self._apply_ordered_unary_op( + ShiftOp(named_exprs=shift_named), + df, + grouper, by_cols, - df.num_rows, - tables, - names, - dtypes, - order_index=order_index, - stream=df.stream, + row_id, + order_by_col=over_order_by_col, + ob_desc=over_ob_desc, + ob_nulls_last=over_ob_nulls_last, + require_sorted_groups=True, ) ) if rolling_named := unary_window_ops["range_rolling"]: - orderby_names = { - rolling_expr.orderby - for ne in rolling_named - for rolling_expr in (ne.value,) - if isinstance(rolling_expr, RollingWindow) - } - if len(orderby_names) != 1: - raise NotImplementedError( - "rolling(...).over(...) only supports one rolling index column" - ) - (orderby_name,) = orderby_names + orderby_name = self._rolling_orderby_name(rolling_named) rolling_order_by_col = df.column_map[orderby_name] - order_index, rolling_by_cols_for_scan, _ = self._grouped_window_scan_setup( - by_cols, - row_id=row_id, - order_by_col=rolling_order_by_col, - ob_desc=False, - ob_nulls_last=False, - grouper=grouper, - stream=df.stream, - require_sorted_groups=True, - ) - names, dtypes, tables = self._apply_unary_op( - RollingWindowOp( - named_exprs=rolling_named, - order_index=order_index, - by_cols_for_scan=rolling_by_cols_for_scan, - ), - df, - grouper, - ) broadcasted_cols.extend( - self._reorder_to_input( - row_id, + self._apply_ordered_unary_op( + RollingWindowOp(named_exprs=rolling_named), + df, + grouper, by_cols, - df.num_rows, - tables, - names, - dtypes, - order_index=order_index, - stream=df.stream, + row_id, + order_by_col=rolling_order_by_col, + require_sorted_groups=True, ) ) if fixed_rolling_named := unary_window_ops["fixed_size_rolling"]: - order_index, rolling_by_cols_for_scan, _ = 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( - FixedSizeRollingOp( - named_exprs=fixed_rolling_named, - order_index=order_index, - by_cols_for_scan=rolling_by_cols_for_scan, - ), - df, - grouper, - ) broadcasted_cols.extend( - self._reorder_to_input( - row_id, + self._apply_ordered_unary_op( + FixedSizeRollingOp(named_exprs=fixed_rolling_named), + df, + grouper, by_cols, - df.num_rows, - tables, - names, - dtypes, - order_index=order_index, - stream=df.stream, + row_id, + order_by_col=over_order_by_col, + ob_desc=over_ob_desc, + ob_nulls_last=over_ob_nulls_last, + require_sorted_groups=True, ) ) From 565efcc0abaa550990be0a296f1aee5184edb478 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 13 Aug 2026 07:52:26 -0700 Subject: [PATCH 4/7] address CI failure --- .../cudf_polars/cudf_polars/dsl/translate.py | 6 +- .../cudf_polars/dsl/utils/aggregations.py | 5 ++ .../tests/expressions/test_rolling.py | 69 ++++++++++++++++++- 3 files changed, 77 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 541edd4300d9..288732182f7c 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -1206,6 +1206,10 @@ def _( named_aggs = [agg for agg, _ in aggs] for named_agg in named_aggs: + if has_order_by and isinstance(named_agg.value, expr.RollingWindow): + raise NotImplementedError( + "rolling(...).over(..., order_by=...) is not supported" + ) if _unsupported_fill_over_window(named_agg.value): raise NotImplementedError( "fill_null with strategy over a window is only supported when " @@ -1229,7 +1233,7 @@ def _( elif isinstance(v, expr.RollingWindow): child_deps.append(v.children[0]) child_deps.append(expr.Col(schema[v.orderby], v.orderby)) - elif isinstance(v, expr.Agg) or ( + elif isinstance(v, (expr.FixedSizeRollingWindow, expr.Agg)) or ( isinstance(v, expr.UnaryFunction) and v.name in { diff --git a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py index f6c88112995b..2c4842cc95b7 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py @@ -38,6 +38,10 @@ def _contains_fixed_size_rolling_window(value: expr.Expr) -> bool: ) +def _contains_range_rolling_window(value: expr.Expr) -> bool: + return any(isinstance(node, expr.RollingWindow) for node in traversal([value])) + + def _contains_window_only_unary(value: expr.Expr) -> bool: return any( isinstance(node, expr.UnaryFunction) @@ -170,6 +174,7 @@ def decompose_single_agg( ) if _contains_window_only_unary(agg.children[0]) or ( _contains_fixed_size_rolling_window(agg.children[0]) + or _contains_range_rolling_window(agg.children[0]) ): raise NotImplementedError( "Range rolling over a window does not support nested window expressions" diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index 9645ef77847d..48e91f3346fb 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -9,6 +9,7 @@ import pytest import polars as pl +from polars.testing import assert_frame_equal from cudf_polars.testing.asserts import ( assert_gpu_result_equal, @@ -227,7 +228,39 @@ def test_rolling_sum_over(engine: pl.GPUEngine) -> None: "notional_before", "volume_after", ) - assert_gpu_result_equal(q, engine=engine) + expected = pl.DataFrame( + { + "ric": ["A", "A", "A", "B", "B", "B"], + "ts": [ + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 1), + dt.datetime(2025, 1, 1, 9, 3), + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 2), + dt.datetime(2025, 1, 1, 9, 3), + ], + "volume_before": [0, 100, 200, 0, 400, 500], + "notional_before": [0.0, 1000.0, 2200.0, 0.0, 8000.0, 10500.0], + "volume_after": [100, 300, 300, 400, 500, 1100], + } + ) + assert_frame_equal(q.collect(engine=engine), expected) + + +@skip_rolling_expr_136_to_138 +def test_rolling_over_with_order_by_raises(engine: pl.GPUEngine) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "A"], + "seq": [1, 2, 3], + "ts": [1, 2, 3], + "x": [10, 20, 30], + } + ) + q = df.select( + pl.col("x").sum().rolling("ts", period="2i").over("g", order_by="seq") + ) + assert_ir_translation_raises(q, engine, NotImplementedError) @skip_rolling_expr_136_to_138 @@ -247,7 +280,39 @@ def test_rolling_common_aggs_over(engine: pl.GPUEngine) -> None: pl.col("x").count().rolling("ts", period="2i").over("g").alias("count"), pl.len().rolling("ts", period="2i").over("g").alias("len"), ) - assert_gpu_result_equal(q, engine=engine) + expected = pl.DataFrame( + { + "sum": [100, 300, 300, 400, 500, 1100], + "min": [100, 100, 300, 400, 500, 500], + "max": [100, 200, 300, 400, 500, 600], + "mean": [100.0, 150.0, 300.0, 400.0, 500.0, 550.0], + "count": pl.Series([1, 2, 1, 1, 1, 2], dtype=pl.UInt32), + "len": pl.Series([1, 2, 1, 1, 1, 2], dtype=pl.UInt32), + } + ) + assert_frame_equal(q.collect(engine=engine), expected) + + +@skip_rolling_expr_136_to_138 +def test_range_rolling_nested_under_range_rolling_over_raises( + engine: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "A"], + "ts": [1, 2, 3], + "x": [10, 20, 30], + } + ) + q = df.select( + pl.col("x") + .sum() + .rolling("ts", period="2i") + .sum() + .rolling("ts", period="2i") + .over("g") + ) + assert_ir_translation_raises(q, engine, NotImplementedError) @pytest.mark.parametrize( From 1b55f52ef51c5f28cbb03e77d1a11678035de23e Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 13 Aug 2026 08:10:13 -0700 Subject: [PATCH 5/7] address more cr comments --- .../tests/expressions/test_rolling.py | 162 ++++++++++++++++++ 1 file changed, 162 insertions(+) diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index 48e91f3346fb..eacca3a3e357 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -293,6 +293,168 @@ def test_rolling_common_aggs_over(engine: pl.GPUEngine) -> None: assert_frame_equal(q.collect(engine=engine), expected) +@skip_rolling_expr_136_to_138 +@pytest.mark.parametrize( + "idx,period", + [ + ([1, 1, 3, 1, 2, 4, 5], "2i"), + ( + [ + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 2), + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 1), + dt.datetime(2025, 1, 1, 9, 3), + dt.datetime(2025, 1, 1, 9, 4), + ], + "2m", + ), + ], + ids=["integer_index", "datetime_index"], +) +def test_rolling_sum_over_index_types_and_group_sizes( + engine: pl.GPUEngine, + idx: list[int] | list[dt.datetime], + period: str, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "B", "B", "C", "C", "C", "C"], + "idx": idx, + "x": [10, 20, 30, 40, 50, 60, 70], + } + ) + q = df.select( + pl.col("x").sum().rolling("idx", period=period).over("g").alias("sum") + ) + expected = pl.DataFrame({"sum": [10, 20, 30, 40, 90, 60, 130]}) + assert_frame_equal(q.collect(engine=engine), expected) + + +@skip_rolling_expr_136_to_138 +@pytest.mark.parametrize( + "lf,expected", + [ + ( + pl.LazyFrame( + { + "g": pl.Series([], dtype=pl.String), + "ts": pl.Series([], dtype=pl.Int64), + "x": pl.Series([], dtype=pl.Int64), + } + ), + pl.DataFrame( + { + "sum": pl.Series([], dtype=pl.Int64), + "min": pl.Series([], dtype=pl.Int64), + "max": pl.Series([], dtype=pl.Int64), + "mean": pl.Series([], dtype=pl.Float64), + "count": pl.Series([], dtype=pl.UInt32), + "len": pl.Series([], dtype=pl.UInt32), + } + ), + ), + ( + pl.LazyFrame( + { + "g": ["A", "A", "B"], + "ts": [1, 2, 1], + "x": pl.Series([None, None, None], dtype=pl.Int64), + } + ), + pl.DataFrame( + { + "sum": [0, 0, 0], + "min": [None, None, None], + "max": [None, None, None], + "mean": [None, None, None], + "count": pl.Series([0, 0, 0], dtype=pl.UInt32), + "len": pl.Series([1, 2, 1], dtype=pl.UInt32), + }, + schema={ + "sum": pl.Int64, + "min": pl.Int64, + "max": pl.Int64, + "mean": pl.Float64, + "count": pl.UInt32, + "len": pl.UInt32, + }, + ), + ), + ( + pl.LazyFrame( + { + "g": ["A", "A", "A", "B", "B"], + "ts": [1, 2, 3, 1, 3], + "x": pl.Series([10, None, 30, None, 50], dtype=pl.Int64), + } + ), + pl.DataFrame( + { + "sum": [10, 10, 30, 0, 50], + "min": [10, 10, 30, None, 50], + "max": [10, 10, 30, None, 50], + "mean": [10.0, 10.0, 30.0, None, 50.0], + "count": pl.Series([1, 1, 1, 0, 1], dtype=pl.UInt32), + "len": pl.Series([1, 2, 2, 1, 1], dtype=pl.UInt32), + }, + schema={ + "sum": pl.Int64, + "min": pl.Int64, + "max": pl.Int64, + "mean": pl.Float64, + "count": pl.UInt32, + "len": pl.UInt32, + }, + ), + ), + ( + pl.LazyFrame( + { + "g": ["A", "B"], + "ts": [1, 1], + "x": pl.Series([10, None], dtype=pl.Int64), + } + ), + pl.DataFrame( + { + "sum": [10, 0], + "min": [10, None], + "max": [10, None], + "mean": [10.0, None], + "count": pl.Series([1, 0], dtype=pl.UInt32), + "len": pl.Series([1, 1], dtype=pl.UInt32), + }, + schema={ + "sum": pl.Int64, + "min": pl.Int64, + "max": pl.Int64, + "mean": pl.Float64, + "count": pl.UInt32, + "len": pl.UInt32, + }, + ), + ), + ], + ids=["empty", "all_null", "mixed_null", "single_row_groups"], +) +def test_rolling_common_aggs_over_edge_cases( + engine: pl.GPUEngine, + lf: pl.LazyFrame, + expected: pl.DataFrame, +) -> None: + q = lf.sort("g", "ts").select( + pl.col("x").sum().rolling("ts", period="2i").over("g").alias("sum"), + pl.col("x").min().rolling("ts", period="2i").over("g").alias("min"), + pl.col("x").max().rolling("ts", period="2i").over("g").alias("max"), + pl.col("x").mean().rolling("ts", period="2i").over("g").alias("mean"), + pl.col("x").count().rolling("ts", period="2i").over("g").alias("count"), + pl.len().rolling("ts", period="2i").over("g").alias("len"), + ) + assert_frame_equal(q.collect(engine=engine), expected) + + @skip_rolling_expr_136_to_138 def test_range_rolling_nested_under_range_rolling_over_raises( engine: pl.GPUEngine, From 4b94a4b2f8ed13667702eb73ac36452c07ac6421 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 13 Aug 2026 17:11:56 -0700 Subject: [PATCH 6/7] another CI fix --- .../tests/streaming/test_rolling.py | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_rolling.py b/python/cudf_polars/tests/streaming/test_rolling.py index 810598ed6a71..6b9a88cf5e49 100644 --- a/python/cudf_polars/tests/streaming/test_rolling.py +++ b/python/cudf_polars/tests/streaming/test_rolling.py @@ -8,6 +8,7 @@ import pytest import polars as pl +from polars.testing import assert_frame_equal from cudf_polars.engine.options import StreamingOptions from cudf_polars.engine.spmd import SPMDEngine @@ -146,7 +147,23 @@ def test_rolling_sum_over(engine): "notional_before", "volume_after", ) - assert_gpu_result_equal(q, engine=engine, check_row_order=True) + expected = pl.DataFrame( + { + "ric": ["A", "A", "A", "B", "B", "B"], + "ts": [ + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 1), + dt.datetime(2025, 1, 1, 9, 3), + dt.datetime(2025, 1, 1, 9, 0), + dt.datetime(2025, 1, 1, 9, 2), + dt.datetime(2025, 1, 1, 9, 3), + ], + "volume_before": [0, 100, 200, 0, 400, 500], + "notional_before": [0.0, 1000.0, 2200.0, 0.0, 8000.0, 10500.0], + "volume_after": [100, 300, 300, 400, 500, 1100], + } + ) + assert_frame_equal(q.collect(engine=engine), expected) @pytest.mark.skipif( @@ -169,7 +186,17 @@ def test_rolling_common_aggs_over(engine): pl.col("x").count().rolling("ts", period="2i").over("g").alias("count"), pl.len().rolling("ts", period="2i").over("g").alias("len"), ) - assert_gpu_result_equal(q, engine=engine, check_row_order=True) + expected = pl.DataFrame( + { + "sum": [100, 300, 300, 400, 500, 1100], + "min": [100, 100, 300, 400, 500, 500], + "max": [100, 200, 300, 400, 500, 600], + "mean": [100.0, 150.0, 300.0, 400.0, 500.0, 550.0], + "count": pl.Series([1, 2, 1, 1, 1, 2], dtype=pl.UInt32), + "len": pl.Series([1, 2, 1, 1, 1, 2], dtype=pl.UInt32), + } + ) + assert_frame_equal(q.collect(engine=engine), expected) @pytest.mark.parametrize("strategy", ["forward", "backward"]) From a6efdd6bf927ced1134ea4eba32f0c28f9f2009a Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 14 Aug 2026 06:00:49 -0700 Subject: [PATCH 7/7] fix code-cov --- .../tests/expressions/test_rolling.py | 75 ++++++++++++++++++- 1 file changed, 73 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index eacca3a3e357..37741278a979 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -11,10 +11,15 @@ import polars as pl from polars.testing import assert_frame_equal +from cudf_polars.containers import DataType +from cudf_polars.dsl import expr +from cudf_polars.dsl.expressions.base import ExecutionContext +from cudf_polars.dsl.utils.aggregations import decompose_single_agg from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, ) +from cudf_polars.typing import Duration from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_139 if TYPE_CHECKING: @@ -43,6 +48,22 @@ def df(): ) +def _range_rolling_sum( + dtype: DataType, orderby: str, child: expr.Expr +) -> expr.RollingWindow: + offset = Duration((0, 0, 0, 0, True, False)) + period = Duration((0, 0, 0, 2, True, False)) + return expr.RollingWindow( + dtype, + dtype.plc_type, + offset, + period, + "right", + orderby, + expr.Agg(dtype, "sum", (), ExecutionContext.WINDOW, child), + ) + + @skip_rolling_expr_136_to_138 @pytest.mark.parametrize("time_unit", ["ns", "us", "ms"]) def test_rolling_datetime(engine: pl.GPUEngine, time_unit): @@ -297,7 +318,7 @@ def test_rolling_common_aggs_over(engine: pl.GPUEngine) -> None: @pytest.mark.parametrize( "idx,period", [ - ([1, 1, 3, 1, 2, 4, 5], "2i"), + (pl.Series("idx", [1, 1, 3, 1, 2, 4, 5], dtype=pl.Int32), "2i"), ( [ dt.datetime(2025, 1, 1, 9, 0), @@ -315,7 +336,7 @@ def test_rolling_common_aggs_over(engine: pl.GPUEngine) -> None: ) def test_rolling_sum_over_index_types_and_group_sizes( engine: pl.GPUEngine, - idx: list[int] | list[dt.datetime], + idx: pl.Series | list[dt.datetime], period: str, ) -> None: df = pl.LazyFrame( @@ -332,6 +353,38 @@ def test_rolling_sum_over_index_types_and_group_sizes( assert_frame_equal(q.collect(engine=engine), expected) +@skip_rolling_expr_136_to_138 +def test_rolling_sum_over_null_index_raises( + engine_raise_on_fail: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "A"], + "idx": pl.Series([1, None, 3], dtype=pl.Int64), + "x": [10, 20, 30], + } + ) + q = df.select(pl.col("x").sum().rolling("idx", period="2i").over("g").alias("sum")) + with pytest.raises( + RuntimeError, match="Index column 'idx' in rolling may not contain nulls" + ): + q.collect(engine=engine_raise_on_fail) + + +def test_rolling_orderby_name_multiple_index_columns_raises() -> None: + dtype = DataType(pl.Int64()) + col = expr.Col(dtype, "x") + named_exprs = [ + expr.NamedExpr("x_sum", _range_rolling_sum(dtype, "t1", col)), + expr.NamedExpr("y_sum", _range_rolling_sum(dtype, "t2", col)), + ] + with pytest.raises( + NotImplementedError, + match=r"rolling\(\.\.\.\)\.over\(\.\.\.\) only supports one rolling index column", + ): + expr.GroupedWindow._rolling_orderby_name(named_exprs) + + @skip_rolling_expr_136_to_138 @pytest.mark.parametrize( "lf,expected", @@ -477,6 +530,24 @@ def test_range_rolling_nested_under_range_rolling_over_raises( assert_ir_translation_raises(q, engine, NotImplementedError) +def test_range_rolling_nested_window_decomposition_raises() -> None: + dtype = DataType(pl.Int64()) + child = expr.Col(dtype, "x") + inner_rolling = _range_rolling_sum(dtype, "ts", child) + outer_rolling = _range_rolling_sum(dtype, "ts", inner_rolling) + + with pytest.raises( + NotImplementedError, + match="Range rolling over a window does not support nested window expressions", + ): + decompose_single_agg( + expr.NamedExpr("out", outer_rolling), + (f"__{i}" for i in range(1)), + is_top=True, + context=ExecutionContext.WINDOW, + ) + + @pytest.mark.parametrize( "expr", [