ENH: Add numba engine to rolling/expanding.std/var (pandas-dev#44461)

Author: mroeschke

asv_bench/benchmarks/rolling.py

@@ -53,7 +53,7 @@ class NumbaEngineMethods:
         ["DataFrame", "Series"],
         ["int", "float"],
         [("rolling", {"window": 10}), ("expanding", {})],
-        ["sum", "max", "min", "median", "mean"],
+        ["sum", "max", "min", "median", "mean", "var", "std"],
         [True, False],
         [None, 100],
     )

doc/source/whatsnew/v1.4.0.rst

@@ -214,6 +214,7 @@ Other enhancements
 - :meth:`Timestamp.isoformat`, now handles the ``timespec`` argument from the base :class:``datetime`` class (:issue:`26131`)
 - :meth:`NaT.to_numpy` ``dtype`` argument is now respected, so ``np.timedelta64`` can be returned (:issue:`44460`)
 - New option ``display.max_dir_items`` customizes the number of columns added to :meth:`Dataframe.__dir__` and suggested for tab completion (:issue:`37996`)
+- :meth:`.Rolling.var`, :meth:`.Expanding.var`, :meth:`.Rolling.std`, :meth:`.Expanding.std` now support `Numba <http://numba.pydata.org/>`_ execution with the ``engine`` keyword (:issue:`44461`)
 
 
 .. ---------------------------------------------------------------------------

pandas/core/_numba/executor.py

@@ -51,10 +51,11 @@ def column_looper(
         start: np.ndarray,
         end: np.ndarray,
         min_periods: int,
+        *args,
     ):
         result = np.empty((len(start), values.shape[1]), dtype=np.float64)
         for i in numba.prange(values.shape[1]):
-            result[:, i] = func(values[:, i], start, end, min_periods)
+            result[:, i] = func(values[:, i], start, end, min_periods, *args)
         return result
 
     return column_looper

pandas/core/_numba/kernels/__init__.py

@@ -1,4 +1,5 @@
 from pandas.core._numba.kernels.mean_ import sliding_mean
 from pandas.core._numba.kernels.sum_ import sliding_sum
+from pandas.core._numba.kernels.var_ import sliding_var
 
-__all__ = ["sliding_mean", "sliding_sum"]
+__all__ = ["sliding_mean", "sliding_sum", "sliding_var"]

pandas/core/_numba/kernels/var_.py

@@ -0,0 +1,116 @@
+"""
+Numba 1D var kernels that can be shared by
+* Dataframe / Series
+* groupby
+* rolling / expanding
+
+Mirrors pandas/_libs/window/aggregation.pyx
+"""
+from __future__ import annotations
+
+import numba
+import numpy as np
+
+from pandas.core._numba.kernels.shared import is_monotonic_increasing
+
+
+@numba.jit(nopython=True, nogil=True, parallel=False)
+def add_var(
+    val: float, nobs: int, mean_x: float, ssqdm_x: float, compensation: float
+) -> tuple[int, float, float, float]:
+    if not np.isnan(val):
+        nobs += 1
+        prev_mean = mean_x - compensation
+        y = val - compensation
+        t = y - mean_x
+        compensation = t + mean_x - y
+        delta = t
+        if nobs:
+            mean_x += delta / nobs
+        else:
+            mean_x = 0
+        ssqdm_x += (val - prev_mean) * (val - mean_x)
+    return nobs, mean_x, ssqdm_x, compensation
+
+
+@numba.jit(nopython=True, nogil=True, parallel=False)
+def remove_var(
+    val: float, nobs: int, mean_x: float, ssqdm_x: float, compensation: float
+) -> tuple[int, float, float, float]:
+    if not np.isnan(val):
+        nobs -= 1
+        if nobs:
+            prev_mean = mean_x - compensation
+            y = val - compensation
+            t = y - mean_x
+            compensation = t + mean_x - y
+            delta = t
+            mean_x -= delta / nobs
+            ssqdm_x -= (val - prev_mean) * (val - mean_x)
+        else:
+            mean_x = 0
+            ssqdm_x = 0
+    return nobs, mean_x, ssqdm_x, compensation
+
+
+@numba.jit(nopython=True, nogil=True, parallel=False)
+def sliding_var(
+    values: np.ndarray,
+    start: np.ndarray,
+    end: np.ndarray,
+    min_periods: int,
+    ddof: int = 1,
+) -> np.ndarray:
+    N = len(start)
+    nobs = 0
+    mean_x = 0.0
+    ssqdm_x = 0.0
+    compensation_add = 0.0
+    compensation_remove = 0.0
+
+    min_periods = max(min_periods, 1)
+    is_monotonic_increasing_bounds = is_monotonic_increasing(
+        start
+    ) and is_monotonic_increasing(end)
+
+    output = np.empty(N, dtype=np.float64)
+
+    for i in range(N):
+        s = start[i]
+        e = end[i]
+        if i == 0 or not is_monotonic_increasing_bounds:
+            for j in range(s, e):
+                val = values[j]
+                nobs, mean_x, ssqdm_x, compensation_add = add_var(
+                    val, nobs, mean_x, ssqdm_x, compensation_add
+                )
+        else:
+            for j in range(start[i - 1], s):
+                val = values[j]
+                nobs, mean_x, ssqdm_x, compensation_remove = remove_var(
+                    val, nobs, mean_x, ssqdm_x, compensation_remove
+                )
+
+            for j in range(end[i - 1], e):
+                val = values[j]
+                nobs, mean_x, ssqdm_x, compensation_add = add_var(
+                    val, nobs, mean_x, ssqdm_x, compensation_add
+                )
+
+        if nobs >= min_periods and nobs > ddof:
+            if nobs == 1:
+                result = 0.0
+            else:
+                result = ssqdm_x / (nobs - ddof)
+        else:
+            result = np.nan
+
+        output[i] = result
+
+        if not is_monotonic_increasing_bounds:
+            nobs = 0
+            mean_x = 0.0
+            ssqdm_x = 0.0
+            compensation_remove = 0.0
+
+    return output

pandas/core/window/doc.py

@@ -98,14 +98,17 @@ def create_section_header(header: str) -> str:
     "extended documentation and performance considerations for the Numba engine.\n\n"
 )
 
-window_agg_numba_parameters = dedent(
-    """
+
+def window_agg_numba_parameters(version: str = "1.3") -> str:
+    return (
+        dedent(
+            """
     engine : str, default None
         * ``'cython'`` : Runs the operation through C-extensions from cython.
         * ``'numba'`` : Runs the operation through JIT compiled code from numba.
         * ``None`` : Defaults to ``'cython'`` or globally setting ``compute.use_numba``
 
-          .. versionadded:: 1.3.0
+          .. versionadded:: {version}.0
 
     engine_kwargs : dict, default None
         * For ``'cython'`` engine, there are no accepted ``engine_kwargs``
@@ -114,6 +117,9 @@ def create_section_header(header: str) -> str:
           ``False``. The default ``engine_kwargs`` for the ``'numba'`` engine is
           ``{{'nopython': True, 'nogil': False, 'parallel': False}}``
 
-          .. versionadded:: 1.3.0\n
+          .. versionadded:: {version}.0\n
     """
-).replace("\n", "", 1)
+        )
+        .replace("\n", "", 1)
+        .replace("{version}", version)
+    )

pandas/core/window/ewm.py

@@ -511,7 +511,7 @@ def aggregate(self, func, *args, **kwargs):
         template_header,
         create_section_header("Parameters"),
         args_compat,
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -565,7 +565,7 @@ def mean(self, *args, engine=None, engine_kwargs=None, **kwargs):
         template_header,
         create_section_header("Parameters"),
         args_compat,
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,

pandas/core/window/expanding.py

@@ -227,7 +227,7 @@ def apply(
         template_header,
         create_section_header("Parameters"),
         args_compat,
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -253,7 +253,7 @@ def sum(
         template_header,
         create_section_header("Parameters"),
         args_compat,
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -279,7 +279,7 @@ def max(
         template_header,
         create_section_header("Parameters"),
         args_compat,
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -305,7 +305,7 @@ def min(
         template_header,
         create_section_header("Parameters"),
         args_compat,
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -330,7 +330,7 @@ def mean(
     @doc(
         template_header,
         create_section_header("Parameters"),
-        window_agg_numba_parameters,
+        window_agg_numba_parameters(),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -361,6 +361,7 @@ def median(
         """
         ).replace("\n", "", 1),
         args_compat,
+        window_agg_numba_parameters("1.4"),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -396,9 +397,18 @@ def median(
         aggregation_description="standard deviation",
         agg_method="std",
     )
-    def std(self, ddof: int = 1, *args, **kwargs):
+    def std(
+        self,
+        ddof: int = 1,
+        *args,
+        engine: str | None = None,
+        engine_kwargs: dict[str, bool] | None = None,
+        **kwargs,
+    ):
         nv.validate_expanding_func("std", args, kwargs)
-        return super().std(ddof=ddof, **kwargs)
+        return super().std(
+            ddof=ddof, engine=engine, engine_kwargs=engine_kwargs, **kwargs
+        )
 
     @doc(
         template_header,
@@ -411,6 +421,7 @@ def std(self, ddof: int = 1, *args, **kwargs):
         """
         ).replace("\n", "", 1),
         args_compat,
+        window_agg_numba_parameters("1.4"),
         kwargs_compat,
         create_section_header("Returns"),
         template_returns,
@@ -446,9 +457,18 @@ def std(self, ddof: int = 1, *args, **kwargs):
         aggregation_description="variance",
         agg_method="var",
     )
-    def var(self, ddof: int = 1, *args, **kwargs):
+    def var(
+        self,
+        ddof: int = 1,
+        *args,
+        engine: str | None = None,
+        engine_kwargs: dict[str, bool] | None = None,
+        **kwargs,
+    ):
         nv.validate_expanding_func("var", args, kwargs)
-        return super().var(ddof=ddof, **kwargs)
+        return super().var(
+            ddof=ddof, engine=engine, engine_kwargs=engine_kwargs, **kwargs
+        )
 
     @doc(
         template_header,