Implemented NumbaExecutionEngine #61487

arthurlw · 2025-05-23T23:51:56Z

~~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 · 2025-05-26T16:42:29Z

pandas/core/apply.py

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


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.

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?

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 · 2025-05-26T16:45:47Z

doc/source/whatsnew/v3.0.0.rst

@@ -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`)


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 · 2025-06-03T15:14:12Z

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 · 2025-06-03T18:28:31Z

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 · 2025-06-16T03:13:58Z

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:

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

datapythonista · 2025-06-16T09:10:10Z

Correct

datapythonista · 2025-05-31T08:53:35Z

pandas/core/apply.py

@@ -13,6 +13,7 @@
    cast,
 )

+import numba


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 · 2025-05-31T09:08:12Z

pandas/core/apply.py

+            if not hasattr(numba.jit, "__pandas_udf__"):
+                numba.jit.__pandas_udf__ = NumbaExecutionEngine


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?

Yes this makes sense thank you

datapythonista · 2025-06-16T09:12:42Z

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

…dating empty or python string condition

…pdated engine checks

arthurlw · 2025-06-18T04:39:41Z

Hey @datapythonista, the current Numba implementation for apply when raw=false is written inside FrameApply e.g., apply_series_numba()), and relies on internal state (self.obj, self.axis, etc.).

As part of moving toward the new engine interface, should we explicitly rewrite this logic inside NumbaExecutionEngine, passing in the required values (like data, func, axis, etc.)?

datapythonista · 2025-06-18T07:25:16Z

As part of moving toward the new engine interface, should we explicitly rewrite this logic inside NumbaExecutionEngine, passing in the required values (like data, func, axis, etc.)?

Yes, the new interface already receives that information as parameters

…alidate_values_for_numba

arthurlw · 2025-06-20T07:39:58Z

Hi @datapythonista, I’ve finished the NumbaExecutionEngine implementation for apply and removed all instances of numba from FrameApply. All tests are passing. Let me know if there’s anything else you'd like me to follow up on!

datapythonista

Thanks @arthurlw for working on this. This is really cool stuff, great job. I added some comments that I think could help make the code a bit simpler, but just minor things, this looks great.

datapythonista · 2025-06-20T19:22:54Z

pandas/core/apply.py

+                "the 'numba' engine doesn't support lists of callables yet"
+            )
+
+        if isinstance(func, str):


No difference in practice, but I'd probably make this an elif also, looks clearer.

datapythonista · 2025-06-20T19:24:40Z

pandas/core/apply.py

+        """
+        Elementwise map for the Numba engine. Currently not supported.
+        """
+        raise NotImplementedError("Numba map is not implemented yet.")


This is the error when users write something like df.map(func, engine=numba.jit). I think it'll be easier to understand for users if the message is something like The Numba engine is not implemented for the map method yet.

datapythonista · 2025-06-20T19:32:54Z

pandas/core/apply.py

+
+        # check for data typing
+        if not isinstance(data, np.ndarray):
+            if len(data.columns) == 0 and len(data.index) == 0:


Why not just data.empty?

Good point, I referred to the logic in FrameApply, which doesn’t have the .empty property. Updated to use data.empty for clarity.

datapythonista · 2025-06-20T19:38:25Z

pandas/core/apply.py

+        # list[Callable[..., Any] | str]]"; expected "Hashable"
+        nb_looper = generate_apply_looper(
+            func,
+            **get_jit_arguments(engine_kwargs),


I think we can make this simpler if you change generate_apply_looper to receive the decorator instead. I checked and it doesn't seem like generate_apply_looper is used elsewhere, so the change should be quite straightforward.

What you are doing now is to extract the nogil... from the numba.jit decorator, and then inside of generate_apply_looper the decorator is created again with the nogil... The idea is to just pass decorator and use it there, so you can forget about engine_kwargs.

I also noticed that generate_numba_apply_func and apply_with_numba for when raw=false also re-created the JIT decorator manually. Will update these functions in the next commit

datapythonista · 2025-06-20T19:40:04Z

pandas/core/apply.py

+        # 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(


I personally wouldn't abbreviate to nb, it's not super clear imho. Just calling this numba_looper seems better. nb is often used for notebook, and probably other things.

datapythonista · 2025-06-20T19:46:51Z

pandas/core/apply.py

+        if not data.index.is_unique or not data.columns.is_unique:
+            raise NotImplementedError(
+                "The index/columns must be unique when raw=False and engine='numba'"
+            )


Personal preference, so feel free to disagree, but I'd personally keep those checks with the rest in apply. You'll have to check whether data is a numpy array or a pandas object there, but I think it'll make the code easier to understand.

Hmm yeah I think I agree with you on this. Will update in the next commit

datapythonista · 2025-06-20T19:48:44Z

pandas/core/apply.py

+        return DataFrame() if isinstance(data, DataFrame) else Series()
+
+    @staticmethod
+    def validate_values_for_numba(obj: Series | DataFrame) -> None:


Same for this, I'd personally validate everything first (maybe all the raise of not supported things can live in a function). But I wouldn't mix the validation with the execution when all is fine, unless it's needed.

datapythonista · 2025-06-20T19:56:27Z

pandas/core/frame.py

+            if engine_kwargs is not None:
+                numba_jit = numba.jit(**engine_kwargs)
+            else:
+                numba_jit = numba.jit()


You can be a bit more concise:

engine = numba.jit(**engine_kwargs or {}) engine.__pandas_udf__ = NumbaExecutionEngine

datapythonista · 2025-06-20T19:59:48Z

pandas/core/frame.py

-                engine = "python"
-
-            if engine not in ["python", "numba"]:
+            if engine not in ["python"] and engine is not None:


You can also do:

if engine not in ["python", None]:

datapythonista · 2025-06-20T20:00:44Z

pandas/core/frame.py

@@ -10632,12 +10641,13 @@ def apply(
                raw=raw,
                result_type=result_type,
                by_row=by_row,
-                engine=engine,
+                engine="python",


If engine used in frame_apply? I think it shouldn't.

As far as I know, the engine parameter is not used in FrameApply anymore, so I updated the call to use engine="python" directly.

arthurlw added 5 commits May 23, 2025 22:47

Implemented NumbaExecutionEngine

aa42037

whatsnew

db9f3b0

precommit

4cb240d

Match function arguments

97d9063

Fix CI

69e0e35

datapythonista reviewed May 26, 2025

View reviewed changes

arthurlw added 2 commits May 28, 2025 16:41

updated whatsnew

7365079

Updated conditions and delegate method to numba.jit

c605857

datapythonista added the Apply Apply, Aggregate, Transform, Map label May 31, 2025

Added try and except to catch ImportError

24a0615

arthurlw and others added 2 commits June 10, 2025 13:28

Use import_optional_dependency to load Numba

b7a2ecb

Merge branch 'main' into numba_execution_engine

545db65

datapythonista reviewed Jun 16, 2025

View reviewed changes

arthurlw added 4 commits June 17, 2025 14:38

Updated engine handling: normalizing numba to a fake decorator and up…

6f4fb50

…dating empty or python string condition

Added check for empty engine_kwargs

221cf7c

Moved checks from Apply.apply to NumbaExecutionEngine.apply

ed8dc7f

Fixed CI, removed unused numba checks, updated raw=false condition, u…

65b9d32

…pdated engine checks

arthurlw added 5 commits June 19, 2025 11:28

Implement and refactor raw=False logic into NumbaExecutionEngine.apply

2703f86

Fix CI, update validate_values_for_numba params

347463e

Adjust error messages

77eb146

Remove Numba-specific logic from FrameApply, added Series import to v…

90f264f

…alidate_values_for_numba

pre-commit

f8f1166

datapythonista reviewed Jun 20, 2025

View reviewed changes

arthurlw added 3 commits June 21, 2025 19:30

Updated with reviewer suggestions and added axis normalizing

bc2939b

Updated executor to accept decorator

176753b

Fix CI and pre-commit

cf3e392

		if not hasattr(numba.jit, "__pandas_udf__"):
		numba.jit.__pandas_udf__ = NumbaExecutionEngine

Uh oh!

Implemented NumbaExecutionEngine #61487

Are you sure you want to change the base?

Implemented NumbaExecutionEngine #61487

Conversation

arthurlw commented May 23, 2025

Uh oh!

datapythonista left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

arthurlw commented Jun 3, 2025

Uh oh!

datapythonista commented Jun 3, 2025

Uh oh!

arthurlw commented Jun 16, 2025

Uh oh!

datapythonista commented Jun 16, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

datapythonista commented Jun 16, 2025

Uh oh!

arthurlw commented Jun 18, 2025

Uh oh!

datapythonista commented Jun 18, 2025

Uh oh!

arthurlw commented Jun 20, 2025

Uh oh!

datapythonista left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!