modify prepare_function_arguments

Author: auderson

pandas/core/apply.py

@@ -1004,14 +1004,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
@@ -1168,9 +1169,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
 
@@ -1313,9 +1314,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

@@ -137,6 +137,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:
@@ -1443,8 +1444,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,
@@ -1479,8 +1481,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

@@ -66,6 +66,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,
@@ -1458,14 +1459,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: