ENH: Adding engine parameter to Series.map

Author: datapythonista

doc/source/whatsnew/v3.0.0.rst

@@ -71,6 +71,7 @@ Other enhancements
 - :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`)
+- :meth:`Series.map` now accepts an ``engine`` parameter to allow execution with a third-party execution engine (:issue:`61125`)
 - :meth:`Series.str.get_dummies` now accepts a  ``dtype`` parameter to specify the dtype of the resulting DataFrame (:issue:`47872`)
 - :meth:`pandas.concat` will raise a ``ValueError`` when ``ignore_index=True`` and ``keys`` is not ``None`` (:issue:`59274`)
 - :py:class:`frozenset` elements in pandas objects are now natively printed (:issue:`60690`)

pandas/core/series.py

@@ -4322,6 +4322,7 @@ def map(
         self,
         arg: Callable | Mapping | Series,
         na_action: Literal["ignore"] | None = None,
+        engine: Callable | None = None,
         **kwargs,
     ) -> Series:
         """
@@ -4338,6 +4339,23 @@ def map(
         na_action : {None, 'ignore'}, default None
             If 'ignore', propagate NaN values, without passing them to the
             mapping correspondence.
+        engine : decorator, optional
+            Choose the execution engine to use. If not provided the function
+            will be executed by the regular Python interpreter.
+
+            Other options include JIT compilers such as Numba and Bodo, which in some
+            cases can speed up the execution. To use an executor you can provide
+            the decorators ``numba.jit``, ``numba.njit`` or ``bodo.jit``. You can
+            also provide the decorator with parameters, like ``numba.jit(nogit=True)``.
+
+            Not all functions can be executed with all execution engines. In general,
+            JIT compilers will require type stability in the function (no variable
+            should change data type during the execution). And not all pandas and
+            NumPy APIs are supported. Check the engine documentation [1]_ and [2]_
+            for limitations.
+
+            .. versionadded:: 3.0.0
+
         **kwargs
             Additional keyword arguments to pass as keywords arguments to
             `arg`.
@@ -4404,6 +4422,18 @@ def map(
         3  I am a rabbit
         dtype: object
         """
+        if engine is not None:
+            if not hasattr(engine, "__pandas_udf__"):
+                raise ValueError(f"Not a valid engine: {engine}")
+            return engine.__pandas_udf__.map(
+                data=self,
+                func=arg,
+                args=(),
+                kwargs=kwargs,
+                decorator=engine,
+                skip_na=na_action == "ignore",
+            ).__finalize__(self, method="map")
+
         if callable(arg):
             arg = functools.partial(arg, **kwargs)
         new_values = self._map_values(arg, na_action=na_action)

pandas/tests/apply/common.py

@@ -1,7 +1,66 @@
+import numpy as np
+
+from pandas import (
+    DataFrame,
+    Series,
+)
+from pandas.api.executors import BaseExecutionEngine
 from pandas.core.groupby.base import transformation_kernels
 
 # There is no Series.cumcount or DataFrame.cumcount
 series_transform_kernels = [
     x for x in sorted(transformation_kernels) if x != "cumcount"
 ]
 frame_transform_kernels = [x for x in sorted(transformation_kernels) if x != "cumcount"]
+
+
+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

pandas/tests/apply/test_frame_apply.py

@@ -17,63 +17,11 @@
     date_range,
 )
 import pandas._testing as tm
-from pandas.api.executors import BaseExecutionEngine
+from pandas.tests.apply.common import MockEngineDecorator
 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():
     """

pandas/tests/apply/test_series_apply.py

@@ -376,13 +376,13 @@ def test_demo():
 
 
 @pytest.mark.parametrize("func", [str, lambda x: str(x)])
-def test_apply_map_evaluate_lambdas_the_same(string_series, func, by_row):
+def test_apply_map_evaluate_lambdas_the_same(string_series, func, by_row, engine):
     # test that we are evaluating row-by-row first if by_row="compat"
     # else vectorized evaluation
     result = string_series.apply(func, by_row=by_row)
 
     if by_row:
-        expected = string_series.map(func)
+        expected = string_series.map(func, engine=engine)
         tm.assert_series_equal(result, expected)
     else:
         assert result == str(string_series)