Support args and kwargs for rolling aggregations (dask#11856)

Author: fjetter

dask/dataframe/dask_expr/_rolling.py

@@ -287,56 +287,56 @@ def _single_agg(self, expr_cls, how_args=(), how_kwargs=None):
         )
 
     @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):

dask/dataframe/dask_expr/tests/test_rolling.py

@@ -14,7 +14,10 @@
 @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