WIP: Add Numba to rolling.apply

Author: Matt Roeschke

pandas/core/window/aggregators/methods.py

@@ -1,3 +1,6 @@
+from typing import Callable
+
+import numba
 import numpy as np
 
 
@@ -33,3 +36,23 @@ def rolling_mean(
             val = total / count
         result[i] = val
     return result
+
+
+@numba.njit(nogil=True)
+def rolling_apply(
+    values: np.ndarray,
+    begin: np.ndarray,
+    end: np.ndarray,
+    minimum_periods: int,
+    numba_func: Callable,
+    args,
+) -> np.ndarray:
+    result = np.empty(len(begin))
+    for i, (start, stop) in enumerate(zip(begin, end)):
+        window = values[start:stop]
+        count_nan = np.sum(np.isnan(window))
+        if len(window) - count_nan >= minimum_periods:
+            result[i] = numba_func(window, *args)
+        else:
+            result[i] = np.nan
+    return result

pandas/core/window/rolling.py

@@ -8,6 +8,7 @@
 from typing import Callable, List, Optional, Set, Union
 import warnings
 
+import numba
 import numpy as np
 
 import pandas._libs.window as libwindow
@@ -1144,7 +1145,20 @@ def f(arg, window, min_periods, closed):
                 kwargs,
             )
 
-        return self._apply(f, func, args=args, kwargs=kwargs, center=False, raw=raw)
+        # Numba doesn't support kwargs in nopython mode
+        # https://github.com/numba/numba/issues/2916
+        numba_func = numba.njit(func)
+        rolling_apply = partial(methods.rolling_apply, numba_func=numba_func, args=args)
+
+        return self._apply(
+            rolling_apply,
+            func,
+            args=args,
+            kwargs=kwargs,
+            center=False,
+            raw=raw,
+            use_numba=True,
+        )
 
     def sum(self, *args, **kwargs):
         nv.validate_window_func("sum", args, kwargs)