[BSE-4737] Implement interface for using bodo as engine in Pandas UDFs.

bodo/decorators.py

@@ -16,6 +16,7 @@
 from numba.extending import lower_builtin, models, register_model
 
 import bodo
+from bodo.pandas_compat import bodo_pandas_udf_execution_engine
 
 # Add Bodo's options to Numba's allowed options/flags
 numba.core.cpu.CPUTargetOptions.all_args_distributed_block = _mapping(
@@ -350,6 +351,8 @@ def return_wrapped_fn(py_func):
             py_func = signature_or_function
             return return_wrapped_fn(py_func)
 
+        return_wrapped_fn.__pandas_udf__ = bodo_pandas_udf_execution_engine
+
         return return_wrapped_fn
 
     elif "propagate_env" in options:
@@ -360,6 +363,9 @@ def return_wrapped_fn(py_func):
     return _jit(signature_or_function, pipeline_class, **options)
 
 
+jit.__pandas_udf__ = bodo_pandas_udf_execution_engine
+
+
 def _jit(signature_or_function=None, pipeline_class=None, **options):
     _init_extensions()
 

bodo/pandas_compat.py

@@ -244,3 +244,53 @@ def get_conversion_factor_to_ns(in_reso: str) -> int:
     else:
         raise ValueError(f"Unsupported resolution {in_reso}")
     return factor * value
+
+
+# Class responsible for executing UDFs using Bodo as the engine in
+# newer version of pandas.
+# https://github.com/pandas-dev/pandas/pull/61032
+bodo_pandas_udf_execution_engine = None
+
+if pandas_version >= (3, 0):
+    from collections.abc import Callable
+    from typing import Any
+
+    from pandas._typing import AggFuncType, Axis
+    from pandas.core.apply import BaseExecutionEngine
+
+    class BodoExecutionEngine(BaseExecutionEngine):
+        @staticmethod
+        def map(
+            data: pd.Series | pd.DataFrame | np.ndarray,
+            func: AggFuncType,
+            args: tuple,
+            kwargs: dict[str, Any],
+            decorator: Callable | None,
+            skip_na: bool,
+        ):
+            raise NotImplementedError("TODO: map")
+
+        @staticmethod
+        def apply(
+            data: pd.Series | pd.DataFrame | np.ndarray,
+            func: AggFuncType,
+            args: tuple,
+            kwargs: dict[str, Any],
+            decorator: Callable,
+            axis: Axis,
+        ):
+            # raw = True converts data to ndarray first
+            if isinstance(data, np.ndarray):
+                raise ValueError(
+                    "Bodo engine does not support the raw=True in DataFrame.apply."
+                )
+
+            jitted_func = decorator(func)
+
+            @decorator
+            def apply_func(data):
+                return data.apply(jitted_func, axis=axis, args=args)
+
+            return apply_func(data)
+
+    bodo_pandas_udf_execution_engine = BodoExecutionEngine

bodo/tests/test_pandas_udf_compat.py

@@ -0,0 +1,32 @@
+"""
+This file tests the Bodo implementation of the Pandas UDF interface.
+See https://github.com/pandas-dev/pandas/pull/61032 for more details.
+
+This feature is only availible on newer versions of Pandas (>=3.0)
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import bodo
+from bodo.pandas_compat import pandas_version
+from bodo.tests.utils import _test_equal, pytest_spawn_mode
+
+pytestmark = [
+    pytest.mark.skipif(
+        pandas_version < (3, 0), reason="Third-party UDF engines requires Pandas >= 3.0"
+    )
+] + pytest_spawn_mode
+
+
+def test_basic_apply_example():
+    """Simplest test to check Pandas UDF apply hook is set up properly"""
+
+    df = pd.DataFrame({"A": np.arange(30)})
+
+    bodo_result = df.apply(lambda x: x.A, axis=1, engine=bodo.jit)
+
+    pandas_result = df.apply(lambda x: x.A, axis=1)
+
+    _test_equal(bodo_result, pandas_result, check_pandas_types=False)

pixi.toml

@@ -70,7 +70,6 @@ pip = "*"
 # Core Python Deps
 numba = "==0.61.0"
 numpy = ">=1.24,<2.2"
-pandas = ">=2.2,<2.3"
 pyarrow = { version = "==19.0.0", channel = "conda-forge" }
 fsspec = ">=2021.09"
 requests = "*"

pyproject.toml

@@ -37,7 +37,7 @@ classifiers = [
 dependencies = [
     "numba==0.61.0",
     "pyarrow==19.0.0",
-    "pandas>=2.2,<2.3",
+    # "pandas>=2.2,<2.3",
     "numpy>=1.24,<2.2",
     # fsspec >= 2021.09 because it includes Arrow filesystem wrappers (useful for fs.glob() for example)
     "fsspec>=2021.09",