Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
36 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
b7a2ecb
Use import_optional_dependency to load Numba
arthurlw Jun 10, 2025
545db65
Merge branch 'main' into numba_execution_engine
arthurlw Jun 10, 2025
6f4fb50
Updated engine handling: normalizing numba to a fake decorator and up…
arthurlw Jun 17, 2025
221cf7c
Added check for empty engine_kwargs
arthurlw Jun 17, 2025
ed8dc7f
Moved checks from Apply.apply to NumbaExecutionEngine.apply
arthurlw Jun 17, 2025
65b9d32
Fixed CI, removed unused numba checks, updated raw=false condition, u…
arthurlw Jun 18, 2025
2703f86
Implement and refactor raw=False logic into NumbaExecutionEngine.apply
arthurlw Jun 19, 2025
347463e
Fix CI, update validate_values_for_numba params
arthurlw Jun 19, 2025
77eb146
Adjust error messages
arthurlw Jun 19, 2025
90f264f
Remove Numba-specific logic from FrameApply, added Series import to v…
arthurlw Jun 19, 2025
f8f1166
pre-commit
arthurlw Jun 19, 2025
bc2939b
Updated with reviewer suggestions and added axis normalizing
arthurlw Jun 21, 2025
176753b
Updated executor to accept decorator
arthurlw Jun 21, 2025
cf3e392
Fix CI and pre-commit
arthurlw Jun 21, 2025
a4bac18
Silence pyright warning for untyped decorator
arthurlw Jun 22, 2025
ca91e89
Revert elif to if
arthurlw Jun 22, 2025
e337cb8
Fix engine_kwargs handling in validate_values_for_numba_raw_false
arthurlw Jun 23, 2025
30703b9
Fix CI and removed references to engine and enginekwargs in FrameApply
arthurlw Jun 24, 2025
b39a8d1
Added test for axis with string input
arthurlw Jun 26, 2025
2b59eeb
pre-commit
arthurlw Jun 26, 2025
f59fb52
Revert test for numba apply with string axis due to docstring validat…
arthurlw Jun 27, 2025
984a008
Merge branch 'main' into numba_execution_engine
datapythonista Jul 2, 2025
b08c361
Updated engine_kwargs extraction logic and added test
arthurlw Jul 4, 2025
8943e0a
Merge
arthurlw Jul 4, 2025
f330473
Fix CI and pre-commit
arthurlw Jul 4, 2025
c05b1a7
Add return type annotation for extract_numba_options
arthurlw Jul 4, 2025
4b650d0
Updated with reviewer suggestions
arthurlw Jul 10, 2025
85e3dd3
Merge branch 'main' into numba_execution_engine
arthurlw Jul 24, 2025
File filter

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix CI and pre-commit
  • Loading branch information
arthurlw committed Jun 21, 2025
commit cf3e39288ca90c17c5c5875936c536fa5ce2b2ff
2 changes: 1 addition & 1 deletion pandas/core/_numba/executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ def generate_apply_looper(func, decorator: Callable):
numba = import_optional_dependency("numba")
nb_compat_func = jit_user_function(func)

@decorator # type: ignore
@decorator
def nb_looper(values, axis, *args):
# Operate on the first row/col in order to get
# the output shape
Expand Down
23 changes: 11 additions & 12 deletions pandas/core/apply.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,9 @@ def map(
"""
Elementwise map for the Numba engine. Currently not supported.
"""
raise NotImplementedError("The Numba engine is not implemented for the map method yet.")
raise NotImplementedError(
"The Numba engine is not implemented for the map method yet."
)

@staticmethod
def apply(
Expand Down Expand Up @@ -222,9 +224,8 @@ def apply(
if data.empty:
return data.copy() # mimic apply_empty_result()
NumbaExecutionEngine.validate_values_for_numba_raw_false(
data,
decorator if isinstance(decorator, dict) else {}
)
data, decorator if isinstance(decorator, dict) else {}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It'd be good to make the second argument a keyword argument, it's not clear from this call what the method is expecting.

And I'm not sure when decorator can be a dictionary? Afaik decorator is the engine, I don't understand the logic here. Is it possible that this always pass an empty dictionary? Feels like if you want to check if numba.jit has been initialized with parallel=True or others you should check an attribute of it.

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 pointing this out, I overlooked it earlier. Will update in the next commit

)

return NumbaExecutionEngine.apply_raw_false(
data, func, args, kwargs, decorator, axis
Expand Down Expand Up @@ -298,7 +299,9 @@ def apply_raw_false(
return DataFrame() if isinstance(data, DataFrame) else Series()

@staticmethod
def validate_values_for_numba_raw_false(data: Series | DataFrame, engine_kwargs: dict[str, bool]) -> None:
def validate_values_for_numba_raw_false(
data: Series | DataFrame, engine_kwargs: dict[str, bool]
) -> None:
from pandas import Series

if engine_kwargs.get("parallel", False):
Expand Down Expand Up @@ -344,7 +347,7 @@ def generate_numba_apply_func(

jitted_udf = numba.extending.register_jitable(func)

@decorator # type: ignore
@decorator
def numba_func(values, col_names_index, index, *args):
results = {}
for i in range(values.shape[1 - axis]):
Expand All @@ -368,16 +371,12 @@ def numba_func(values, col_names_index, index, *args):
return numba_func

@staticmethod
def apply_with_numba(
data, func, args, kwargs, decorator, axis=0
) -> dict[int, Any]:
def apply_with_numba(data, func, args, kwargs, decorator, axis=0) -> dict[int, Any]:
func = cast(Callable, func)
args, kwargs = prepare_function_arguments(
func, args, kwargs, num_required_args=1
)
nb_func = NumbaExecutionEngine.generate_numba_apply_func(
func, axis, decorator
)
nb_func = NumbaExecutionEngine.generate_numba_apply_func(func, axis, decorator)

from pandas.core._numba.extensions import set_numba_data

Expand Down
7 changes: 4 additions & 3 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -10621,10 +10621,11 @@ def apply(
"""
if engine == "numba":
numba = import_optional_dependency("numba")
engine = numba.jit(**engine_kwargs or {})
engine.__pandas_udf__ = NumbaExecutionEngine
numba_jit = numba.jit(**engine_kwargs or {})
numba_jit.__pandas_udf__ = NumbaExecutionEngine
engine = numba_jit

if engine is None or isinstance(engine, str):
elif engine is None or isinstance(engine, str):
from pandas.core.apply import frame_apply

if engine not in ["python", None]:
Expand Down
Loading