Skip to content

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

Open
wants to merge 7 commits into
base: main
Choose a base branch
from

Conversation

arthurlw
Copy link
Member

  • 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.

Copy link
Member

@datapythonista datapythonista left a 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.

args, kwargs = prepare_function_arguments(
self.func, # type: ignore[arg-type]
engine_obj = NumbaExecutionEngine()
result = engine_obj.apply(
Copy link
Member

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.

Copy link
Member Author

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?

Copy link
Member

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.

@@ -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`)
Copy link
Member

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants