[PoC] Allow JIT compilation with an internal API

Author: datapythonista

pandas/core/bodo_patched.py

@@ -0,0 +1,57 @@
+"""
+This file is here as an example, this code will live in the Numba and
+Bodo libraries.
+"""
+from __future__ import annotations
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Literal, Any
+
+import pandas as pd
+import bodo
+
+if TYPE_CHECKING:
+    from pandas._typing import Axis, AggFuncType
+
+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.__class__.__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,12 @@ 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
+            Numba or Bodo 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 ``numba.jit(parallel=True)`` etc.
+
         **kwargs
             Additional keyword arguments to pass as keywords arguments to
             `func`.
@@ -10435,7 +10442,33 @@ 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,
+                axis=axis,
+                raw=raw,
+                result_type=result_type,
+                by_row=by_row,
+                args=args,
+                kwargs=kwargs)
+
         from pandas.core.apply import frame_apply
 
         op = frame_apply(