Leverage new pylibcudf grouped_range_rolling_window for cuDF classic rolling(window: int)

ci/test_narwhals.sh

@@ -25,7 +25,22 @@ rapids-pip-retry install -U -e .
 rapids-logger "Check narwhals versions"
 python -c "import narwhals; print(narwhals.show_versions())"
 
-TESTS_TO_ALWAYS_SKIP_CUDF="test_date_lit[cudf]" # pa.date32() is not supported in cudf
+# test_rolling_mean_expr_lazy_grouped: xpassing in Narwhals
+# test_rolling_std_expr_lazy_grouped: xpassing in Narwhals
+# test_rolling_sum_expr_lazy_grouped: xpassing in Narwhals
+# test_rolling_var_expr_lazy_grouped: xpassing in Narwhals
+TESTS_THAT_NEED_NARWHALS_FIX_FOR_CUDF="not test_rolling_mean_expr_lazy_grouped[cudf-expected_a4-3-1-True] \
+and not test_rolling_mean_expr_lazy_grouped[cudf-expected_a5-4-1-True] \
+and not test_rolling_mean_expr_lazy_grouped[cudf-expected_a6-5-1-True] \
+and not test_rolling_std_expr_lazy_grouped[cudf-expected_a4-3-1-True-1] \
+and not test_rolling_std_expr_lazy_grouped[cudf-expected_a5-4-1-True-1] \
+and not test_rolling_std_expr_lazy_grouped[cudf-expected_a6-5-1-True-0] \
+and not test_rolling_sum_expr_lazy_grouped[cudf-expected_a4-3-1-True] \
+and not test_rolling_sum_expr_lazy_grouped[cudf-expected_a5-4-1-True] \
+and not test_rolling_sum_expr_lazy_grouped[cudf-expected_a6-5-1-True] \
+and not test_rolling_var_expr_lazy_grouped[cudf-expected_a4-3-1-True-1] \
+and not test_rolling_var_expr_lazy_grouped[cudf-expected_a5-4-1-True-1] \
+and not test_rolling_var_expr_lazy_grouped[cudf-expected_a6-5-1-True-0]"
 
 rapids-logger "Run narwhals tests for cuDF"
 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest \
@@ -35,9 +50,7 @@ PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest \
     -p env \
     -p no:pytest_benchmark \
     -p cudf.testing.narwhals_test_plugin \
-    -k "not ( \
-        ${TESTS_TO_ALWAYS_SKIP_CUDF} \
-    )" \
+    -k "$TESTS_THAT_NEED_NARWHALS_FIX_FOR_CUDF" \
     --numprocesses=8 \
     --dist=worksteal \
     --constructors=cudf

python/cudf/cudf/core/window/rolling.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2020-2025, NVIDIA CORPORATION
 from __future__ import annotations
 
+import itertools
 import warnings
 from typing import TYPE_CHECKING
 
@@ -23,6 +24,7 @@
 
 if TYPE_CHECKING:
     from cudf.core.dataframe import DataFrame
+    from cudf.core.index import Index
     from cudf.core.series import Series
 
 
@@ -32,6 +34,7 @@ class _RollingBase:
     """
 
     obj: DataFrame | Series
+    _group_keys: Index | None = None
 
     def _apply_agg_column(
         self, source_column: ColumnBase, agg_name: str
@@ -194,6 +197,7 @@ class Rolling(GetAttrGetItemMixin, _RollingBase, Reducible):
     _PROTECTED_KEYS = frozenset(("obj",))
 
     _time_window = False
+    _group_keys = None
 
     _VALID_REDUCTIONS = {
         "sum",
@@ -259,69 +263,91 @@ def __getitem__(self, arg):
             center=self.center,
         )
 
-    def _apply_agg_column(self, source_column, agg_name):
+    def _apply_agg_column(self, source_column, agg_name) -> ColumnBase:
         min_periods = self.min_periods or 1
         if isinstance(self.window, int):
-            preceding_window = None
-            following_window = None
-            window = self.window
-        elif isinstance(self.window, BaseIndexer):
-            start, end = self.window.get_window_bounds(
-                num_values=len(self.obj),
-                min_periods=self.min_periods,
-                center=self.center,
-                closed=None,
-                step=None,
-            )
-            start = as_column(start, dtype=SIZE_TYPE_DTYPE)
-            end = as_column(end, dtype=SIZE_TYPE_DTYPE)
-
-            idx = as_column(range(len(start)))
-            preceding_window = (idx - start + np.int32(1)).astype(
-                SIZE_TYPE_DTYPE
-            )
-            following_window = (end - idx - np.int32(1)).astype(
-                SIZE_TYPE_DTYPE
+            if self.center:
+                pre = (self.window // 2) + 1
+                fwd = self.window - (pre)
+            else:
+                pre = self.window
+                fwd = 0
+            rolling_request = plc.rolling.RollingRequest(
+                source_column.to_pylibcudf(mode="read"),
+                min_periods,
+                aggregation.make_aggregation(
+                    agg_name,
+                    {"dtype": source_column.dtype}
+                    if callable(agg_name)
+                    else self.agg_params,
+                ).plc_obj,
             )
-            window = None
+            orderby_obj = as_column(range(len(source_column)))
+            if self._group_keys is not None:
+                group_cols: list[plc.Column] = [
+                    col.to_pylibcudf(mode="read")
+                    for col in self._group_keys._columns
+                ]
+            else:
+                group_cols = []
+            group_keys = plc.Table(group_cols)
+            with acquire_spill_lock():
+                (plc_result,) = plc.rolling.grouped_range_rolling_window(
+                    group_keys,
+                    orderby_obj.to_pylibcudf(mode="read"),
+                    plc.types.Order.ASCENDING,
+                    plc.types.NullOrder.BEFORE,
+                    plc.rolling.BoundedOpen(plc.Scalar.from_py(pre)),
+                    plc.rolling.BoundedClosed(plc.Scalar.from_py(fwd)),
+                    [rolling_request],
+                ).columns()
+            return ColumnBase.from_pylibcudf(plc_result)
         else:
-            preceding_window = as_column(self.window)
-            following_window = as_column(
-                0, length=self.window.size, dtype=self.window.dtype
-            )
-            window = None
+            if isinstance(self.window, BaseIndexer):
+                start, end = self.window.get_window_bounds(
+                    num_values=len(self.obj),
+                    min_periods=self.min_periods,
+                    center=self.center,
+                    closed=None,
+                    step=None,
+                )
+                start = as_column(start, dtype=SIZE_TYPE_DTYPE)
+                end = as_column(end, dtype=SIZE_TYPE_DTYPE)
 
-        with acquire_spill_lock():
-            if window is None:
+                idx = as_column(range(len(start)))
+                preceding_window = (idx - start + np.int32(1)).astype(
+                    SIZE_TYPE_DTYPE
+                )
+                following_window = (end - idx - np.int32(1)).astype(
+                    SIZE_TYPE_DTYPE
+                )
+            else:
                 if self.center:
                     # TODO: we can support this even though Pandas currently does not
                     raise NotImplementedError(
                         "center is not implemented for offset-based windows"
                     )
-                pre = preceding_window.to_pylibcudf(mode="read")
-                fwd = following_window.to_pylibcudf(mode="read")
-            else:
-                if self.center:
-                    pre = (window // 2) + 1
-                    fwd = window - (pre)
-                else:
-                    pre = window
-                    fwd = 0
-
-            return ColumnBase.from_pylibcudf(
-                plc.rolling.rolling_window(
-                    source_column.to_pylibcudf(mode="read"),
-                    pre,
-                    fwd,
-                    min_periods,
-                    aggregation.make_aggregation(
-                        agg_name,
-                        {"dtype": source_column.dtype}
-                        if callable(agg_name)
-                        else self.agg_params,
-                    ).plc_obj,
+                preceding_window = as_column(self.window)
+                following_window = as_column(
+                    0, length=self.window.size, dtype=self.window.dtype
+                )
+            pre = preceding_window.to_pylibcudf(mode="read")
+            fwd = following_window.to_pylibcudf(mode="read")
+            with acquire_spill_lock():
+                return ColumnBase.from_pylibcudf(
+                    plc.rolling.rolling_window(
+                        source_column.to_pylibcudf(mode="read"),
+                        pre,
+                        fwd,
+                        min_periods,
+                        aggregation.make_aggregation(
+                            agg_name,
+                            {"dtype": source_column.dtype}
+                            if callable(agg_name)
+                            else self.agg_params,
+                        ).plc_obj,
+                    )
                 )
-            )
 
     def _reduce(
         self,
@@ -560,31 +586,39 @@ def __init__(self, groupby, window, min_periods=None, center=False):
             sort_order
         )
 
-        gb_size = groupby.size().sort_index()
-        self._group_starts = (
-            gb_size.cumsum().shift(1).fillna(0).repeat(gb_size)
-        )
+        if not is_integer(window):
+            gb_size = groupby.size().sort_index()
+            self._group_starts = (
+                gb_size.cumsum().shift(1).fillna(0).repeat(gb_size)
+            )
 
         super().__init__(obj, window, min_periods=min_periods, center=center)
 
-    @acquire_spill_lock()
     def _window_to_window_sizes(self, window):
         if is_integer(window):
-            return cudautils.grouped_window_sizes_from_offset(
-                as_column(range(len(self.obj))).data_array_view(mode="read"),
-                self._group_starts,
-                window,
-            )
+            return super()._window_to_window_sizes(window)
         else:
-            return cudautils.grouped_window_sizes_from_offset(
-                self.obj.index._column.data_array_view(mode="read"),
-                self._group_starts,
-                window,
-            )
+            with acquire_spill_lock():
+                return cudautils.grouped_window_sizes_from_offset(
+                    self.obj.index._column.data_array_view(mode="read"),
+                    self._group_starts,
+                    window,
+                )
 
     def _apply_agg(self, agg_name):
         index = MultiIndex._from_data(
-            {**self._group_keys._data, **self.obj.index._data}
+            dict(
+                enumerate(
+                    itertools.chain(
+                        self._group_keys._columns, self.obj.index._columns
+                    )
+                )
+            )
+        )
+        index.names = list(
+            itertools.chain(
+                self._group_keys._column_names, self.obj.index._column_names
+            )
         )
         result = super()._apply_agg(agg_name)
         result.index = index

python/cudf/cudf/pandas/scripts/conftest-patch.py

@@ -16874,21 +16874,9 @@ def pytest_unconfigure(config):
     "tests/window/test_groupby.py::TestRolling::test_datelike_on_monotonic_within_each_group",
     "tests/window/test_groupby.py::TestRolling::test_getitem",
     "tests/window/test_groupby.py::TestRolling::test_getitem_multiple",
-    "tests/window/test_groupby.py::TestRolling::test_groupby_level",
     "tests/window/test_groupby.py::TestRolling::test_groupby_monotonic",
     "tests/window/test_groupby.py::TestRolling::test_groupby_rolling_custom_indexer",
-    "tests/window/test_groupby.py::TestRolling::test_groupby_rolling_group_keys[False]",
-    "tests/window/test_groupby.py::TestRolling::test_groupby_rolling_group_keys[True]",
     "tests/window/test_groupby.py::TestRolling::test_groupby_rolling_no_sort",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[int16]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[int32]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[int64]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[int8]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[int]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[uint16]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[uint32]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[uint64]",
-    "tests/window/test_groupby.py::TestRolling::test_nan_and_zero_endpoints[uint8]",
     "tests/window/test_groupby.py::test_rolling_corr_with_single_integer_in_index",
     "tests/window/test_groupby.py::test_rolling_corr_with_tuples_in_index",
     "tests/window/test_pairwise.py::TestPairwise::test_no_flex[pairwise_frames1-<lambda>1]",

python/cudf/cudf/testing/narwhals_test_plugin.py

@@ -12,6 +12,7 @@
 
 EXPECTED_FAILURES: Mapping[str, str] = {
     "tests/frame/select_test.py::test_select_duplicates[cudf]": "cuDF doesn't support having multiple columns with same names",
+    "tests/expr_and_series/lit_test.py::test_date_lit[cudf]": "cuDF does not support pa.date32()",
 }