Added tests, fixed some bugs and added a release note

Author: datapythonista

doc/source/whatsnew/v3.0.0.rst

@@ -65,6 +65,7 @@ Other enhancements
 - :class:`Rolling` and :class:`Expanding` now support aggregations ``first`` and ``last`` (:issue:`33155`)
 - :func:`read_parquet` accepts ``to_pandas_kwargs`` which are forwarded to :meth:`pyarrow.Table.to_pandas` which enables passing additional keywords to customize the conversion to pandas, such as ``maps_as_pydicts`` to read the Parquet map data type as python dictionaries (:issue:`56842`)
 - :meth:`.DataFrameGroupBy.transform`, :meth:`.SeriesGroupBy.transform`, :meth:`.DataFrameGroupBy.agg`, :meth:`.SeriesGroupBy.agg`, :meth:`.SeriesGroupBy.apply`, :meth:`.DataFrameGroupBy.apply` now support ``kurt`` (:issue:`40139`)
+- :meth:`DataFrame.apply` supports using third-party execution engines like the Bodo.ai JIT compiler (:issue:`60668`)
 - :meth:`DataFrameGroupBy.transform`, :meth:`SeriesGroupBy.transform`, :meth:`DataFrameGroupBy.agg`, :meth:`SeriesGroupBy.agg`, :meth:`RollingGroupby.apply`, :meth:`ExpandingGroupby.apply`, :meth:`Rolling.apply`, :meth:`Expanding.apply`, :meth:`DataFrame.apply` with ``engine="numba"`` now supports positional arguments passed as kwargs (:issue:`58995`)
 - :meth:`Rolling.agg`, :meth:`Expanding.agg` and :meth:`ExponentialMovingWindow.agg` now accept :class:`NamedAgg` aggregations through ``**kwargs`` (:issue:`28333`)
 - :meth:`Series.map` can now accept kwargs to pass on to func (:issue:`59814`)

pandas/core/apply.py

@@ -120,7 +120,7 @@ def map(
             For JIT compilers and other engines that need to decorate the
             function ``func``, this is the decorator to use. While the
             executor may already know which is the decorator to use, this
-            is useful as for a single executor the user can specify for a
+            is useful as for a single executor the user can specify for
             example ``numba.jit`` or ``numba.njit(nogil=True)``, and this
             decorator parameter will contain the exact decortor from the
             executor the user wants to use.
@@ -163,7 +163,7 @@ def apply(
             For JIT compilers and other engines that need to decorate the
             function ``func``, this is the decorator to use. While the
             executor may already know which is the decorator to use, this
-            is useful as for a single executor the user can specify for a
+            is useful as for a single executor the user can specify for
             example ``numba.jit`` or ``numba.njit(nogil=True)``, and this
             decorator parameter will contain the exact decortor from the
             executor the user wants to use.

pandas/core/frame.py

@@ -10459,6 +10459,9 @@ def apply(
             if engine is None:
                 engine = "python"
 
+            if engine not in ["python", "numba"]:
+                raise ValueError(f"Unknown engine '{engine}'")
+
             op = frame_apply(
                 self,
                 func=func,
@@ -10478,21 +10481,56 @@ def apply(
                     f"{result_type=} only implemented for the default engine"
                 )
 
+            agg_axis = self._get_agg_axis(axis)
+
+            # one axis is empty
+            if not all(self.shape):
+                try:
+                    if axis == 0:
+                        r = func(Series([], dtype=np.float64), *args, **kwargs)
+                    else:
+                        r = func(
+                            Series(index=self.columns, dtype=np.float64),
+                            *args,
+                            **kwargs,
+                        )
+                except Exception:
+                    pass
+                else:
+                    if not isinstance(r, Series):
+                        if len(agg_axis):
+                            r = func(Series([], dtype=np.float64), *args, **kwargs)
+                        else:
+                            r = np.nan
+
+                        return self._constructor_sliced(r, index=agg_axis)
+                return self.copy()
+
             data = self
             if raw:
                 # This will upcast the whole DataFrame to the same type,
                 # and likely result in an object 2D array.
                 # We should probably pass a list of 1D arrays instead, at
                 # lest for ``axis=0``
                 data = data.values
-            return engine.__pandas_udf__.apply(
+            result = engine.__pandas_udf__.apply(
                 data=data,
                 func=func,
                 args=args,
                 kwargs=kwargs,
                 decorator=engine,
                 axis=axis,
             )
+            if raw:
+                if result.ndim == 2:
+                    return self._constructor(
+                        result, index=self.index, columns=self.columns
+                    )
+                else:
+                    return self._constructor_sliced(result, index=agg_axis)
+            return result
+        else:
+            raise ValueError(f"Unknown engine {engine}")
 
     def map(
         self, func: PythonFuncType, na_action: Literal["ignore"] | None = None, **kwargs

pandas/tests/apply/test_frame_apply.py

@@ -17,10 +17,63 @@
     date_range,
 )
 import pandas._testing as tm
+from pandas.api.executors import BaseExecutionEngine
 from pandas.tests.frame.common import zip_frames
 from pandas.util.version import Version
 
 
+class MockExecutionEngine(BaseExecutionEngine):
+    """
+    Execution Engine to test if the execution engine interface receives and
+    uses all parameters provided by the user.
+
+    Making this engine work as the default Python engine by calling it, no extra
+    functionality is implemented here.
+
+    When testing, this will be called when this engine is provided, and then the
+    same pandas.map and pandas.apply function will be called, but without engine,
+    executing the default behavior from the python engine.
+    """
+
+    def map(data, func, args, kwargs, decorator, skip_na):
+        kwargs_to_pass = kwargs if isinstance(data, DataFrame) else {}
+        return data.map(
+            func, action_na="ignore" if skip_na else False, **kwargs_to_pass
+        )
+
+    def apply(data, func, args, kwargs, decorator, axis):
+        if isinstance(data, Series):
+            return data.apply(func, convert_dtype=True, args=args, by_row=False)
+        elif isinstance(data, DataFrame):
+            return data.apply(
+                func,
+                axis=axis,
+                raw=False,
+                result_type=None,
+                args=args,
+                by_row="compat",
+                **kwargs,
+            )
+        else:
+            assert isinstance(data, np.ndarray)
+
+            def wrap_function(func):
+                # https://github.com/numpy/numpy/issues/8352
+                def wrapper(*args, **kwargs):
+                    result = func(*args, **kwargs)
+                    if isinstance(result, str):
+                        result = np.array(result, dtype=object)
+                    return result
+
+                return wrapper
+
+            return np.apply_along_axis(wrap_function(func), axis, data, *args, **kwargs)
+
+
+class MockEngineDecorator:
+    __pandas_udf__ = MockExecutionEngine
+
+
 @pytest.fixture
 def int_frame_const_col():
     """
@@ -35,7 +88,13 @@ def int_frame_const_col():
     return df
 
 
-@pytest.fixture(params=["python", pytest.param("numba", marks=pytest.mark.single_cpu)])
+@pytest.fixture(
+    params=[
+        "python",
+        pytest.param("numba", marks=pytest.mark.single_cpu),
+        MockEngineDecorator,
+    ]
+)
 def engine(request):
     if request.param == "numba":
         pytest.importorskip("numba")
@@ -1079,12 +1138,21 @@ def test_result_type_broadcast(int_frame_const_col, request, engine):
         mark = pytest.mark.xfail(reason="numba engine doesn't support list return")
         request.node.add_marker(mark)
     df = int_frame_const_col
-    # broadcast result
-    result = df.apply(
-        lambda x: [1, 2, 3], axis=1, result_type="broadcast", engine=engine
-    )
-    expected = df.copy()
-    tm.assert_frame_equal(result, expected)
+    if engine is MockEngineDecorator:
+        with pytest.raises(
+            NotImplementedError,
+            match="result_type='broadcast' only implemented for the default engine",
+        ):
+            df.apply(
+                lambda x: [1, 2, 3], axis=1, result_type="broadcast", engine=engine
+            )
+    else:
+        # broadcast result
+        result = df.apply(
+            lambda x: [1, 2, 3], axis=1, result_type="broadcast", engine=engine
+        )
+        expected = df.copy()
+        tm.assert_frame_equal(result, expected)
 
 
 def test_result_type_broadcast_series_func(int_frame_const_col, engine, request):
@@ -1097,14 +1165,27 @@ def test_result_type_broadcast_series_func(int_frame_const_col, engine, request)
         request.node.add_marker(mark)
     df = int_frame_const_col
     columns = ["other", "col", "names"]
-    result = df.apply(
-        lambda x: Series([1, 2, 3], index=columns),
-        axis=1,
-        result_type="broadcast",
-        engine=engine,
-    )
-    expected = df.copy()
-    tm.assert_frame_equal(result, expected)
+
+    if engine is MockEngineDecorator:
+        with pytest.raises(
+            NotImplementedError,
+            match="result_type='broadcast' only implemented for the default engine",
+        ):
+            df.apply(
+                lambda x: Series([1, 2, 3], index=columns),
+                axis=1,
+                result_type="broadcast",
+                engine=engine,
+            )
+    else:
+        result = df.apply(
+            lambda x: Series([1, 2, 3], index=columns),
+            axis=1,
+            result_type="broadcast",
+            engine=engine,
+        )
+        expected = df.copy()
+        tm.assert_frame_equal(result, expected)
 
 
 def test_result_type_series_result(int_frame_const_col, engine, request):
@@ -1791,3 +1872,9 @@ def test_agg_dist_like_and_nonunique_columns():
     result = df.agg({"A": "count"})
     expected = df["A"].count()
     tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize("engine_name", ["unknown", 25])
+def test_wrong_engine(engine_name):
+    with pytest.raises(ValueError, match="Unknown engine "):
+        DataFrame().apply(lambda x: x, engine=engine_name)