ENH: Allow JIT compilation with an internal API

pandas/core/bodo_patched.py

@@ -0,0 +1,68 @@
+"""
+This file is here as an example, this code will live in the Numba and
+Bodo libraries.
+"""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+)
+
+import bodo
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from pandas._typing import (
+        AggFuncType,
+        Axis,
+    )
+
+
+def __pandas_udf__(
+    jit_decorator: Callable,
+    obj: pd.Series | pd.DataFrame,
+    method: Literal["apply", "map"],
+    func: AggFuncType,
+    axis: Axis,
+    raw: bool,
+    result_type: Literal["expand", "reduce", "broadcast"] | None,
+    args: tuple,
+    kwargs: dict[str, Any],
+    by_row: Literal[False, "compat"],
+):
+    if isinstance(obj, pd.DataFrame) and method == "apply":
+        if result_type is not None:
+            raise NotImplementedError(
+                "engine='bodo' not supported when result_type is not None"
+            )
+
+        if raw:
+            raise NotImplementedError("engine='bodo' not supported when raw=True")
+        if isinstance(func, str) and axis != 1:
+            raise NotImplementedError(
+                "engine='bodo' only supports axis=1 when func is the name of a "
+                "user-defined function"
+            )
+        if args or kwargs:
+            raise NotImplementedError(
+                "engine='bodo' not supported when args or kwargs are specified"
+            )
+
+        @jit_decorator
+        def jit_func(df, func, axis):
+            return df.apply(func, axis=axis)
+
+        return jit_func(obj, func, axis)
+    else:
+        raise NotImplementedError(
+            f"engine='bodo' not supported for {obj.__name__}.{method}"
+        )
+
+
+bodo.jit.__pandas_udf__ = __pandas_udf__

pandas/core/frame.py

@@ -10256,6 +10256,7 @@ def apply(
         by_row: Literal[False, "compat"] = "compat",
         engine: Literal["python", "numba"] = "python",
         engine_kwargs: dict[str, bool] | None = None,
+        jit: Callable | None = None,
         **kwargs,
     ):
         """
@@ -10345,6 +10346,15 @@ def apply(
             Pass keyword arguments to the engine.
             This is currently only used by the numba engine,
             see the documentation for the engine argument for more information.
+
+        jit : function, optional
+            Decorator to JIT compile the execution. The main available options are
+            ``numba.jit``, ``numba.njit`` or ``bodo.jit``. Parameters can be used in
+            the same way as the decorators, for example ``numba.jit(parallel=True)``.
+
+            Refer to the the [1]_ and [2]_ documentation to learn about limitations
+            on what code can be JIT compiled.
+
         **kwargs
             Additional keyword arguments to pass as keywords arguments to
             `func`.
@@ -10367,6 +10377,13 @@ def apply(
         behavior or errors and are not supported. See :ref:`gotchas.udf-mutation`
         for more details.
 
+        References
+        ----------
+        .. [1] `Numba documentation
+                <https://numba.readthedocs.io/en/stable/index.html>`_
+        .. [2] `Bodo documentation
+                <https://docs.bodo.ai/latest/>`/
+
         Examples
         --------
         >>> df = pd.DataFrame([[4, 9]] * 3, columns=["A", "B"])
@@ -10435,7 +10452,34 @@ def apply(
         0  1  2
         1  1  2
         2  1  2
-        """
+
+        Advanced users can speed up their code by using a Just-in-time (JIT) compiler
+        with ``apply``. The main JIT compilers available for pandas are Numba and Bodo.
+        In general, JIT compilation is only possible when the function passed to
+        ``apply`` has type stability (variables in the function do not change their
+        type during the execution).
+
+        >>> import bodo
+        >>> df.apply(lambda x: x.A + x.B, axis=1, jit=bodo.jit(parallel=True))
+
+        Note that JIT compilation is only recommended for functions that take a
+        significant amount of time to run. Fast functions are unlikely to run faster
+        with JIT compilation.
+        """
+        if hasattr(jit, "__pandas_udf__"):
+            return jit.__pandas_udf__(
+                jit_decorator=jit,
+                obj=self,
+                method="apply",
+                func=func,
+                args=args,
+                kwargs=kwargs,
+                axis=axis,
+                raw=raw,
+                result_type=result_type,
+                by_row=by_row,
+            )
+
         from pandas.core.apply import frame_apply
 
         op = frame_apply(
@@ -10567,9 +10611,11 @@ def _append(
 
             index = Index(
                 [other.name],
-                name=self.index.names
-                if isinstance(self.index, MultiIndex)
-                else self.index.name,
+                name=(
+                    self.index.names
+                    if isinstance(self.index, MultiIndex)
+                    else self.index.name
+                ),
             )
             row_df = other.to_frame().T
             # infer_objects is needed for