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.

Sign up for GitHub

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 8 commits into
base: main
Choose a base branch
from
+80 −17
Open

Implemented NumbaExecutionEngine #61487

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
aa42037
Implemented NumbaExecutionEngine
arthurlw May 23, 2025
db9f3b0
whatsnew
arthurlw May 23, 2025
4cb240d
precommit
arthurlw May 23, 2025
97d9063
Match function arguments
arthurlw May 24, 2025
69e0e35
Fix CI
arthurlw May 24, 2025
7365079
updated whatsnew
arthurlw May 28, 2025
c605857
Updated conditions and delegate method to numba.jit
arthurlw May 29, 2025
24a0615
Added try and except to catch ImportError
arthurlw Jun 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/source/whatsnew/v3.0.0.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

- Added :meth:`.Styler.to_typst` to write Styler objects to file, buffer or string in Typst format (:issue:`57617`)
- Added missing :meth:`pandas.Series.info` to API reference (:issue:`60926`)
- :class:`pandas.api.typing.NoDefault` is available for typing ``no_default``
Expand Down
74 changes: 60 additions & 14 deletions pandas/core/apply.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,60 @@ def apply(
"""


class NumbaExecutionEngine(BaseExecutionEngine):
"""
Numba-based execution engine for pandas apply and map operations.
"""

@staticmethod
def map(
data: np.ndarray | Series | DataFrame,
func,
args: tuple,
kwargs: dict,
decorator: Callable | None,
skip_na: bool,
):
"""
Elementwise map for the Numba engine. Currently not supported.
"""
raise NotImplementedError("Numba map is not implemented yet.")

@staticmethod
def apply(
data: np.ndarray | Series | DataFrame,
func,
args: tuple,
kwargs: dict,
decorator: Callable,
axis: int | str,
):
"""
Apply `func` along the given axis using Numba.
"""
engine_kwargs: dict[str, bool] | None = (
decorator if isinstance(decorator, dict) else None
)

looper_args, looper_kwargs = prepare_function_arguments(
func,
args,
kwargs,
num_required_args=1,
)
# error: Argument 1 to "__call__" of "_lru_cache_wrapper" has
# incompatible type "Callable[..., Any] | str | list[Callable
# [..., Any] | str] | dict[Hashable,Callable[..., Any] | str |
# list[Callable[..., Any] | str]]"; expected "Hashable"
nb_looper = generate_apply_looper(
func,
**get_jit_arguments(engine_kwargs),
)
result = nb_looper(data, axis, *looper_args)
# If we made the result 2-D, squeeze it back to 1-D
return np.squeeze(result)


def frame_apply(
obj: DataFrame,
func: AggFuncType,
Expand Down Expand Up @@ -1094,23 +1148,15 @@ def wrapper(*args, **kwargs):
return wrapper

if engine == "numba":
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.

self.values,
self.func,
self.args,
self.kwargs,
num_required_args=1,
)
# error: Argument 1 to "__call__" of "_lru_cache_wrapper" has
# incompatible type "Callable[..., Any] | str | list[Callable
# [..., Any] | str] | dict[Hashable,Callable[..., Any] | str |
# list[Callable[..., Any] | str]]"; expected "Hashable"
nb_looper = generate_apply_looper(
self.func, # type: ignore[arg-type]
**get_jit_arguments(engine_kwargs),
engine_kwargs,
self.axis,
)
result = nb_looper(self.values, self.axis, *args)
# If we made the result 2-D, squeeze it back to 1-D
result = np.squeeze(result)
else:
result = np.apply_along_axis(
wrap_function(self.func),
Expand Down
Loading