ENH: numba apply supports positional arguments passed as **kwargs

doc/source/whatsnew/v3.0.0.rst

@@ -62,6 +62,8 @@ Other enhancements
 - Support passing a :class:`Iterable[Hashable]` input to :meth:`DataFrame.drop_duplicates` (:issue:`59237`)
 - Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
+- numba apply now supports positional arguments passed as kwargs (:issue:`58995`)
+
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_300.notable_bug_fixes:

pandas/core/apply.py

@@ -994,14 +994,15 @@ def wrapper(*args, **kwargs):
                 self.func,  # type: ignore[arg-type]
                 self.args,
                 self.kwargs,
+                1,
             )
             # error: Argument 1 to "__call__" of "_lru_cache_wrapper" has
             # incompatible type "Callable[..., Any] | str | list[Callable
             # [..., Any] | str] | dict[Hashable,Callable[..., Any] | str |
             # list[Callable[..., Any] | str]]"; expected "Hashable"
             nb_looper = generate_apply_looper(
                 self.func,  # type: ignore[arg-type]
-                **get_jit_arguments(engine_kwargs, kwargs),
+                **get_jit_arguments(engine_kwargs),
             )
             result = nb_looper(self.values, self.axis, *args)
             # If we made the result 2-D, squeeze it back to 1-D
@@ -1158,9 +1159,9 @@ def numba_func(values, col_names, df_index, *args):
 
     def apply_with_numba(self) -> dict[int, Any]:
         func = cast(Callable, self.func)
-        args, kwargs = prepare_function_arguments(func, self.args, self.kwargs)
+        args, kwargs = prepare_function_arguments(func, self.args, self.kwargs, 1)
         nb_func = self.generate_numba_apply_func(
-            func, **get_jit_arguments(self.engine_kwargs, kwargs)
+            func, **get_jit_arguments(self.engine_kwargs)
         )
         from pandas.core._numba.extensions import set_numba_data
 
@@ -1298,9 +1299,9 @@ def numba_func(values, col_names_index, index, *args):
 
     def apply_with_numba(self) -> dict[int, Any]:
         func = cast(Callable, self.func)
-        args, kwargs = prepare_function_arguments(func, self.args, self.kwargs)
+        args, kwargs = prepare_function_arguments(func, self.args, self.kwargs, 1)
         nb_func = self.generate_numba_apply_func(
-            func, **get_jit_arguments(self.engine_kwargs, kwargs)
+            func, **get_jit_arguments(self.engine_kwargs)
         )
 
         from pandas.core._numba.extensions import set_numba_data

pandas/core/groupby/groupby.py

@@ -136,6 +136,7 @@ class providing the base-class of operations.
 from pandas.core.util.numba_ import (
     get_jit_arguments,
     maybe_use_numba,
+    prepare_function_arguments,
 )
 
 if TYPE_CHECKING:
@@ -1275,8 +1276,9 @@ def _transform_with_numba(self, func, *args, engine_kwargs=None, **kwargs):
 
         starts, ends, sorted_index, sorted_data = self._numba_prep(df)
         numba_.validate_udf(func)
+        args, kwargs = prepare_function_arguments(func, args, kwargs, 2)
         numba_transform_func = numba_.generate_numba_transform_func(
-            func, **get_jit_arguments(engine_kwargs, kwargs)
+            func, **get_jit_arguments(engine_kwargs)
         )
         result = numba_transform_func(
             sorted_data,
@@ -1311,8 +1313,9 @@ def _aggregate_with_numba(self, func, *args, engine_kwargs=None, **kwargs):
 
         starts, ends, sorted_index, sorted_data = self._numba_prep(df)
         numba_.validate_udf(func)
+        args, kwargs = prepare_function_arguments(func, args, kwargs, 2)
         numba_agg_func = numba_.generate_numba_agg_func(
-            func, **get_jit_arguments(engine_kwargs, kwargs)
+            func, **get_jit_arguments(engine_kwargs)
         )
         result = numba_agg_func(
             sorted_data,

pandas/core/util/numba_.py

@@ -29,18 +29,14 @@ def set_use_numba(enable: bool = False) -> None:
     GLOBAL_USE_NUMBA = enable
 
 
-def get_jit_arguments(
-    engine_kwargs: dict[str, bool] | None = None, kwargs: dict | None = None
-) -> dict[str, bool]:
+def get_jit_arguments(engine_kwargs: dict[str, bool] | None = None) -> dict[str, bool]:
     """
     Return arguments to pass to numba.JIT, falling back on pandas default JIT settings.
 
     Parameters
     ----------
     engine_kwargs : dict, default None
         user passed keyword arguments for numba.JIT
-    kwargs : dict, default None
-        user passed keyword arguments to pass into the JITed function
 
     Returns
     -------
@@ -55,16 +51,6 @@ def get_jit_arguments(
         engine_kwargs = {}
 
     nopython = engine_kwargs.get("nopython", True)
-    if kwargs:
-        # Note: in case numba supports keyword-only arguments in
-        # a future version, we should remove this check. But this
-        # seems unlikely to happen soon.
-
-        raise NumbaUtilError(
-            "numba does not support keyword-only arguments"
-            "https://github.com/numba/numba/issues/2916, "
-            "https://github.com/numba/numba/issues/6846"
-        )
     nogil = engine_kwargs.get("nogil", False)
     parallel = engine_kwargs.get("parallel", False)
     return {"nopython": nopython, "nogil": nogil, "parallel": parallel}
@@ -109,7 +95,7 @@ def jit_user_function(func: Callable) -> Callable:
 
 
 def prepare_function_arguments(
-    func: Callable, args: tuple, kwargs: dict
+    func: Callable, args: tuple, kwargs: dict, num_required_args: int
 ) -> tuple[tuple, dict]:
     """
     Prepare arguments for jitted function. As numba functions do not support kwargs,
@@ -123,6 +109,8 @@ def prepare_function_arguments(
         user input positional arguments
     kwargs : dict
         user input keyword arguments
+    num_required_args : int
+        the number of required leading positional arguments for udf.
 
     Returns
     -------
@@ -133,17 +121,26 @@ def prepare_function_arguments(
     if not kwargs:
         return args, kwargs
 
-    # the udf should have this pattern: def udf(value, *args, **kwargs):...
+    # the udf should have this pattern: def udf(arg1, arg2, ..., *args, **kwargs):...
     signature = inspect.signature(func)
-    arguments = signature.bind(_sentinel, *args, **kwargs)
+    arguments = signature.bind(*[_sentinel] * num_required_args, *args, **kwargs)
     arguments.apply_defaults()
     # Ref: https://peps.python.org/pep-0362/
     # Arguments which could be passed as part of either *args or **kwargs
     # will be included only in the BoundArguments.args attribute.
     args = arguments.args
     kwargs = arguments.kwargs
 
-    assert args[0] is _sentinel
-    args = args[1:]
+    if kwargs:
+        # Note: in case numba supports keyword-only arguments in
+        # a future version, we should remove this check. But this
+        # seems unlikely to happen soon.
+
+        raise NumbaUtilError(
+            "numba does not support keyword-only arguments"
+            "https://github.com/numba/numba/issues/2916, "
+            "https://github.com/numba/numba/issues/6846"
+        )
 
+    args = args[num_required_args:]
     return args, kwargs

pandas/core/window/rolling.py

@@ -65,6 +65,7 @@
 from pandas.core.util.numba_ import (
     get_jit_arguments,
     maybe_use_numba,
+    prepare_function_arguments,
 )
 from pandas.core.window.common import (
     flex_binary_moment,
@@ -1472,14 +1473,14 @@ def apply(
         if maybe_use_numba(engine):
             if raw is False:
                 raise ValueError("raw must be `True` when using the numba engine")
-            numba_args = args
+            numba_args, kwargs = prepare_function_arguments(func, args, kwargs, 1)
             if self.method == "single":
                 apply_func = generate_numba_apply_func(
-                    func, **get_jit_arguments(engine_kwargs, kwargs)
+                    func, **get_jit_arguments(engine_kwargs)
                 )
             else:
                 apply_func = generate_numba_table_func(
-                    func, **get_jit_arguments(engine_kwargs, kwargs)
+                    func, **get_jit_arguments(engine_kwargs)
                 )
         elif engine in ("cython", None):
             if engine_kwargs is not None:

pandas/tests/apply/test_frame_apply.py

@@ -94,6 +94,16 @@ def test_apply_args(float_frame, axis, raw, engine, nopython):
     tm.assert_frame_equal(result, expected)
 
     if engine == "numba":
+        # py signature binding
+        with pytest.raises(TypeError, match="missing a required argument: 'a'"):
+            float_frame.apply(
+                lambda x, a: x + a,
+                b=2,
+                raw=raw,
+                engine=engine,
+                engine_kwargs=engine_kwargs,
+            )
+
         # keyword-only arguments are not supported in numba
         with pytest.raises(
             pd.errors.NumbaUtilError,

pandas/tests/groupby/aggregate/test_numba.py

@@ -35,18 +35,43 @@ def incorrect_function(x):
 def test_check_nopython_kwargs():
     pytest.importorskip("numba")
 
-    def incorrect_function(values, index):
-        return sum(values) * 2.7
+    def incorrect_function(values, index, *, a):
+        return sum(values) * 2.7 + a
+
+    def correct_function(values, index, a):
+        return sum(values) * 2.7 + a
 
     data = DataFrame(
         {"key": ["a", "a", "b", "b", "a"], "data": [1.0, 2.0, 3.0, 4.0, 5.0]},
         columns=["key", "data"],
     )
+    expected = data.groupby("key").sum() * 2.7
+
+    # py signature binding
+    with pytest.raises(
+        TypeError, match="missing a required (keyword-only argument|argument): 'a'"
+    ):
+        data.groupby("key").agg(incorrect_function, engine="numba", b=1)
+    with pytest.raises(TypeError, match="missing a required argument: 'a'"):
+        data.groupby("key").agg(correct_function, engine="numba", b=1)
+
+    with pytest.raises(
+        TypeError, match="missing a required (keyword-only argument|argument): 'a'"
+    ):
+        data.groupby("key")["data"].agg(incorrect_function, engine="numba", b=1)
+    with pytest.raises(TypeError, match="missing a required argument: 'a'"):
+        data.groupby("key")["data"].agg(correct_function, engine="numba", b=1)
+
+    # numba signature check after binding
     with pytest.raises(NumbaUtilError, match="numba does not support"):
         data.groupby("key").agg(incorrect_function, engine="numba", a=1)
+    actual = data.groupby("key").agg(correct_function, engine="numba", a=1)
+    tm.assert_frame_equal(expected + 1, actual)
 
     with pytest.raises(NumbaUtilError, match="numba does not support"):
         data.groupby("key")["data"].agg(incorrect_function, engine="numba", a=1)
+    actual = data.groupby("key")["data"].agg(correct_function, engine="numba", a=1)
+    tm.assert_series_equal(expected["data"] + 1, actual)
 
 
 @pytest.mark.filterwarnings("ignore")

pandas/tests/groupby/transform/test_numba.py

@@ -33,18 +33,43 @@ def incorrect_function(x):
 def test_check_nopython_kwargs():
     pytest.importorskip("numba")
 
-    def incorrect_function(values, index):
-        return values + 1
+    def incorrect_function(values, index, *, a):
+        return values + a
+
+    def correct_function(values, index, a):
+        return values + a
 
     data = DataFrame(
         {"key": ["a", "a", "b", "b", "a"], "data": [1.0, 2.0, 3.0, 4.0, 5.0]},
         columns=["key", "data"],
     )
+    # py signature binding
+    with pytest.raises(
+        TypeError, match="missing a required (keyword-only argument|argument): 'a'"
+    ):
+        data.groupby("key").transform(incorrect_function, engine="numba", b=1)
+    with pytest.raises(TypeError, match="missing a required argument: 'a'"):
+        data.groupby("key").transform(correct_function, engine="numba", b=1)
+
+    with pytest.raises(
+        TypeError, match="missing a required (keyword-only argument|argument): 'a'"
+    ):
+        data.groupby("key")["data"].transform(incorrect_function, engine="numba", b=1)
+    with pytest.raises(TypeError, match="missing a required argument: 'a'"):
+        data.groupby("key")["data"].transform(correct_function, engine="numba", b=1)
+
+    # numba signature check after binding
     with pytest.raises(NumbaUtilError, match="numba does not support"):
         data.groupby("key").transform(incorrect_function, engine="numba", a=1)
+    actual = data.groupby("key").transform(correct_function, engine="numba", a=1)
+    tm.assert_frame_equal(data[["data"]] + 1, actual)
 
     with pytest.raises(NumbaUtilError, match="numba does not support"):
         data.groupby("key")["data"].transform(incorrect_function, engine="numba", a=1)
+    actual = data.groupby("key")["data"].transform(
+        correct_function, engine="numba", a=1
+    )
+    tm.assert_series_equal(data["data"] + 1, actual)
 
 
 @pytest.mark.filterwarnings("ignore")

pandas/tests/window/test_numba.py

@@ -319,13 +319,24 @@ def f(x):
 
 @td.skip_if_no("numba")
 def test_invalid_kwargs_nopython():
+    with pytest.raises(TypeError, match="got an unexpected keyword argument 'a'"):
+        Series(range(1)).rolling(1).apply(
+            lambda x: x, kwargs={"a": 1}, engine="numba", raw=True
+        )
     with pytest.raises(
         NumbaUtilError, match="numba does not support keyword-only arguments"
     ):
         Series(range(1)).rolling(1).apply(
-            lambda x: x, kwargs={"a": 1}, engine="numba", raw=True
+            lambda x, *, a: x, kwargs={"a": 1}, engine="numba", raw=True
         )
 
+    tm.assert_series_equal(
+        Series(range(1), dtype=float) + 1,
+        Series(range(1))
+        .rolling(1)
+        .apply(lambda x, a: (x + a).sum(), kwargs={"a": 1}, engine="numba", raw=True),
+    )
+
 
 @td.skip_if_no("numba")
 @pytest.mark.slow