ENH: numba apply supports positional arguments passed as **kwargs by auderson · Pull Request #58995 · pandas-dev/pandas (original) (raw)

Tests added and passed if fixing a bug or adding a new feature
All code checks passed.
Added type annotations to new arguments/methods/functions.
Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

This is a followup PR that enhances prepare_function_arguments introduced in #58767. prepare_function_arguments helps to move **kwargs to *args where possible using inspect.Signature. And now it's also more reasonable to move kwargs check from get_jit_arguments to prepare_function_arguments.

ref: #58767 (comment)

We can also extend prepare_function_arguments by adding a option: num_required_args, so that all the numba functions' args & kwargs stuff can be handled by this function, instead of in get_jit_arguments . For example in groupby, we have numba_func(group, group_index, *args), in this case num_required_args = 2:

@numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)

def group_agg(

values: np.ndarray,

index: np.ndarray,

begin: np.ndarray,

end: np.ndarray,

num_columns: int,

*args: Any,

) -> np.ndarray:

assert len(begin) == len(end)

num_groups = len(begin)

result = np.empty((num_groups, num_columns))

for i in numba.prange(num_groups):

group_index = index[begin[i] : end[i]]

for j in numba.prange(num_columns):

group = values[begin[i] : end[i], j]

result[i, j] = numba_func(group, group_index, *args)

return result

@numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
def group_agg(
values: np.ndarray,
index: np.ndarray,
begin: np.ndarray,
end: np.ndarray,
num_columns: int,
*args: Any,
) -> np.ndarray:
assert len(begin) == len(end)
num_groups = len(begin)

result = np.empty((num_groups, num_columns))
for i in numba.prange(num_groups):
group_index = index[begin[i] : end[i]]
for j in numba.prange(num_columns):
group = values[begin[i] : end[i], j]
result[i, j] = numba_func(group, group_index, *args)
return result

There are several places where the usage of get_jit_arguments is updated and prepare_function_arguments is added : groupby.agg, groupby.transform, rolling.apply and df.apply.