ENH: Add numba engine for rolling apply

[mroeschke]

• closes Add Numba as an optional dependency for rolling.apply for pandas 1.0 #28987
• tests added / passed
• passes black pandas
• passes git diff upstream/master -u -- "*.py" | flake8 --diff
• whatsnew entry

[TomAugspurger]

Random question: do you have benchmarks on this anywhere? It would be good to give guidance on when this is helpful (data sizes, type of user-defined function, ratio of time spent in windowing vs. the UDF, etc.).

[jreback]

args,kwargs should be at the end

[jreback]

type new arguments

[mroeschke]

@TomAugspurger performance testing is lightly described in twosigma#29, but good point to have a dedicated section in the documentation describing when the numba engine would be advantageous

[jreback]

can you add some comments here

[alimcmaster1]

Code check failures are related to https://github.com/pandas-dev/pandas/pull/30370/files#r361222880
I can submit a PR to fix this

Check for non-standard imports
##[error]pandas/io/stata.py:12:from collections.abc import Iterator
##[error]pandas/io/sas/sas7bdat.py:16:from collections.abc import Iterator
##[error]pandas/io/sas/sas_xport.py:10:from collections.abc import Iterator
##[error]pandas/io/common.py:5:from collections.abc import Iterator
##[error]pandas/io/parsers.py:6:from collections.abc import Iterator
##[error]pandas/io/json/_json.py:2:from collections.abc import Iterator

PR here: #30455

[WillAyd]

Few more comments. I don't want to go overboard on annotations so those in particular aren't blockers, but generally if subtypes can be added for anything they would be preferred

[WillAyd]

Not a blocker here since this is large enough, but would be nice to annotate this in a follow up

[WillAyd]

I think I missed this in previous review but similar comment as rest; not a block but subtypes make annotations much more insightful to the reader

[WillAyd]

This can just be annotated as bool no? Or do we need to explicitly handle None?

[WillAyd]

I think this class would be more logically placed in test_apply

[WillAyd]

Can we just skip or xfail instead? I think more indicative to reader that it is an invalid combination rather than re-assigning

[jbrockmendel]

could make engine_and_raw fixture using the pattern in #30456 if this is going to come up repeatedly

[WillAyd]

Can just keep as skip_if_no("numba") since the minimum is already enforced by requirements files

[jbrockmendel]

"the data set is larger" here is pretty vague. is the perf gain a function of the array size or more about the user-defined function?

[mroeschke]

Somewhat both, but more obvious with the data size. I can make this more specific.

[jbrockmendel]

add an X for this

[jbrockmendel]

This looks really nice. way to go

[mroeschke]

Thanks for the review everyone. Think I addressed all the comments.

[jreback]

lgtm. small questions.

cc @jorisvandenbossche @TomAugspurger @jbrockmendel

if any comments

[jreback]

can you add a doc-string that says what this function does (the parameters are already documented elsewhere, maybe just mention that)

[WillAyd]

lgtm

[mroeschke]

Thanks for another round of review. Addressed all the comments.

[TomAugspurger]

Looks good.

[jreback]

thanks @mroeschke very nice!

[jbrockmendel]

@mroeschke heads up that im now seeing this in test output locally:

pandas/tests/window/test_numba.py::TestApply::test_cache[False-True-False-False]
  /Users/bmendel/Desktop/pd/dtmean/pandas/core/window/numba_.py:75: NumbaPerformanceWarning: 
  The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible.
  
  To find out why, try turning on parallel diagnostics, see http://numba.pydata.org/numba-doc/latest/user/parallel.html#diagnostics for help.
  
  File "pandas/core/window/numba_.py", line 59:
  
              def impl(window, *_args):
              ^
  
    result[i] = numba_func(window, *args)

[mroeschke]

I'll make a followup to suppress this warning (no baring on what is being test).

[mroeschke]

@stuartarchibald sorry for the ping, but I see that generated_jit has been deprecated in numba 0.57. IIRC you helped me add this a while back and am lost on how to write this in terms of overload

[stuartarchibald]

@mroeschke no problem, I can try and help with this. I think it needs to look a bit like this (for reference, this is untested, I am just guessing from the context! Also, the pandas variant is obviously wrapped to close over some configuration which I've omitted, so consider this as the function body of make_rolling_apply. I've left comments inline to try and explain what's going on):

import types
import numpy as np
from numba.extending import overload, is_jitted
from numba import njit
import numba


# this provides a local definition to overload
def overload_target(window, *_args):
    # If JIT is disabled, this function will run, so write the implementation here!
    pass


nopython = True
nogil = True
parallel = False

# pretend this is an arg to `make_rolling_apply`
def func(window, *args):
    return window * 2 + args[0]


@overload(overload_target, jit_options={'nopython':nopython, 'nogil':nogil,
                                        'parallel':parallel})
def ol_overload_target(window, *_args):
    # This function "overloads" `overload_target`, whenever the Numba compiler
    # "sees" `overload_target` it will use this function.

    # Using `is_jitted` to avoid `isinstance` on
    # `numba.targets.registry.CPUDispatcher` as that may be considered an
    # internal Numba detail.
    if is_jitted(func):
        # it's already JIT compiled so just reference it
        overload_target_impl = func
    elif getattr(np, func.__name__, False) is func or isinstance(
        func, types.BuiltinFunctionType
    ):
        # it's a NumPy function or builtin so just reference it
        overload_target_impl = func
    else:
        # it's a Python function, so register it as JIT compilable and reference
        # that
        overload_target_impl = numba.jit(func, nopython=nopython, nogil=nogil)

    # This is the Numba implementation of the overload, it will just be JIT
    # compiled whenever the compiler "sees" a reference to "overload_target" in
    # code it is compiling.
    def impl(window, *_args):
        return overload_target_impl(window, *_args)

    return impl


# demo

@njit
def roll_apply(window, *_args):
    return overload_target(window, *_args)


print(roll_apply(np.arange(10.), 1.23))

@overload is basically saying to Numba "when you see this specific python function (the one in the first argument in the @overload decorator) use this implementation". The concept about there being a "typing" part that can be used to dispatch different variants based on type is exactly the same as in @generated_jit. The largest difference is what happens if the JIT compiler is turned off. In the case of @overload the python function being overloaded will run, i.e., the code just executes as would be expected in the interpreter. Whereas in the case of @generated_jit, because the pure python implementation and the Numba implementation are the same function, if you turn the JIT compiler off it will just break (the value returned when calling a @generated_jit function is a function implementing the Numba specialisation). Essentially, @generated_jit is like doing @overload but the function being decorated is also the function being overloaded.

Hope this helps?

[mroeschke]

Thanks for the reply! We had a PR recently that refactored this to use extending.register_jittable. Would that be a sufficient alternative? https://github.com/pandas-dev/pandas/pull/53455/files

[stuartarchibald]

No problem! I just took a look at the patch above, I think it'd work but think it might lose some of the dispatch ability offered by generated_jit/overload. As I understand it, the original code would have let a NumPy function or a built-in be passed in as the "user function", whereas I think the register_jittable version requires a user defined Python function. It may be that the register_jittable version is a sufficient alternative for the need/use cases in practice, in which case, it seems appropriate.

[mroeschke]

Great thanks for the context! Yeah this function should expect a custom UDF so thanks for the confirmation

[stuartarchibald]

Glad to get this resolved, thanks for confirming too! It sounds like the replacement above is appropriate. If there are any more issues/queries feel free to open issues on the Numba issue tracker (or ping here!).

[kartiksubbarao]

@stuartarchibald I'm running into a rolling apply issue with pandas 2.1.1 and numba 0.58 that might be related. Discussion is here:
https://numba.discourse.group/t/pandas-source-of-old-style-error-capturing-warning/2169/8