-
-
Notifications
You must be signed in to change notification settings - Fork 18.5k
Implemented NumbaExecutionEngine #61487
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for working on this @arthurlw. Great start, I added a couple of comments that I think should be useful.
pandas/core/apply.py
Outdated
args, kwargs = prepare_function_arguments( | ||
self.func, # type: ignore[arg-type] | ||
engine_obj = NumbaExecutionEngine() | ||
result = engine_obj.apply( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
doc/source/whatsnew/v3.0.0.rst
Outdated
@@ -31,6 +31,7 @@ Other enhancements | |||
- :class:`pandas.api.typing.FrozenList` is available for typing the outputs of :attr:`MultiIndex.names`, :attr:`MultiIndex.codes` and :attr:`MultiIndex.levels` (:issue:`58237`) | |||
- :class:`pandas.api.typing.SASReader` is available for typing the output of :func:`read_sas` (:issue:`55689`) | |||
- :meth:`pandas.api.interchange.from_dataframe` now uses the `PyCapsule Interface <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html>`_ if available, only falling back to the Dataframe Interchange Protocol if that fails (:issue:`60739`) | |||
- Added :class:`pandas.core.apply.NumbaExecutionEngine` as the built-in ``numba`` execution engine for ``apply`` and ``map`` operations (:issue:`61458`) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
closes #xxxx (Replace xxxx with the GitHub issue number)Added type annotations to new arguments/methods/functions.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.