ENH: Allow JIT compilation with an internal API
State: closed

- [X] closes #60668, supersedes #60622

I've been exploring what @MarcoGorelli proposed in https://github.com/pandas-dev/pandas/pull/60622#pullrequestreview-2584711239 regarding allowing JIT compilers to implement their own apply engine. I think it is a much better approach than the current implementation. I've been testing different APIs, both the user facing one, and the internal (what Numba and Bodo have to implement that we will call), and this PR is what I think makes more sense...

The approach here is to use a `jit` parameter for any function that could make sense to JIT in pandas (`DataFrame.apply`, `Series.map`, `SeriesGroupBy.transform`...) that delegates to the JIT compiler (Numba or Bodo) 100% of the logic. So far I just implement `DataFrame.apply` for simplicity, but happy to add the rest once there is agreement.

Final user API would look like next examples:

```python
df.apply(lambda x: x.A + x.B, axis=1, jit=bodo.njit)
df.apply(lambda x: x.A + x.B, axis=1, jit=bodo.jit(parallel=True))
```

Which I think it's very simple and intuitive, and at the same time it makes users import `numba` and `bodo` themselves, creating the right impression that they are using those libraries to JIT compile, and it's not something provided by pandas. Also, it makes users install the library themselves, which I think makes things easier (any installation or importing error is not reported to pandas, no need for a soft dependency, no need to even add them to our CI if we don't want).

I think this approach is very convenient for us, as maintaining the code in pandas is trivial. Which I think it should address the concerns @jbrockmendel expressed, and I think many of us share. It should also be very convenient for Numba and Bodo, which should be able to fix any issue much faster, as well as performance improvements, and support for new use cases . Numba and Bodo can probably release faster than us, so any needed work in the JIT functionality shouldn't consume pandas resources, make Numba and Bodo life easier, and end up in the users faster.

The exact internal API (the `__pandas_udf__` function in this PR) can probably be improved. To make it looks simple and reasonable, but happy to know other points of view.

As said in the code, the `bodo_patched.py` file is only in the PR so there is more context on the API that we will call. But the code will be properly implemented in Numba and Bodo, and removed from this PR before it's merged.

@ehsantn @scott-routledge2 @pandas-dev/pandas-core 

Labels:
name: Enhancement
name: Apply
description: Apply, Aggregate, Transform, Map

Comments:
This looks great to me!

> The exact internal API (the `__pandas_udf__` function in this PR) can probably be improved. To make it looks simple and reasonable, but happy to know other points of view.

I think it would be good to include `result_dtype` (can be `None` for current behavior), but that may be difficult with how flexible apply is. Offhand I would imagine it could be a single dtype (we'll raise if result has more than one column), list of dtypes (we'll raise if the number of columns is different from the length of the list), or a dictionary (we'll raise if key is missing). But this might require a bit more thought.
@DrTodd13 @sklam @stuartarchibald What do you think of implementing this interface in Numba?
@ehsantn, To clarify, do you mean that Bodo will be the main interface (given it has the support for pandas datastructure) and Numba may need to add things to properly support such use?
@sklam, the proposal here is that pandas provides a jit parameter to the functions that makes sense (`DataFrame.apply`, `Series.map`...). And if the user provides numba (the `numba.jit` decorator in the proposal, but can be discussed), then pandas will call numba with the parameters of the call (the dataframe, the function to be jitted, the axis...) and numba will be the one deciding how to execute that code and return the result to pandas.
@datapythonista wouldn't you need to add a test for this?
> @datapythonista wouldn't you need to add a test for this?

Yes, thanks for pointing that out. I wanted to make the PR simple first, mainly in case there is feedback on the API between pandas and the JIT compilers. I think this first version of the PR should make reviewing the API simple. Once the API has been discussed and we are happy with it, I'll be adding the tests, a release note, and I'll review if there is anywhere else in the docs where this feature should be mentioned.
pandas still could cache the jit function under this protocol, correct? Currently pandas does that for numba to avoid overhead of jitting the same function more than once
> I think it would be good to include result_dtype [...]

Exactly what would be useful is going to differ by implementation, and may not be stable.  This is a reason for the correct usage to be `whatever.apply(func, obj, specific_keywords=...)`
> @sklam, the proposal here is that pandas provides a jit parameter to the functions that makes sense (`DataFrame.apply`, `Series.map`...). And if the user provides numba (the `numba.jit` decorator in the proposal, but can be discussed), then pandas will call numba with the parameters of the call (the dataframe, the function to be jitted, the axis...) and numba will be the one deciding how to execute that code and return the result to pandas.

I guess I have the same question as @sklam.  When you say "parameters of the call" and list dataframe as an example then given that numba doesn't support dataframes, I question what you expect to happen.  I can imagine a conversion layer in Python that could convert a dataframe to a set of numpy arrays and convert references to columns in the lambda to those individual arrays and then call numba to compile and run that code and convert the result back.  (The rest of this message assumes this conversion layer is necessary.  Please disregard if this is somehow already happening and I'm not aware of it.)  The question is where would this conversion layer live.  If in numba, wouldn't this require numba to take a pandas dependency?  I'm pretty sure they wouldn't want to do that, partially for the same reason that pandas doesn't want to take a numba dependency but also because it would be considered out-of-scope for numba.

I think the issue of maintenance was hinted at in the first post.  If this doesn't work for some reason, then I guess the pandas devs' hope would be for the user to complain directly to bodo or numba.  Some may but I guarantee others won't.  The pandas shim is pretty thin so it is unlikely an error would be there but it is theoretically possible.  If there is a need for a conversion layer then there is a lot of room for error there in addition to the base numba compilation layers.  It is sounding to me like neither numba or pandas is going to want to be responsible for this conversion layer.  Perhaps a separate org or separate repo in the pandas org?  If that separate package is listed in the pandas docs as a suggested option then you might get some uptake but still have to identify who is willing to maintain it.
> pandas still could cache the jit function under this protocol, correct? Currently pandas does that for numba to avoid overhead of jitting the same function more than once

@mroeschke I think the way the interface is designed, and what makes more sense to me, is that are Numba and Bodo that cache their compiled functions. pandas will send the data, the function, and anything else needed to the execution backends, and will get the new data. For the Python/pandas backend there is nothing to cache, for Numba the function itself is jitted for what I've seen, and for Bodo they create a new function with the call to apply again, and that's what they jit. So, I think intuitively it could make sense that we take the jit decorator and the function in a generic way, and we cache it in pandas, but in practice I think it would make things very convoluted, compared to just let Numba and Bodo take care of it.
> for Numba the function itself is jitted for what I've seen

and we create another jit function to apply the UDF over each column/row which is the jit function we cache i.e.

1. We jit the UDF passed from `apply`
2. We embed this jitted UDF in another jit function designed to apply this UDF over each row/column and cache this function

The specifics can be found in https://github.com/pandas-dev/pandas/blob/56847c56c59bc3edf8607aaacb4f2d1b544204b2/pandas/core/_numba/executor.py#L21

(and as you can see, there are different flavors further down whether we're doing an `apply` or `groupby.apply` or `rolling.apply`)

So are you suggesting that Numba/Bodo cache the larger jitted function (2 in the list above) as well? 

> > for Numba the function itself is jitted for what I've seen
> 
> and we create another jit function to apply the UDF over each column/row which is the jit function we cache i.e.
> 
> 1. We jit the UDF passed from `apply`
> 2. We embed this jitted UDF in another jit function designed to apply this UDF over each row/column and cache this function
> 
> The specifics can be found in
> 
> https://github.com/pandas-dev/pandas/blob/56847c56c59bc3edf8607aaacb4f2d1b544204b2/pandas/core/_numba/executor.py#L21
> 
> (and as you can see, there are different flavors further down whether we're doing an `apply` or `groupby.apply` or `rolling.apply`)
> 
> So are you suggesting that Numba/Bodo cache the larger jitted function (2 in the list above) as well?

If you want the Numba/Bodo function to be cached, then you can just do functools.partial on the decorator you pass to jit/executor?  However, note that if you cache an outer function that the inner function it calls is also included in that module even if it isn't given as cache=True.  So, caching the inner one would help only in the case where the inner is the same but the outer function is different for some reason.
Some updates here:

- I changed the approach here and instead of trying to fully replace the existing approach with a new `jit` parameter, I'm implementing it as a hybrid approach of what we have now with Numba, and the new interface. So, the `engine` parameter accepts the strings `python` and `numba`, as well as the `numba.jit` and `bodo.jit` decorators when they implement the interface. Also, I leave for a follow up making the python and numba engines use the new interface. 
- Instead of passing the value of `raw` to the engines, now when `raw=True` pandas will convert the Series/DataFrame to a NumPy array before passing it to the engine. While for now the Numba engine will stay the same, this should help move it to the Numba codebase without making them depend on pandas.
- I improved the API that Numba and Bodo will implement, which before was just a function accepting what `DataFrame.apply` accepts. Now it's a base class with simpler methods.

Feedback welcome
Thanks all for the feedback, it was very helpful. This should be now ready, I added tests, a release note and CI should be green (other than the unrelated pyarrow failure).

In a follow up, we should move the python and numba engine to the new interface (in the pandas codebase), which I think should make the code much simpler. And then have the discussion on moving the numba engine to the numba codebase, which I think it should be quite straightforward at that point.

Also, seems like blosc2 implemented a jit for numpy operations, that can probably also benefit from this new interface: https://www.blosc.org/python-blosc2/reference/autofiles/utilities/blosc2.jit.html
> Also, seems like blosc2 implemented a jit for numpy operations, that can probably also benefit from this new interface: https://www.blosc.org/python-blosc2/reference/autofiles/utilities/blosc2.jit.html

Thanks for the intro, @datapythonista.  FWIW, I did a quick [micro-benchmark comparing speeds of blosc2 vs numba vs numpy](https://gist.github.com/FrancescAlted/ac5713abf7bc8501d867a328e71ab826) and run it on my desktop machine.

One of the advantages of the new compute engine in blosc2 is the reduced compile time (at least when compared to numba; see first point in plot above).  Then, thanks to its CPU cache-sensible blocking, it also scales quite well when operands are large.  The [operators and functions](https://www.blosc.org/python-blosc2/reference/lazy_functions.html) that are accelerated are essentially those in [numexpr](https://numexpr.readthedocs.io/en/latest/user_guide.html) (+, -, *, **, /, sin, cos...), plus [reductions](https://www.blosc.org/python-blosc2/reference/reduction_functions.html) (sum, prod, mean, std, var, min, max, all, any); others not in that subset should fallback and executed via plain numpy.  Last, but not least, dtype inference mimics more numpy (although not perfectly yet) than the numexpr one.

In case you find something that you may need to use `blosc2.jit` from pandas, tell us and we will try to help.
@DrTodd13 when you have the chance, can you check if your concerns for this PR are addressed (we surely need few follow ups to leave all this perfect). And if it's the case, remove the "requested chanches" tick.

@pandas-dev/pandas-core please let me know if there is any additional feedback worth discussing before merging this, and starting work on the follow up PRs mentioned above. Otherwise I'll be merging this soon.
Thanks all for the reviews and feedback. Merging this so we can continue with the follow ups. If there is any further comment, please let me know, happy to discuss and address in a follow up.