Updated with reviewer suggestions

Author: arthurlw

pandas/core/apply.py

@@ -216,11 +216,8 @@ def apply(
 
         NumbaExecutionEngine.check_numba_support(func)
 
-        # normalize axis values
-        if axis in (0, "index"):
-            axis = 0
-        else:
-            axis = 1
+        if not isinstance(data, np.ndarray):
+            axis = data._get_axis_number(cast(Axis, axis))
 
         # check for data typing
         if not isinstance(data, np.ndarray):
@@ -230,7 +227,7 @@ def apply(
             engine_kwargs = extract_numba_options(decorator)
 
             NumbaExecutionEngine.validate_values_for_numba_raw_false(
-                data, get_jit_arguments(engine_kwargs)
+                data, **get_jit_arguments(engine_kwargs)
             )
 
             return NumbaExecutionEngine.apply_raw_false(
@@ -288,9 +285,29 @@ def apply_raw_false(
             Series,
         )
 
-        results = NumbaExecutionEngine.apply_with_numba(
-            data, func, args, kwargs, decorator, axis
+        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)
+
+        from pandas.core._numba.extensions import set_numba_data
+
+        # Convert from numba dict to regular dict
+        # Our isinstance checks in the df constructor don't pass for numbas typed dict
+
+        if axis == 0:
+            col_names_index = data.index
+            result_index = data.columns
+        else:
+            col_names_index = data.columns
+            result_index = data.index
+
+        with (
+            set_numba_data(result_index) as index,
+            set_numba_data(col_names_index) as columns,
+        ):
+            results = dict(nb_func(data.values, columns, index, *args))
 
         if results:
             sample = next(iter(results.values()))
@@ -306,11 +323,14 @@ def apply_raw_false(
 
     @staticmethod
     def validate_values_for_numba_raw_false(
-        data: Series | DataFrame, engine_kwargs: dict[str, bool]
+        data: Series | DataFrame,
+        nopython: bool | None = None,
+        nogil: bool | None = None,
+        parallel: bool | None = None,
     ) -> None:
         from pandas import Series
 
-        if engine_kwargs.get("parallel", False):
+        if parallel:
             raise NotImplementedError(
                 "Parallel apply is not supported when raw=False and engine='numba'"
             )
@@ -376,34 +396,6 @@ 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]:
-        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)
-
-        from pandas.core._numba.extensions import set_numba_data
-
-        # Convert from numba dict to regular dict
-        # Our isinstance checks in the df constructor don't pass for numbas typed dict
-
-        if axis == 0 or axis == "index":
-            col_names_index = data.index
-            result_index = data.columns
-        else:
-            col_names_index = data.columns
-            result_index = data.index
-
-        with (
-            set_numba_data(result_index) as index,
-            set_numba_data(col_names_index) as columns,
-        ):
-            res = dict(nb_func(data.values, columns, index, *args))
-
-        return res
-
 
 def frame_apply(
     obj: DataFrame,

pandas/core/util/numba_.py

@@ -152,20 +152,37 @@ def prepare_function_arguments(
 
 def extract_numba_options(decorator: Callable) -> dict:
     """
-    Extract targetoptions from a numba.jit decorator
+    Extract the `targetoptions` dictionary from a numba.jit decorator.
+
+    The `targetoptions` attribute stores the keyword arguments
+    passed to the `@numba.jit` decorator when it is created.
+
+    This function returns a dictionary with the following keys,
+    if present in the decorator:
+        - nopython
+        - nogil
+        - parallel
+
+    Parameters
+    ----------
+    decorator : Callable
+        A numba.jit decorated function or a numba dispatcher object.
+
+    Returns
+    -------
+    dict
+        A dictionary with the extracted numba compilation options.
     """
-    try:
-        closure = decorator.__closure__
-        if closure is None:
-            return {}
-        freevars = decorator.__code__.co_freevars
-        if "targetoptions" not in freevars:
-            return {}
-        idx = freevars.index("targetoptions")
-        cell = closure[idx]
-        targetoptions = cell.cell_contents
-        if isinstance(targetoptions, dict):
-            return targetoptions
+    closure = decorator.__closure__
+    if closure is None:
         return {}
-    except Exception:
+    freevars = decorator.__code__.co_freevars
+    if "targetoptions" not in freevars:
         return {}
+    idx = freevars.index("targetoptions")
+    cell = closure[idx]
+    targetoptions = cell.cell_contents
+    if isinstance(targetoptions, dict):
+        relevant_keys = {"nopython", "nogil", "parallel"}
+        return {k: v for k, v in targetoptions.items() if k in relevant_keys}
+    return {}

pandas/tests/apply/test_numba.py

@@ -141,3 +141,7 @@ def test_extract_numba_options_from_user_decorated_function(jit_args):
     extracted = extract_numba_options(numba.jit(**jit_args))
     for k, v in jit_args.items():
         assert extracted.get(k) == v
+
+    extracted = extract_numba_options(numba.njit(**jit_args))
+    for k, v in jit_args.items():
+        assert extracted.get(k) == v