ENH: Add method argument to rolling constructor to allow table-wise rolling

doc/source/user_guide/window.rst

@@ -37,14 +37,14 @@ pandas supports 4 types of windowing operations:
 #. Expanding window: Accumulating window over the values.
 #. Exponentially Weighted window: Accumulating and exponentially weighted window over the values.
 
-=============================   =================  ===========================   ===========================  ========================
-Concept                         Method             Returned Object               Supports time-based windows  Supports chained groupby
-=============================   =================  ===========================   ===========================  ========================
-Rolling window                  ``rolling``        ``Rolling``                   Yes                          Yes
-Weighted window                 ``rolling``        ``Window``                    No                           No
-Expanding window                ``expanding``      ``Expanding``                 No                           Yes
-Exponentially Weighted window   ``ewm``            ``ExponentialMovingWindow``   No                           Yes (as of version 1.2)
-=============================   =================  ===========================   ===========================  ========================
+=============================   =================  ===========================   ===========================  ========================  ===================================
+Concept                         Method             Returned Object               Supports time-based windows  Supports chained groupby  Supports table method (version 1.3)
+=============================   =================  ===========================   ===========================  ========================  ===================================
+Rolling window                  ``rolling``        ``Rolling``                   Yes                          Yes                       Yes
+Weighted window                 ``rolling``        ``Window``                    No                           No                        No
+Expanding window                ``expanding``      ``Expanding``                 No                           Yes                       Yes
+Exponentially Weighted window   ``ewm``            ``ExponentialMovingWindow``   No                           Yes (as of version 1.2)   No
+=============================   =================  ===========================   ===========================  ========================  ===================================
 
 As noted above, some operations support specifying a window based on a time offset:
 
@@ -76,6 +76,27 @@ which will first group the data by the specified keys and then perform a windowi
     to compute the rolling sums to preserve accuracy as much as possible.
 
 
+.. versionadded:: 1.3
+
+Some windowing operations also support the ``method='table'`` option in the constructor which
+performs the windowing operaion over an entire :class:`DataFrame` instead of a single column or row at a time.
+This can provide a useful performance benefit for :class:`DataFrame`s with a many columns or rows
+(with the corresponding ``axis`` argument) or the ability to utilize other columns during the windowing
+operation. The ``method='table'`` option can only be used if ``engine='numba'`` is specified
+in the corresponding method call.
+
+
+.. ipython:: python
+
+   def weighted_mean(x):
+       arr = np.ones((1, x.shape[1]))
+       arr[:, :2] = (x[:, :2] * x[:, 2]).sum(axis=0) / x[:, 2].sum()
+       return arr
+
+   df = pd.DataFrame([[1, 2, 0.6], [2, 3, 0.4], [3, 4, 0.2], [4, 5, 0.7]])
+   df.rolling(2, method="table", min_periods=0).apply(weighted_mean, raw=True, engine="numba")
+
+
 All windowing operations support a ``min_periods`` argument that dictates the minimum amount of
 non-``np.nan`` values a window must have; otherwise, the resulting value is ``np.nan``.
 ``min_peridos`` defaults to 1 for time-based windows and ``window`` for fixed windows

doc/source/whatsnew/v1.3.0.rst

@@ -13,6 +13,11 @@ including other versions of pandas.
 Enhancements
 ~~~~~~~~~~~~
 
+.. _whatsnew_130.window_method_table:
+
+:class:`Rolling` and :class:`Expanding` now support a ``method`` argument with a
+``'table'`` option that performs the windowing operation over an entire :class:`DataFrame`.
+See ref:`window.overview` for performance and functional benefits. (:issue:`15095`)
 
 .. _whatsnew_130.enhancements.other:
 

pandas/core/generic.py

@@ -11210,6 +11210,7 @@ def rolling(
         on: Optional[str] = None,
         axis: Axis = 0,
         closed: Optional[str] = None,
+        method: str = "single",
     ):
         axis = self._get_axis_number(axis)
 
@@ -11223,6 +11224,7 @@ def rolling(
                 on=on,
                 axis=axis,
                 closed=closed,
+                method=method,
             )
 
         return Rolling(
@@ -11234,12 +11236,17 @@ def rolling(
             on=on,
             axis=axis,
             closed=closed,
+            method=method,
         )
 
     @final
     @doc(Expanding)
     def expanding(
-        self, min_periods: int = 1, center: Optional[bool_t] = None, axis: Axis = 0
+        self,
+        min_periods: int = 1,
+        center: Optional[bool_t] = None,
+        axis: Axis = 0,
+        method: str = "single",
     ) -> Expanding:
         axis = self._get_axis_number(axis)
         if center is not None:
@@ -11251,7 +11258,9 @@ def expanding(
         else:
             center = False
 
-        return Expanding(self, min_periods=min_periods, center=center, axis=axis)
+        return Expanding(
+            self, min_periods=min_periods, center=center, axis=axis, method=method
+        )
 
     @final
     @doc(ExponentialMovingWindow)

pandas/core/window/ewm.py

@@ -241,6 +241,7 @@ def __init__(
         self.on = None
         self.center = False
         self.closed = None
+        self.method = "single"
         if times is not None:
             if isinstance(times, str):
                 times = self._selected_obj[times]

pandas/core/window/expanding.py

@@ -24,6 +24,14 @@ class Expanding(RollingAndExpandingMixin):
     center : bool, default False
         Set the labels at the center of the window.
     axis : int or str, default 0
+    method : str, default 'single'
+        Execute the expanding operation per single column or row (`'single'`)
+        or over the entire object (`'table'`).
+
+        This argument is only implemented when specifying ``engine='numba'``
+        in the method call.
+
+        .. versionadded:: 1.3.0
 
     Returns
     -------
@@ -59,10 +67,14 @@ class Expanding(RollingAndExpandingMixin):
     4  7.0
     """
 
-    _attributes = ["min_periods", "center", "axis"]
+    _attributes = ["min_periods", "center", "axis", "method"]
 
-    def __init__(self, obj, min_periods=1, center=None, axis=0, **kwargs):
-        super().__init__(obj=obj, min_periods=min_periods, center=center, axis=axis)
+    def __init__(
+        self, obj, min_periods=1, center=None, axis=0, method="single", **kwargs
+    ):
+        super().__init__(
+            obj=obj, min_periods=min_periods, center=center, axis=axis, method=method
+        )
 
     @property
     def _constructor(self):

pandas/core/window/numba_.py

@@ -17,6 +17,7 @@ def generate_numba_apply_func(
     kwargs: Dict[str, Any],
     func: Callable[..., Scalar],
     engine_kwargs: Optional[Dict[str, bool]],
+    name: str,
 ):
     """
     Generate a numba jitted apply function specified by values from engine_kwargs.
@@ -37,14 +38,16 @@ def generate_numba_apply_func(
         function to be applied to each window and will be JITed
     engine_kwargs : dict
         dictionary of arguments to be passed into numba.jit
+    name: str
+        name of the caller (Rolling/Expanding)
 
     Returns
     -------
     Numba function
     """
     nopython, nogil, parallel = get_jit_arguments(engine_kwargs, kwargs)
 
-    cache_key = (func, "rolling_apply")
+    cache_key = (func, f"{name}_apply_single")
     if cache_key in NUMBA_FUNC_CACHE:
         return NUMBA_FUNC_CACHE[cache_key]
 
@@ -161,3 +164,67 @@ def groupby_ewma(
         return result
 
     return groupby_ewma
+
+
+def generate_numba_table_func(
+    args: Tuple,
+    kwargs: Dict[str, Any],
+    func: Callable[..., Scalar],
+    engine_kwargs: Optional[Dict[str, bool]],
+    name: str,
+):
+    """
+    Generate a numba jitted function to apply window calculations table-wise
+
+    1. jit the user's function
+    2. Return a rolling apply function with the jitted function inline
+
+    Parameters
+    ----------
+    args : tuple
+        *args to be passed into the function
+    kwargs : dict
+        **kwargs to be passed into the function
+    func : function
+        function to be applied to each window and will be JITed
+    engine_kwargs : dict
+        dictionary of arguments to be passed into numba.jit
+    name : str
+        caller (Rolling/Expanding) and original method name for numba cache key
+
+    Returns
+    -------
+    Numba function
+    """
+    nopython, nogil, parallel = get_jit_arguments(engine_kwargs, kwargs)
+
+    cache_key = (func, f"{name}_table")
+    if cache_key in NUMBA_FUNC_CACHE:
+        return NUMBA_FUNC_CACHE[cache_key]
+
+    numba_func = jit_user_function(func, nopython, nogil, parallel)
+    numba = import_optional_dependency("numba")
+    if parallel:
+        loop_range = numba.prange
+    else:
+        loop_range = range
+
+    @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
+    def roll_table(
+        values: np.ndarray, begin: np.ndarray, end: np.ndarray, minimum_periods: int
+    ):
+        result = np.empty(values.shape)
+        min_periods_mask = np.empty(values.shape)
+        for i in loop_range(len(result)):
+            start = begin[i]
+            stop = end[i]
+            window = values[start:stop]
+            count_nan = np.sum(np.isnan(window), axis=0)
+            sub_result = numba_func(window, *args)
+            nan_mask = len(window) - count_nan >= minimum_periods
+            min_periods_mask[i, :] = nan_mask
+            result[i, :] = sub_result
+        result = np.where(min_periods_mask, result, np.nan)
+        return result
+
+    return roll_table

pandas/core/window/rolling.py

@@ -66,7 +66,10 @@
     GroupbyIndexer,
     VariableWindowIndexer,
 )
-from pandas.core.window.numba_ import generate_numba_apply_func
+from pandas.core.window.numba_ import (
+    generate_numba_apply_func,
+    generate_numba_table_func,
+)
 
 if TYPE_CHECKING:
     from pandas import DataFrame, Series
@@ -84,6 +87,7 @@ class BaseWindow(ShallowMixin, SelectionMixin):
         "axis",
         "on",
         "closed",
+        "method",
     ]
     exclusions: Set[str] = set()
 
@@ -97,6 +101,7 @@ def __init__(
         axis: Axis = 0,
         on: Optional[Union[str, Index]] = None,
         closed: Optional[str] = None,
+        method: str = "single",
         **kwargs,
     ):
 
@@ -110,6 +115,7 @@ def __init__(
         self.win_type = win_type
         self.win_freq = None
         self.axis = obj._get_axis_number(axis) if axis is not None else None
+        self.method = method
         self.validate()
 
     @property
@@ -159,6 +165,8 @@ def validate(self) -> None:
                     f"{type(self.window).__name__} does not implement "
                     f"the correct signature for get_window_bounds"
                 )
+        if self.method not in ["table", "single"]:
+            raise ValueError("method must be 'table' or 'single")
 
     def _create_data(self, obj: FrameOrSeries) -> FrameOrSeries:
         """
@@ -407,6 +415,26 @@ def hfunc(bvalues: ArrayLike) -> ArrayLike:
         self._insert_on_column(out, obj)
         return out
 
+    def _apply_tablewise(
+        self, homogeneous_func: Callable[..., ArrayLike], name: Optional[str] = None
+    ) -> FrameOrSeriesUnion:
+        if self._selected_obj.ndim == 1:
+            raise ValueError("method='table' not applicable for Series objects.")
+        obj = self._create_data(self._selected_obj)
+        values = self._prep_values(obj.to_numpy())
+        values = values.T if self.axis == 1 else values
+        result = homogeneous_func(values)
+        result = result.T if self.axis == 1 else result
+        out = obj._constructor(result, index=obj.index, columns=obj.columns)
+
+        if out.shape[1] == 0 and obj.shape[1] > 0:
+            raise DataError("No numeric types to aggregate")
+        elif out.shape[1] == 0:
+            return obj.astype("float64")
+
+        self._insert_on_column(out, obj)
+        return out
+
     def _apply(
         self,
         func: Callable[..., Any],
@@ -455,7 +483,7 @@ def calc(x):
                 return func(x, start, end, min_periods)
 
             with np.errstate(all="ignore"):
-                if values.ndim > 1:
+                if values.ndim > 1 and self.method == "single":
                     result = np.apply_along_axis(calc, self.axis, values)
                 else:
                     result = calc(values)
@@ -466,7 +494,10 @@ def calc(x):
 
             return result
 
-        return self._apply_blockwise(homogeneous_func, name)
+        if self.method == "single":
+            return self._apply_blockwise(homogeneous_func, name)
+        else:
+            return self._apply_tablewise(homogeneous_func, name)
 
     def aggregate(self, func, *args, **kwargs):
         result, how = aggregate(self, func, *args, **kwargs)
@@ -887,6 +918,14 @@ class Window(BaseWindow):
         .. versionchanged:: 1.2.0
 
             The closed parameter with fixed windows is now supported.
+    method : str, default 'single'
+        Execute the rolling operation per single column or row (`'single'`)
+        or over the entire object (`'table'`).
+
+        This argument is only implemented when specifying ``engine='numba'``
+        in the method call.
+
+        .. versionadded:: 1.3.0
 
     Returns
     -------
@@ -1033,6 +1072,9 @@ def validate(self):
         else:
             raise ValueError(f"Invalid window {self.window}")
 
+        if self.method != "single":
+            raise NotImplementedError("'single' is the only supported method type.")
+
     def _center_window(self, result: np.ndarray, offset: int) -> np.ndarray:
         """
         Center the result in the window for weighted rolling aggregations.
@@ -1307,8 +1349,17 @@ def apply(
         if maybe_use_numba(engine):
             if raw is False:
                 raise ValueError("raw must be `True` when using the numba engine")
-            apply_func = generate_numba_apply_func(args, kwargs, func, engine_kwargs)
-            numba_cache_key = (func, "rolling_apply")
+            caller_name = type(self).__name__
+            if self.method == "single":
+                apply_func = generate_numba_apply_func(
+                    args, kwargs, func, engine_kwargs, caller_name
+                )
+                numba_cache_key = (func, f"{caller_name}_apply_single")
+            else:
+                apply_func = generate_numba_table_func(
+                    args, kwargs, func, engine_kwargs, f"{caller_name}_apply"
+                )
+                numba_cache_key = (func, f"{caller_name}_apply_table")
         elif engine in ("cython", None):
             if engine_kwargs is not None:
                 raise ValueError("cython engine does not accept engine_kwargs")