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.

[arthurlw]

Hi @datapythonista, I was seeing a CI error because Numba isn’t installed in the test environment, so I tried to guard against it using a try and catch method, but it seems not to work. Do you have any advice on how to move forward?

[datapythonista]

This is what we use for optional dependencies in tests: https://github.com/pandas-dev/pandas/blob/main/pandas/tests/io/test_iceberg.py#L21

[arthurlw]

Thanks for the clarification in the comment @datapythonista

Just to clarify, for the NumbaExecutionEngine, this means that instead of having the condition if engine == "numba" inside of apply_raw, we should be calling NumbaExecutionEngine.apply directly in Dataframe.apply and let the engine handle all the unsupported cases by Numba.

This would involve:

1. removing all instances of if engine == "numba" in apply.py and moving them into NumbaExecutionEngine.apply
2. calling NumbaExecutionEngine.apply directly in Dataframe.apply instead of apply_raw

[datapythonista]

Correct

[datapythonista]

numba is not a required pandas dependency, I think this will make import pandas raise an ImportError for users with no numba installed, even if they don't call any numba functionality.

I think we've got an import_optional_dependency function that we call inside the functions where numba is used to avoid this problem.

[datapythonista]

What I think it'd be a simpler approach is to implement this logic here:

https://github.com/pandas-dev/pandas/blob/main/pandas/core/frame.py#L10563

There, now we are considering two cases:

• No engine (default engine) or string engine (numba engine)
• engine with __pandas_udf__

I would simplify that and just support engines with the engine interface __pandas_udf__:

• No engine
• __pandas_udf__

Since we want to support engine="numba" for now, for compatibility reasons, what I would do is immediately after DataFrame.apply is called, convert the "numba" string to a "fake" numba decorator with the __pandas_udf__ containing the the NumaExecutionEngine class. Something like:

def apply(...):
    if engine == "numba":
        numba = import_optional_dependency("numba")
        numba_jit = numba.jit(**engine_kwargs)
        numba_jit.__pandas_udf__ = NumbaExecutionEngine

From this point, all the code can pretend engine is going to be None for the default Python engine, or a __pandas_udf__ class, which should make things significantly.

The challenge is that numba and the default engine share some code, and with this approach they'll be running independently. The default engine won't know anything about an engine parameter, and the numba engine will run NumbaExecutionEngine.apply. If we don't want to repeat code, we'll probably have to restructure a bit the code, so some functions are generic and called by both engines.

When we move the default engine to a PythonExecutionEngine class, maybe it's a good idea to have the base class with the code reused by different engines. But I think that change is to big to address in a single PR, so I'd see what can be done for now that it's not too big of a change.

Does this approach makes sense to you?

[datapythonista]

Sorry, I added some comments before, but I just realized now I didn't submit the review with them. I explained there the idea on how to call NumbaExecutionEngine from DataFrame.apply