Implemented NumbaExecutionEngine

[arthurlw]

• closes #xxxx (Replace xxxx with the GitHub issue number)
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

Implements NumbaExecutionEngine for #61458

Docstring is currently a placeholder.

[datapythonista]

Thanks for working on this @arthurlw. Great start, I added a couple of comments that I think should be useful.

[datapythonista]

This is the ideal, but what would be even better is that apply itself delegates to the new executor class the numba execution, in the same way we do for other engines.

So, if we call df.apply(func, engine=bodo.jit), apply will delegate the execution to bodo.jit.__pandas_udf__.apply. Same will hopefully be true for numba.jit at some point. For that, Numba will be the one implementing the executor you are coding now. So, until that happens, we'll have it in pandas, but it'd be better if it behaves like a third-party execution engine already.

So, an idea is that before we delegate the execution to a third-party executin engine, we could do something like:

if engine == "numba":
    numba,jit.__pandas_udf__ = NumbaExecutorEngine

This way, when apply checks if the engine has a __pandas_udf__ attribute, it will already use the Numba executor like any other.

[arthurlw]

Thanks for the suggestion! The implementation would look something like this correct?

numba.jit.__pandas_udf__ = NumbaExecutionEngine
result = numba.jit.__pandas_udf__.apply(
    self.values,
    self.func,
    self.args,
    self.kwargs,
    engine_kwargs,
    self.axis,
)

Also just to clarify, this implementation should wait until Numba writes its own executor?

[datapythonista]

Yes, that's correct. But since we already support numba, I wouldn't wait until it's implemented in numba. I would create the execution engine class ourselves, and just simulate that things work the way you describe.

So, in the future we expect numba.jit to have the __pandas_udf__ attribute. But for now, if we receive engine=numba.jit and numba.jit doesn't have the attribute, we add the attribute ourselves with the engine class we implement, and we continue with the execution as if it had.

There may be other options, but this approach will keep background compatibility for now when engine="numba", will implement the executor class so numba can easily copy in their repo, and things will already work when they do. To me, this is the best. Only thing is that when engine='numba`` was implemented we didn't have the execution engine interface, so it was implemented with some if engine == "numba":` mixed with the default executor. That's what I think we should revert now. And keep things well organized with the new interface.

For reference, this is the implementation of the interface for blosc2, another jit compiler: https://github.com/Blosc/python-blosc2/pull/418/files. There are differences, since blosc2 is mostly for vectorized numpy operations, and numba should work well with jitting loops over numpy arrays. but the idea is somehow similar.

[datapythonista]

This note is for final pandas users, I think we don't need to share too much about the internal implementation (users in general won't know about NumbaExecutionEngine. What the change in this PR will ideally mean for users is that they'll be able to use df.apply(func, engine=numba.jit). I'd mention that instead.