Repository: dask/dask
Pull request #11856 (2025-04-01 10:09:03)

Support args and kwargs for rolling aggregations
Closes https://github.com/dask/dask/issues/11851

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ±    0        9 suites  ±0   3h 25m 40s ⏱️ - 3m 20s
 17 999 tests +  180   16 784 ✅ +  180   1 215 💤 ±0  0 ❌ ±0 
161 065 runs  +1 620  148 952 ✅ +1 622  12 113 💤  - 2  0 ❌ ±0 

Results for commit 37372969. ± Comparison against base commit 4399aef0.

<details>
  <summary>This pull request <b>removes</b> 132 and <b>adds</b> 312 tests. <i>Note that renamed tests count towards both.</i></summary>

```
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-count-how_args0]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-kurt-how_args10]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-max-how_args4]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-mean-how_args1]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-median-how_args7]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-min-how_args3]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-quantile-how_args9]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-skew-how_args8]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-std-how_args6]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-sum-how_args2]
…
```

```
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-count-how_args0]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-kurt-how_args12]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-max-how_args4]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-mean-how_args1]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-median-how_args9]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-min-how_args3]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-quantile-how_args11]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-skew-how_args10]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-std-how_args7]
dask.dataframe.dask_expr.tests.test_rolling ‑ test_rolling_apis[1-False-1-None-False-std-how_args8]
…
```
</details>

[test-results]:data:application/gzip;base64,H4sIAOHE62cC/02Myw6DIBBFf8Ww7oJheDj9GUMQElIfDcKq6b8XjVKX59zc82EhTn5jz44eHdtKzA3GkmyO61IRBEpeVR3zPoMhoouHrTi3S216+Zev+D6eoJoLNk7VtdLgU1rTaVJZjrQGrtXJLS17UqLJlga83C198L3s1nmOuQJDg0aQJq8Cgh6lJ+CjCAJrPkhDCNZb5QL7/gDOtk3qGAEAAA==


Patch:
diff --git a/dask/dataframe/dask_expr/_rolling.py b/dask/dataframe/dask_expr/_rolling.py
index fd72ba447..2bf4f9cf5 100644
--- a/dask/dataframe/dask_expr/_rolling.py
+++ b/dask/dataframe/dask_expr/_rolling.py
@@ -287,56 +287,56 @@ class Rolling:
         )
 
     @derived_from(pd_Rolling)
-    def cov(self):
-        return self._single_agg(RollingCov)
+    def cov(self, *args, **kwargs):
+        return self._single_agg(RollingCov, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
     def apply(self, func, *args, **kwargs):
         return self._single_agg(RollingApply, how_args=(func, *args), how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def count(self):
-        return self._single_agg(RollingCount)
+    def count(self, *args, **kwargs):
+        return self._single_agg(RollingCount, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def sum(self):
-        return self._single_agg(RollingSum)
+    def sum(self, *args, **kwargs):
+        return self._single_agg(RollingSum, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def mean(self):
-        return self._single_agg(RollingMean)
+    def mean(self, *args, **kwargs):
+        return self._single_agg(RollingMean, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def min(self):
-        return self._single_agg(RollingMin)
+    def min(self, *args, **kwargs):
+        return self._single_agg(RollingMin, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def max(self):
-        return self._single_agg(RollingMax)
+    def max(self, *args, **kwargs):
+        return self._single_agg(RollingMax, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def var(self):
-        return self._single_agg(RollingVar)
+    def var(self, *args, **kwargs):
+        return self._single_agg(RollingVar, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def std(self):
-        return self._single_agg(RollingStd)
+    def std(self, *args, **kwargs):
+        return self._single_agg(RollingStd, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def median(self):
-        return self._single_agg(RollingMedian)
+    def median(self, *args, **kwargs):
+        return self._single_agg(RollingMedian, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def quantile(self, q):
-        return self._single_agg(RollingQuantile, how_args=(q,))
+    def quantile(self, q, *args, **kwargs):
+        return self._single_agg(RollingQuantile, how_args=(q, *args), how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def skew(self):
-        return self._single_agg(RollingSkew)
+    def skew(self, *args, **kwargs):
+        return self._single_agg(RollingSkew, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def kurt(self):
-        return self._single_agg(RollingKurt)
+    def kurt(self, *args, **kwargs):
+        return self._single_agg(RollingKurt, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
     def agg(self, func, *args, **kwargs):


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_rolling.py b/dask/dataframe/dask_expr/tests/test_rolling.py
index 35026d0d9..a1384ce61 100644
--- a/dask/dataframe/dask_expr/tests/test_rolling.py
+++ b/dask/dataframe/dask_expr/tests/test_rolling.py
@@ -14,7 +14,10 @@ pd = _backend_library()
 @pytest.fixture
 def pdf():
     idx = pd.date_range("2000-01-01", periods=12, freq="min")
-    pdf = pd.DataFrame({"foo": range(len(idx))}, index=idx)
+    pdf = pd.DataFrame(
+        {"foo": range(len(idx)), "bar": idx},
+        index=idx,
+    )
     pdf["bar"] = 1
     yield pdf
 
@@ -34,30 +37,44 @@ def df(pdf, request):
         ("min", ()),
         ("max", ()),
         ("var", ()),
+        ("var", (2,)),  # ddof
         ("std", ()),
+        ("std", (2,)),  # ddof
         ("median", ()),
         ("skew", ()),
         ("quantile", (0.5,)),
         ("kurt", ()),
     ],
 )
+@pytest.mark.parametrize("numeric_only", [True, False])
 @pytest.mark.parametrize("window,min_periods", ((1, None), (3, 2), (3, 3)))
 @pytest.mark.parametrize("center", (True, False))
 @pytest.mark.parametrize("df", (1, 2), indirect=True)
-def test_rolling_apis(df, pdf, window, api, how_args, min_periods, center):
+def test_rolling_apis(
+    df, pdf, window, api, how_args, min_periods, center, numeric_only
+):
     args = (window,)
-    kwargs = dict(min_periods=min_periods, center=center)
-
-    result = getattr(df.rolling(*args, **kwargs), api)(*how_args)
-    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args)
+    kwargs = dict(
+        min_periods=min_periods,
+        center=center,
+    )
+    how_kwargs = dict(
+        numeric_only=numeric_only,
+    )
+    result = getattr(df.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)
+    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)
     assert_eq(result, expected)
 
-    result = getattr(df.rolling(*args, **kwargs), api)(*how_args)["foo"]
-    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args)["foo"]
+    result = getattr(df.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)["foo"]
+    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)[
+        "foo"
+    ]
     assert_eq(result, expected)
 
     q = result.simplify()
-    eq = getattr(df["foo"].rolling(*args, **kwargs), api)(*how_args).simplify()
+    eq = getattr(df["foo"].rolling(*args, **kwargs), api)(
+        *how_args, **how_kwargs
+    ).simplify()
     assert q._name == eq._name