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.5.7              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/numba_.py

@@ -0,0 +1,77 @@
+from typing import Callable, Dict, Optional, Tuple
+
+import numpy as np
+
+from pandas.compat._optional import import_optional_dependency
+
+
+def _generate_numba_apply_func(
+    args: Tuple, kwargs: Dict, func: Callable, engine_kwargs: Optional[Dict]
+):
+    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)

pandas/core/window/rolling.py

@@ -54,6 +54,7 @@
     FixedWindowIndexer,
     VariableWindowIndexer,
 )
+from pandas.core.window.numba_ import _generate_numba_apply_func
 
 
 class _Window(PandasObject, ShallowMixin, SelectionMixin):
@@ -1239,9 +1240,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.
+    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}``
+    args : tuple, default None
+        Positional arguments to be passed into func
+    kwargs : dict, default None
+        Keyword arguments to be passed into func
 
     Returns
     -------
@@ -1255,16 +1268,46 @@ def count(self):
     """
     )
 
-    def apply(self, func, raw=False, args=(), kwargs={}):
-        from pandas import Series
-
+    def apply(
+        self,
+        func: Callable,
+        raw: bool = False,
+        engine: str = "cython",
+        engine_kwargs: Optional[Dict] = None,
+        args: Optional[Tuple] = None,
+        kwargs: Optional[Dict] = 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 = _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,
@@ -1279,9 +1322,7 @@ 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 sum(self, *args, **kwargs):
         nv.validate_window_func("sum", args, kwargs)
@@ -1927,8 +1968,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,
+        engine="cython",
+        engine_kwargs=None,
+        args=None,
+        kwargs=None,
+    ):
+        return super().apply(
+            func,
+            raw=raw,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+            args=args,
+            kwargs=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
+            )

requirements-dev.txt

@@ -47,6 +47,7 @@ jinja2
 matplotlib>=2.2.2
 numexpr>=2.6.8
 scipy>=1.1
+numba>=0.46.0
 beautifulsoup4>=4.6.0
 fastparquet>=0.3.2
 html5lib