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

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