PERF: Rolling/Expanding.cov/corr (#39591)

Author: mroeschke

asv_bench/benchmarks/rolling.py

@@ -140,8 +140,11 @@ class Pairwise:
 
     def setup(self, window, method, pairwise):
         N = 10 ** 4
+        n_groups = 20
+        groups = [i for _ in range(N // n_groups) for i in range(n_groups)]
         arr = np.random.random(N)
         self.df = pd.DataFrame(arr)
+        self.df_group = pd.DataFrame({"A": groups, "B": arr}).groupby("A")
 
     def time_pairwise(self, window, method, pairwise):
         if window is None:
@@ -150,6 +153,13 @@ def time_pairwise(self, window, method, pairwise):
             r = self.df.rolling(window=window)
         getattr(r, method)(self.df, pairwise=pairwise)
 
+    def time_groupby(self, window, method, pairwise):
+        if window is None:
+            r = self.df_group.expanding()
+        else:
+            r = self.df_group.rolling(window=window)
+        getattr(r, method)(self.df, pairwise=pairwise)
+
 
 class Quantile:
     params = (

doc/source/whatsnew/v1.3.0.rst

@@ -249,7 +249,8 @@ Performance improvements
 - Performance improvement in :meth:`Series.mean` for nullable data types (:issue:`34814`)
 - Performance improvement in :meth:`Series.isin` for nullable data types (:issue:`38340`)
 - Performance improvement in :meth:`DataFrame.corr` for method=kendall (:issue:`28329`)
-- Performance improvement in :meth:`core.window.Rolling.corr` and :meth:`core.window.Rolling.cov` (:issue:`39388`)
+- Performance improvement in :meth:`core.window.rolling.Rolling.corr` and :meth:`core.window.rolling.Rolling.cov` (:issue:`39388`)
+- Performance improvement in :meth:`core.window.rolling.RollingGroupby.corr`, :meth:`core.window.expanding.ExpandingGroupby.corr`, :meth:`core.window.expanding.ExpandingGroupby.corr` and :meth:`core.window.expanding.ExpandingGroupby.cov` (:issue:`39591`)
 
 .. ---------------------------------------------------------------------------
 
@@ -407,6 +408,8 @@ Groupby/resample/rolling
 - Bug in :meth:`.Resampler.aggregate` and :meth:`DataFrame.transform` raising ``TypeError`` instead of ``SpecificationError`` when missing keys had mixed dtypes (:issue:`39025`)
 - Bug in :meth:`.DataFrameGroupBy.idxmin` and :meth:`.DataFrameGroupBy.idxmax` with ``ExtensionDtype`` columns (:issue:`38733`)
 - Bug in :meth:`Series.resample` would raise when the index was a :class:`PeriodIndex` consisting of ``NaT`` (:issue:`39227`)
+- Bug in :meth:`core.window.rolling.RollingGroupby.corr` and :meth:`core.window.expanding.ExpandingGroupby.corr` where the groupby column would return 0 instead of ``np.nan`` when providing ``other`` that was longer than each group (:issue:`39591`)
+- Bug in :meth:`core.window.expanding.ExpandingGroupby.corr` and :meth:`core.window.expanding.ExpandingGroupby.cov` where 1 would be returned instead of ``np.nan`` when providing ``other`` that was longer than each group (:issue:`39591`)
 
 Reshaping
 ^^^^^^^^^

pandas/core/window/ewm.py

@@ -3,14 +3,14 @@
 import datetime
 from functools import partial
 from textwrap import dedent
-from typing import TYPE_CHECKING, Optional, Union
+from typing import Optional, Union
 import warnings
 
 import numpy as np
 
 from pandas._libs.tslibs import Timedelta
 import pandas._libs.window.aggregations as window_aggregations
-from pandas._typing import FrameOrSeries, TimedeltaConvertibleTypes
+from pandas._typing import FrameOrSeries, FrameOrSeriesUnion, TimedeltaConvertibleTypes
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import doc
 
@@ -19,7 +19,7 @@
 
 import pandas.core.common as common
 from pandas.core.util.numba_ import maybe_use_numba
-from pandas.core.window.common import flex_binary_moment, zsqrt
+from pandas.core.window.common import zsqrt
 from pandas.core.window.doc import (
     _shared_docs,
     args_compat,
@@ -35,10 +35,7 @@
     GroupbyIndexer,
 )
 from pandas.core.window.numba_ import generate_numba_groupby_ewma_func
-from pandas.core.window.rolling import BaseWindow, BaseWindowGroupby, dispatch
-
-if TYPE_CHECKING:
-    from pandas import Series
+from pandas.core.window.rolling import BaseWindow, BaseWindowGroupby
 
 
 def get_center_of_mass(
@@ -74,13 +71,20 @@ def get_center_of_mass(
     return float(comass)
 
 
-def wrap_result(obj: Series, result: np.ndarray) -> Series:
+def dispatch(name: str, *args, **kwargs):
     """
-    Wrap a single 1D result.
+    Dispatch to groupby apply.
     """
-    obj = obj._selected_obj
 
-    return obj._constructor(result, obj.index, name=obj.name)
+    def outer(self, *args, **kwargs):
+        def f(x):
+            x = self._shallow_copy(x, groupby=self._groupby)
+            return getattr(x, name)(*args, **kwargs)
+
+        return self._groupby.apply(f)
+
+    outer.__name__ = name
+    return outer
 
 
 class ExponentialMovingWindow(BaseWindow):
@@ -443,36 +447,30 @@ def var_func(values, begin, end, min_periods):
     )
     def cov(
         self,
-        other: Optional[Union[np.ndarray, FrameOrSeries]] = None,
+        other: Optional[FrameOrSeriesUnion] = None,
         pairwise: Optional[bool] = None,
         bias: bool = False,
         **kwargs,
     ):
-        if other is None:
-            other = self._selected_obj
-            # only default unset
-            pairwise = True if pairwise is None else pairwise
-        other = self._shallow_copy(other)
-
-        def _get_cov(X, Y):
-            X = self._shallow_copy(X)
-            Y = self._shallow_copy(Y)
-            cov = window_aggregations.ewmcov(
-                X._prep_values(),
+        from pandas import Series
+
+        def cov_func(x, y):
+            x_array = self._prep_values(x)
+            y_array = self._prep_values(y)
+            result = window_aggregations.ewmcov(
+                x_array,
                 np.array([0], dtype=np.int64),
                 np.array([0], dtype=np.int64),
                 self.min_periods,
-                Y._prep_values(),
+                y_array,
                 self.com,
                 self.adjust,
                 self.ignore_na,
                 bias,
             )
-            return wrap_result(X, cov)
+            return Series(result, index=x.index, name=x.name)
 
-        return flex_binary_moment(
-            self._selected_obj, other._selected_obj, _get_cov, pairwise=bool(pairwise)
-        )
+        return self._apply_pairwise(self._selected_obj, other, pairwise, cov_func)
 
     @doc(
         template_header,
@@ -502,45 +500,37 @@ def _get_cov(X, Y):
     )
     def corr(
         self,
-        other: Optional[Union[np.ndarray, FrameOrSeries]] = None,
+        other: Optional[FrameOrSeriesUnion] = None,
         pairwise: Optional[bool] = None,
         **kwargs,
     ):
-        if other is None:
-            other = self._selected_obj
-            # only default unset
-            pairwise = True if pairwise is None else pairwise
-        other = self._shallow_copy(other)
+        from pandas import Series
 
-        def _get_corr(X, Y):
-            X = self._shallow_copy(X)
-            Y = self._shallow_copy(Y)
+        def cov_func(x, y):
+            x_array = self._prep_values(x)
+            y_array = self._prep_values(y)
 
-            def _cov(x, y):
+            def _cov(X, Y):
                 return window_aggregations.ewmcov(
-                    x,
+                    X,
                     np.array([0], dtype=np.int64),
                     np.array([0], dtype=np.int64),
                     self.min_periods,
-                    y,
+                    Y,
                     self.com,
                     self.adjust,
                     self.ignore_na,
                     1,
                 )
 
-            x_values = X._prep_values()
-            y_values = Y._prep_values()
             with np.errstate(all="ignore"):
-                cov = _cov(x_values, y_values)
-                x_var = _cov(x_values, x_values)
-                y_var = _cov(y_values, y_values)
-                corr = cov / zsqrt(x_var * y_var)
-            return wrap_result(X, corr)
-
-        return flex_binary_moment(
-            self._selected_obj, other._selected_obj, _get_corr, pairwise=bool(pairwise)
-        )
+                cov = _cov(x_array, y_array)
+                x_var = _cov(x_array, x_array)
+                y_var = _cov(y_array, y_array)
+                result = cov / zsqrt(x_var * y_var)
+            return Series(result, index=x.index, name=x.name)
+
+        return self._apply_pairwise(self._selected_obj, other, pairwise, cov_func)
 
 
 class ExponentialMovingWindowGroupby(BaseWindowGroupby, ExponentialMovingWindow):