PERF: Use Indexers to implement groupby rolling

pandas/core/window/indexers.py

@@ -1,5 +1,5 @@
 """Indexer objects for computing start/end window bounds for rolling operations"""
-from typing import Optional, Tuple
+from typing import Dict, Optional, Tuple
 
 import numpy as np
 
@@ -170,3 +170,49 @@ def get_window_bounds(
         end = np.concatenate([end_s, end_e])
 
         return start, end
+
+
+class GroupbyRollingIndexer(BaseIndexer):
+    """Calculate bounds to compute groupby rolling, mimicking df.groupby().rolling()"""
+
+    def __init__(
+        self,
+        index_array: Optional[np.ndarray] = None,
+        window_size: int = 0,
+        groupby_indicies: Optional[Dict] = None,
+        rolling_indexer=None,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        **kwargs :
+            keyword arguments that will be available when get_window_bounds is called
+        """
+        self.groupby_indicies = groupby_indicies
+        self.rolling_indexer = rolling_indexer
+        super().__init__(index_array, window_size, **kwargs)
+
+    @Appender(get_window_bounds_doc)
+    def get_window_bounds(
+        self,
+        num_values: int = 0,
+        min_periods: Optional[int] = None,
+        center: Optional[bool] = None,
+        closed: Optional[str] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        start_arrays = []
+        end_arrays = []
+        for key, indicies in self.groupby_indicies.items():
+            indexer = self.rolling_indexer(
+                index_array=indicies, window_size=self.window_size
+            )
+            start, end = indexer.get_window_bounds(
+                len(indicies), min_periods, center, closed
+            )
+            indicies = np.append(indicies, [indicies[-1] + 1])  # Check this assumption
+            start_arrays.append(indicies.take(start))
+            end_arrays.append(indicies.take(end))
+        start = np.concatenate(start_arrays)
+        end = np.concatenate(end_arrays)
+        return start, end

pandas/core/window/rolling.py

@@ -37,7 +37,7 @@
 from pandas.core.base import DataError, PandasObject, SelectionMixin, ShallowMixin
 import pandas.core.common as com
 from pandas.core.construction import extract_array
-from pandas.core.indexes.api import Index, ensure_index
+from pandas.core.indexes.api import Index, MultiIndex, ensure_index
 from pandas.core.util.numba_ import NUMBA_FUNC_CACHE
 from pandas.core.window.common import (
     WindowGroupByMixin,
@@ -51,6 +51,7 @@
 )
 from pandas.core.window.indexers import (
     BaseIndexer,
+    GroupbyRollingIndexer,
     FixedWindowIndexer,
     VariableWindowIndexer,
 )
@@ -2086,12 +2087,64 @@ class RollingGroupby(WindowGroupByMixin, Rolling):
     Provide a rolling groupby implementation.
     """
 
+    def _apply(
+        self,
+        func: Callable,
+        center: bool,
+        require_min_periods: int = 0,
+        floor: int = 1,
+        is_weighted: bool = False,
+        name: Optional[str] = None,
+        use_numba_cache: bool = False,
+        **kwargs,
+    ):
+        result = Rolling._apply(
+            self,
+            func,
+            center,
+            require_min_periods,
+            floor,
+            is_weighted,
+            name,
+            use_numba_cache,
+            **kwargs,
+        )
+        # _wrap_outputs does not know about what the result index should be
+        index_name = [self._selected_obj.index.name]
+        groupby_names = [grouping.name for grouping in self._groupby.grouper._groupings]
+        index_names = groupby_names + index_name
+        index_data = []
+        for key, values in self._groupby.grouper.indices.items():
+            for value in values:
+                if not is_list_like(key):
+                    data = (key, value)
+                else:
+                    data = (*key, value)
+                index_data.append(data)
+        result_index = MultiIndex.from_tuples(index_data, names=index_names)
+        result.index = result_index
+        return result
+
     @property
     def _constructor(self):
         return Rolling
 
-    def _gotitem(self, key, ndim, subset=None):
+    def _get_window_indexer(self, window: int) -> BaseIndexer:
+        """
+        Return an indexer class that will compute the window start and end bounds
+        """
+        if self.is_freq_type:
+            rolling_indexer = VariableWindowIndexer
+        else:
+            rolling_indexer = FixedWindowIndexer
+        window_indexer = GroupbyRollingIndexer(
+            window_size=window,
+            groupby_indicies=self._groupby.indices,
+            rolling_indexer=rolling_indexer,
+        )
+        return window_indexer
 
+    def _gotitem(self, key, ndim, subset=None):
         # we are setting the index on the actual object
         # here so our index is carried thru to the selected obj
         # when we do the splitting for the groupby