diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index bdcd6c73764..37c5380331c 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -7,10 +7,12 @@ 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 +import polars as pl + import pylibcudf as plc from cudf_polars.containers import Column, DataFrame, DataType @@ -72,8 +74,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 +96,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 +116,20 @@ 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: + 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, + 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 +386,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 +774,98 @@ 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_name = self._rolling_orderby_name(op.named_exprs) + orderby = df.column_map[orderby_name] + if orderby.obj.null_count() != 0: + raise RuntimeError( + f"Index column '{orderby_name}' in rolling may not contain nulls" + ) + if plc.traits.is_integral(orderby.obj.type()): + orderby = orderby.astype(DataType(pl.Int64()), stream=df.stream) + sorted_orderby_obj = plc.copying.gather( + plc.Table([orderby.obj]), + op.order_index, + plc.copying.OutOfBoundsPolicy.NULLIFY, + stream=df.stream, + ).columns()[0] + + 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, 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] = plc.Table([output_col]) + + names = [ne.name for ne in op.named_exprs] + dtypes = [ne.value.dtype for ne in op.named_exprs] + tables = [results_by_name[name] for name in names] + return names, dtypes, tables + def _reorder_to_input( self, row_id: plc.Column, @@ -812,6 +926,7 @@ def _split_named_expr( "cum_sum": [], "shift": [], "fixed_size_rolling": [], + "range_rolling": [], } for ne in self.named_aggs: @@ -832,10 +947,25 @@ 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 + @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 orderby_names.pop() + def _build_window_order_index( self, by_cols: list[Column], @@ -984,6 +1114,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], @@ -1143,6 +1317,14 @@ def do_evaluate( # noqa: D102 plc.Scalar.from_py(1, plc.types.SIZE_TYPE, stream=df.stream), stream=df.stream, ) + if self._order_by_expr is not None: + over_order_by_col, over_ob_desc, over_ob_nulls_last = ( + order_by_col, + self.options[2], + self.options[3], + ) + else: + over_order_by_col, over_ob_desc, over_ob_nulls_last = (None, False, False) if rank_named := unary_window_ops["rank"]: if self._order_by_expr is not None: @@ -1270,121 +1452,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( + broadcasted_cols.extend( + self._apply_ordered_unary_op( + ShiftOp(named_exprs=shift_named), + df, + grouper, 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, + 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, ) ) - 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, - ) + + if rolling_named := unary_window_ops["range_rolling"]: + orderby_name = self._rolling_orderby_name(rolling_named) + rolling_order_by_col = df.column_map[orderby_name] 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, ) ) diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 8316879b0c1..288732182f7 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 " @@ -1216,30 +1220,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.FixedSizeRollingWindow, 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 3b68c05ab50..839f8c6f91b 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) @@ -163,6 +167,19 @@ 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]) + or _contains_range_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 268cca2f1cb..5cc4c3d5329 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -3,16 +3,23 @@ from __future__ import annotations +import datetime as dt from typing import TYPE_CHECKING, Literal, cast import pytest 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: @@ -41,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): @@ -185,6 +208,348 @@ 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", + ) + 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], + } + ) + # Polars <1.36 cannot collect this expression on CPU. Switch to + # assert_gpu_result_equal after we drop Polars 1.35. + 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 +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"), + ) + 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 +@pytest.mark.parametrize( + "idx,period", + [ + (pl.Series("idx", [1, 1, 3, 1, 2, 4, 5], dtype=pl.Int32), "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: pl.Series | 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 +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", + [ + ( + 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, +) -> 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) + + +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", [ diff --git a/python/cudf_polars/tests/streaming/test_rolling.py b/python/cudf_polars/tests/streaming/test_rolling.py index de1682c173e..6b9a88cf5e4 100644 --- a/python/cudf_polars/tests/streaming/test_rolling.py +++ b/python/cudf_polars/tests/streaming/test_rolling.py @@ -3,9 +3,12 @@ from __future__ import annotations +import datetime as dt + 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 @@ -100,6 +103,102 @@ 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", + ) + 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( + 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"), + ) + 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"]) def test_over_cum_sum_fill_null_per_partition(engine, strategy): df = pl.LazyFrame(