PERF: expanding/rolling.min/max with engine='numba' (#45170)

Author: mroeschke

doc/source/whatsnew/v1.4.0.rst

@@ -665,7 +665,7 @@ Performance improvements
 - :meth:`SparseArray.min` and :meth:`SparseArray.max` no longer require converting to a dense array (:issue:`43526`)
 - Indexing into a :class:`SparseArray` with a ``slice`` with ``step=1`` no longer requires converting to a dense array (:issue:`43777`)
 - Performance improvement in :meth:`SparseArray.take` with ``allow_fill=False`` (:issue:`43654`)
-- Performance improvement in :meth:`.Rolling.mean`, :meth:`.Expanding.mean`, :meth:`.Rolling.sum`, :meth:`.Expanding.sum` with ``engine="numba"`` (:issue:`43612`, :issue:`44176`)
+- Performance improvement in :meth:`.Rolling.mean`, :meth:`.Expanding.mean`, :meth:`.Rolling.sum`, :meth:`.Expanding.sum`, :meth:`.Rolling.max`, :meth:`.Expanding.max`, :meth:`.Rolling.min` and :meth:`.Expanding.min` with ``engine="numba"`` (:issue:`43612`, :issue:`44176`, :issue:`45170`)
 - Improved performance of :meth:`pandas.read_csv` with ``memory_map=True`` when file encoding is UTF-8 (:issue:`43787`)
 - Performance improvement in :meth:`RangeIndex.sort_values` overriding :meth:`Index.sort_values` (:issue:`43666`)
 - Performance improvement in :meth:`RangeIndex.insert` (:issue:`43988`)

pandas/_libs/window/aggregations.pyx

@@ -960,7 +960,7 @@ cdef _roll_min_max(ndarray[numeric_t] values,
     with nogil:
 
         # This is using a modified version of the C++ code in this
-        # SO post: http://bit.ly/2nOoHlY
+        # SO post: https://stackoverflow.com/a/12239580
         # The original impl didn't deal with variable window sizes
         # So the code was optimized for that
 

pandas/core/_numba/kernels/__init__.py

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

pandas/core/_numba/kernels/min_max_.py

@@ -0,0 +1,70 @@
+"""
+Numba 1D min/max 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
+
+
+@numba.jit(nopython=True, nogil=True, parallel=False)
+def sliding_min_max(
+    values: np.ndarray,
+    start: np.ndarray,
+    end: np.ndarray,
+    min_periods: int,
+    is_max: bool,
+) -> np.ndarray:
+    N = len(start)
+    nobs = 0
+    output = np.empty(N, dtype=np.float64)
+    # Use deque once numba supports it
+    # https://github.com/numba/numba/issues/7417
+    Q: list = []
+    W: list = []
+    for i in range(N):
+
+        curr_win_size = end[i] - start[i]
+        if i == 0:
+            st = start[i]
+        else:
+            st = end[i - 1]
+
+        for k in range(st, end[i]):
+            ai = values[k]
+            if not np.isnan(ai):
+                nobs += 1
+            elif is_max:
+                ai = -np.inf
+            else:
+                ai = np.inf
+            # Discard previous entries if we find new min or max
+            if is_max:
+                while Q and ((ai >= values[Q[-1]]) or values[Q[-1]] != values[Q[-1]]):
+                    Q.pop()
+            else:
+                while Q and ((ai <= values[Q[-1]]) or values[Q[-1]] != values[Q[-1]]):
+                    Q.pop()
+            Q.append(k)
+            W.append(k)
+
+        # Discard entries outside and left of current window
+        while Q and Q[0] <= start[i] - 1:
+            Q.pop(0)
+        while W and W[0] <= start[i] - 1:
+            if not np.isnan(values[W[0]]):
+                nobs -= 1
+            W.pop(0)
+
+        # Save output based on index in input value array
+        if Q and curr_win_size > 0 and nobs >= min_periods:
+            output[i] = values[Q[0]]
+        else:
+            output[i] = np.nan
+
+    return output

pandas/core/window/rolling.py

@@ -1385,15 +1385,18 @@ def max(
         if maybe_use_numba(engine):
             if self.method == "table":
                 func = generate_manual_numpy_nan_agg_with_axis(np.nanmax)
+                return self.apply(
+                    func,
+                    raw=True,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
             else:
-                func = np.nanmax
+                from pandas.core._numba.kernels import sliding_min_max
 
-            return self.apply(
-                func,
-                raw=True,
-                engine=engine,
-                engine_kwargs=engine_kwargs,
-            )
+                return self._numba_apply(
+                    sliding_min_max, "rolling_max", engine_kwargs, True
+                )
         window_func = window_aggregations.roll_max
         return self._apply(window_func, name="max", **kwargs)
 
@@ -1408,15 +1411,18 @@ def min(
         if maybe_use_numba(engine):
             if self.method == "table":
                 func = generate_manual_numpy_nan_agg_with_axis(np.nanmin)
+                return self.apply(
+                    func,
+                    raw=True,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
             else:
-                func = np.nanmin
+                from pandas.core._numba.kernels import sliding_min_max
 
-            return self.apply(
-                func,
-                raw=True,
-                engine=engine,
-                engine_kwargs=engine_kwargs,
-            )
+                return self._numba_apply(
+                    sliding_min_max, "rolling_min", engine_kwargs, False
+                )
         window_func = window_aggregations.roll_min
         return self._apply(window_func, name="min", **kwargs)
 

pandas/tests/window/test_numba.py

@@ -61,7 +61,7 @@ def test_numba_vs_cython_rolling_methods(
         expected = getattr(roll, method)(engine="cython", **kwargs)
 
         # Check the cache
-        if method not in ("mean", "sum", "var", "std"):
+        if method not in ("mean", "sum", "var", "std", "max", "min"):
             assert (
                 getattr(np, f"nan{method}"),
                 "Rolling_apply_single",
@@ -88,7 +88,7 @@ def test_numba_vs_cython_expanding_methods(
         expected = getattr(expand, method)(engine="cython", **kwargs)
 
         # Check the cache
-        if method not in ("mean", "sum", "var", "std"):
+        if method not in ("mean", "sum", "var", "std", "max", "min"):
             assert (
                 getattr(np, f"nan{method}"),
                 "Expanding_apply_single",
@@ -150,15 +150,16 @@ def test_dont_cache_args(
         def add(values, x):
             return np.sum(values) + x
 
+        engine_kwargs = {"nopython": nopython, "nogil": nogil, "parallel": parallel}
         df = DataFrame({"value": [0, 0, 0]})
-        result = getattr(df, window)(**window_kwargs).apply(
-            add, raw=True, engine="numba", args=(1,)
+        result = getattr(df, window)(method=method, **window_kwargs).apply(
+            add, raw=True, engine="numba", engine_kwargs=engine_kwargs, args=(1,)
         )
         expected = DataFrame({"value": [1.0, 1.0, 1.0]})
         tm.assert_frame_equal(result, expected)
 
-        result = getattr(df, window)(**window_kwargs).apply(
-            add, raw=True, engine="numba", args=(2,)
+        result = getattr(df, window)(method=method, **window_kwargs).apply(
+            add, raw=True, engine="numba", engine_kwargs=engine_kwargs, args=(2,)
         )
         expected = DataFrame({"value": [2.0, 2.0, 2.0]})
         tm.assert_frame_equal(result, expected)