Add apply function cache to save compiled numba functions

pandas-dev · jreback · Dec 27, 2019 · Dec 8, 2019 · Dec 8, 2019 · Dec 9, 2019
commit 0c140330b36ad25ff86d1b4ce3d615d050517928
diff --git a/pandas/core/window/common.py b/pandas/core/window/common.py
@@ -70,6 +70,7 @@ def _apply(
         floor: int = 1,
         is_weighted: bool = False,
         name: Optional[str] = None,
+        use_numba_cache: Optional = False,
         **kwargs,
     ):
         """

diff --git a/pandas/core/window/numba_.py b/pandas/core/window/numba_.py
@@ -6,7 +6,11 @@
 
 
 def _generate_numba_apply_func(
-    args: Tuple, kwargs: Dict, func: Callable, engine_kwargs: Optional[Dict]
+    args: Tuple,
+    kwargs: Dict,
-    kwargs: Dict,
+    kwargs: Dict[str, Any],
-    kwargs: Dict,
+    kwargs: Dict[str, Any],
+    func: Callable,
+    engine_kwargs: Optional[Dict],
-    engine_kwargs: Optional[Dict],
+    engine_kwargs: Optional[Dict[str, Any]],
-    engine_kwargs: Optional[Dict],
+    engine_kwargs: Optional[Dict[str, Any]],
+    function_cache: Dict,
 ):
     """
     Generate a numba jitted apply function specified by values from engine_kwargs.
@@ -37,6 +41,10 @@ def _generate_numba_apply_func(
     else:
         loop_range = range
 
+    # Return an already compiled version of roll_apply if available
+    if func in function_cache:
+        return function_cache[func]
+
     def make_rolling_apply(func):
 
         if isinstance(func, numba.targets.registry.CPUDispatcher):

diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py
@@ -93,6 +93,7 @@ def __init__(
         self.win_freq = None
         self.axis = obj._get_axis_number(axis) if axis is not None else None
         self.validate()
+        self._numba_func_cache = dict()
 
     @property
     def _constructor(self):
@@ -443,6 +444,7 @@ def _apply(
         floor: int = 1,
         is_weighted: bool = False,
         name: Optional[str] = None,
+        use_numba_cache: Optional[bool] = False,
         **kwargs,
     ):
         """
@@ -455,10 +457,11 @@ def _apply(
         func : callable function to apply
         center : bool
         require_min_periods : int
-        floor: int
-        is_weighted
-        name: str,
+        floor : int
+        is_weighted : bool
+        name : str,
             compatibility with groupby.rolling
+        use_numba_cache : bool
         **kwargs
             additional arguments for rolling function and window function
 
@@ -533,6 +536,9 @@ def calc(x):
                     result = calc(values)
                     result = np.asarray(result)
 
+            if use_numba_cache:
+                self._numba_func_cache[name] = func
+
             if center:
                 result = self._center_window(result, window)
 
@@ -1303,13 +1309,21 @@ def apply(
         elif engine == "numba":
             if raw is False:
                 raise ValueError("raw must be `True` when using the numba engine")
-            apply_func = _generate_numba_apply_func(args, kwargs, func, engine_kwargs)
+            apply_func = _generate_numba_apply_func(
+                args, kwargs, func, engine_kwargs, self._numba_func_cache
+            )
         else:
             raise ValueError("engine must be either 'numba' or 'cython'")
 
         # TODO: Why do we always pass center=False?
         # name=func for WindowGroupByMixin._apply
-        return self._apply(apply_func, center=False, floor=0, name=func)
+        return self._apply(
+            apply_func,
+            center=False,
+            floor=0,
+            name=func,
+            use_numba_cache=engine == "numba",
+        )
 
     def _generate_cython_apply_func(self, args, kwargs, raw, offset, func):
         from pandas import Series