ENH: Add method argument to rolling constructor to allow table-wise rolling (pandas-dev#38417)

Author: mroeschke

asv_bench/benchmarks/rolling.py

@@ -252,4 +252,22 @@ def time_groupby_mean(self, engine):
         self.gb_ewm.mean(engine=engine)
 
 
+def table_method_func(x):
+    return np.sum(x, axis=0) + 1
+
+
+class TableMethod:
+
+    params = ["single", "table"]
+    param_names = ["method"]
+
+    def setup(self, method):
+        self.df = pd.DataFrame(np.random.randn(10, 1000))
+
+    def time_apply(self, method):
+        self.df.rolling(2, method=method).apply(
+            table_method_func, raw=True, engine="numba"
+        )
+
+
 from .pandas_vb_common import setup  # noqa: F401 isort:skip

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
+=============================   =================  ===========================   ===========================  ========================  ===================================
+Rolling window                  ``rolling``        ``Rolling``                   Yes                          Yes                       Yes (as of version 1.3)
+Weighted window                 ``rolling``        ``Window``                    No                           No                        No
+Expanding window                ``expanding``      ``Expanding``                 No                           Yes                       Yes (as of version 1.3)
+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,29 @@ 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 a :class:`DataFrame` with 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.
+
+For example, a `weighted mean <https://en.wikipedia.org/wiki/Weighted_arithmetic_mean>`__ calculation can
+be calculated with :meth:`~Rolling.apply` by specifying a separate column of weights.
+
+.. 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")  # noqa:E501
+
+
 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

@@ -33,6 +33,11 @@ For example:
         storage_options=headers
     )
 
+.. _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

@@ -10996,6 +10996,7 @@ def rolling(
         on: Optional[str] = None,
         axis: Axis = 0,
         closed: Optional[str] = None,
+        method: str = "single",
     ):
         axis = self._get_axis_number(axis)
 
@@ -11009,6 +11010,7 @@ def rolling(
                 on=on,
                 axis=axis,
                 closed=closed,
+                method=method,
             )
 
         return Rolling(
@@ -11020,12 +11022,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:
@@ -11037,7 +11044,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

@@ -242,6 +242,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 {'single', 'table'}, 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
     -------
@@ -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
+        )
 
     def _get_window_indexer(self) -> BaseIndexer:
         """

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]
 
@@ -153,3 +156,67 @@ def groupby_ewma(
         return result
 
     return groupby_ewma
+
+
+def generate_numba_table_func(
+    args: Tuple,
+    kwargs: Dict[str, Any],
+    func: Callable[..., np.ndarray],
+    engine_kwargs: Optional[Dict[str, bool]],
+    name: str,
+):
+    """
+    Generate a numba jitted function to apply window calculations table-wise.
+
+    Func will be passed a M window size x N number of columns array, and
+    must return a 1 x N number of columns array. Func is intended to operate
+    row-wise, but the result will be transposed for axis=1.
+
+    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")
+
+    @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 numba.prange(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