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):