Implemented NumbaExecutionEngine

Author: arthurlw

pandas/core/apply.py

@@ -45,9 +45,9 @@
     ABCSeries,
 )
 
-from pandas.core._numba.executor import generate_apply_looper
 import pandas.core.common as com
 from pandas.core.construction import ensure_wrapped_if_datetimelike
+from pandas.core._numba.executor import generate_apply_looper
 from pandas.core.util.numba_ import (
     get_jit_arguments,
     prepare_function_arguments,
@@ -178,6 +178,57 @@ 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,
+        engine_kwargs: dict | 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,
+        engine_kwargs: dict | None,
+        axis: int | str,
+    ):
+        """
+        Apply `func` along the given axis using Numba.
+        """
+
+        looper_args, looper_kwargs = prepare_function_arguments(
+            func,   # type: ignore[arg-type]
+            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,   # type: ignore[arg-type]
+            **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,
@@ -1094,23 +1145,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(
+                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),