ENH: Add numba engine for rolling apply

doc/source/getting_started/install.rst

@@ -255,6 +255,7 @@ gcsfs                     0.2.2              Google Cloud Storage access
 html5lib                                     HTML parser for read_html (see :ref:`note <optional_html>`)
 lxml                      3.8.0              HTML parser for read_html (see :ref:`note <optional_html>`)
 matplotlib                2.2.2              Visualization
+numba                     0.46.0             Alternative execution engine for rolling operations
 openpyxl                  2.4.8              Reading / writing for xlsx files
 pandas-gbq                0.8.0              Google Big Query access
 psycopg2                                     PostgreSQL engine for sqlalchemy

environment.yml

@@ -72,6 +72,7 @@ dependencies:
   - matplotlib>=2.2.2  # pandas.plotting, Series.plot, DataFrame.plot
   - numexpr>=2.6.8
   - scipy>=1.1
+  - numba>=0.46.0
 
   # optional for io
   - beautifulsoup4>=4.6.0  # pandas.read_html

pandas/compat/_optional.py

@@ -27,6 +27,7 @@
     "xlrd": "1.1.0",
     "xlwt": "1.2.0",
     "xlsxwriter": "0.9.8",
+    "numba": "0.46.0",
 }
 
 

pandas/core/window/rolling.py

@@ -1246,9 +1246,21 @@ def count(self):
           objects instead.
           If you are just applying a NumPy reduction function this will
           achieve much better performance.
-
-    *args, **kwargs
-        Arguments and keyword arguments to be passed into func.
+    args : tuple, default None
+        Positional arguments to be passed into func
+    kwargs : dict, default None
+        Keyword arguments to be passed into func
+    engine : str, default 'cython'
+        Execution engine for the applied function.
+        * ``'cython'`` : Runs rolling apply through C-extensions from cython.
+        * ``'numba'`` : Runs rolling apply through JIT compiled code from numba.
+          Only available when ``raw`` is set to ``True``.
+    engine_kwargs : dict, default None
+        Arguments to specify for the execution engine.
+        * For ``'cython'`` engine, there are no accepted ``engine_kwargs``
+        * For ``'numba'`` engine, the engine can accept ``nopython``, ``nogil``
+          and ``parallel``. The default ``engine_kwargs`` for the ``'numba'`` engine is
+          ``{'nopython': True, 'nogil': False, 'parallel': False}``
 
     Returns
     -------
@@ -1262,16 +1274,48 @@ def count(self):
     """
     )
 
-    def apply(self, func, raw=False, args=(), kwargs={}):
-        from pandas import Series
-
+    def apply(
+        self,
+        func,
+        raw=False,
+        args=None,
+        kwargs=None,
+        engine="cython",
+        engine_kwargs=None,
+    ):
+        if args is None:
+            args = ()
+        if kwargs is None:
+            kwargs = {}
         kwargs.pop("_level", None)
         kwargs.pop("floor", None)
         window = self._get_window()
         offset = _offset(window, self.center)
         if not is_bool(raw):
             raise ValueError("raw parameter must be `True` or `False`")
 
+        if engine == "cython":
+            if engine_kwargs is not None:
+                raise ValueError("cython engine does not accept engine_kwargs")
+            apply_func = self._generate_cython_apply_func(
+                args, kwargs, raw, offset, func
+            )
+        elif engine == "numba":
+            if raw is False:
+                raise ValueError("raw must be `True` when using the numba engine")
+            apply_func = self._generate_numba_apply_func(
+                args, kwargs, func, engine_kwargs
+            )
+        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)
+
+    def _generate_cython_apply_func(self, args, kwargs, raw, offset, func):
+        from pandas import Series
+
         window_func = partial(
             self._get_cython_func_type("roll_generic"),
             args=args,
@@ -1286,9 +1330,76 @@ def apply_func(values, begin, end, min_periods, raw=raw):
                 values = Series(values, index=self.obj.index)
             return window_func(values, begin, end, min_periods)
 
-        # 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 apply_func
+
+    def _generate_numba_apply_func(self, args, kwargs, func, engine_kwargs):
+        numba = import_optional_dependency("numba")
+
+        if engine_kwargs is None:
+            engine_kwargs = {"nopython": True, "nogil": False, "parallel": False}
+
+        nopython = engine_kwargs.get("nopython", True)
+        nogil = engine_kwargs.get("nogil", False)
+        parallel = engine_kwargs.get("parallel", False)
+
+        if kwargs and nopython:
+            raise ValueError(
+                "numba does not support kwargs with nopython=True: "
+                "https://github.com/numba/numba/issues/2916"
+            )
+
+        if parallel:
+            loop_range = numba.prange
+        else:
+            loop_range = range
+
+        def make_rolling_apply(func):
+            """
+            1. jit the user's function
+            2. Return a rolling apply function with the jitted function inline
+
+            Configurations specified in engine_kwargs apply to both the user's
+            function _AND_ the rolling apply function.
+            """
+
+            @numba.generated_jit(nopython=nopython)
+            def numba_func(window, *_args):
+                if getattr(np, func.__name__, False) is func:
+
+                    def impl(window, *_args):
+                        return func(window, *_args)
+
+                    return impl
+                else:
+                    jf = numba.jit(func, nopython=nopython)
+
+                    def impl(window, *_args):
+                        return jf(window, *_args)
+
+                    return impl
+
+            @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
+            def roll_apply(
+                values: np.ndarray,
+                begin: np.ndarray,
+                end: np.ndarray,
+                minimum_periods: int,
+            ):
+                result = np.empty(len(begin))
+                for i in loop_range(len(result)):
+                    start = begin[i]
+                    stop = end[i]
+                    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
+
+            return roll_apply
+
+        return make_rolling_apply(func)
 
     def sum(self, *args, **kwargs):
         nv.validate_window_func("sum", args, kwargs)
@@ -1934,8 +2045,23 @@ def count(self):
 
     @Substitution(name="rolling")
     @Appender(_shared_docs["apply"])
-    def apply(self, func, raw=False, args=(), kwargs={}):
-        return super().apply(func, raw=raw, args=args, kwargs=kwargs)
+    def apply(
+        self,
+        func,
+        raw=False,
+        args=None,
+        kwargs=None,
+        engine="cython",
+        engine_kwargs=None,
+    ):
+        return super().apply(
+            func,
+            raw=raw,
+            args=args,
+            kwargs=kwargs,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
 
     @Substitution(name="rolling")
     @Appender(_shared_docs["sum"])

pandas/tests/window/test_api.py

@@ -342,3 +342,31 @@ def test_multiple_agg_funcs(self, func, window_size, expected_vals):
         )
 
         tm.assert_frame_equal(result, expected)
+
+
+class TestEngine:
+    def test_invalid_engine(self):
+        with pytest.raises(
+            ValueError, match="engine must be either 'numba' or 'cython'"
+        ):
+            Series(range(1)).rolling(1).apply(lambda x: x, engine="foo")
+
+    def test_invalid_engine_kwargs_cython(self):
+        with pytest.raises(
+            ValueError, match="cython engine does not accept engine_kwargs"
+        ):
+            Series(range(1)).rolling(1).apply(
+                lambda x: x, engine="cython", engine_kwargs={"nopython": False}
+            )
+
+    def test_invalid_raw_numba(self):
+        with pytest.raises(
+            ValueError, match="raw must be `True` when using the numba engine"
+        ):
+            Series(range(1)).rolling(1).apply(lambda x: x, raw=False, engine="numba")
+
+    def test_invalid_kwargs_nopython(self):
+        with pytest.raises(ValueError, match="numba does not support kwargs with"):
+            Series(range(1)).rolling(1).apply(
+                lambda x: x, kwargs={"a": 1}, engine="numba", raw=True
+            )