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

[auderson]

• 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.

@lithomas1

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:

pandas/pandas/core/groupby/numba_.py

Lines 102 to 120 in 3b48b17

	@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.

[github-actions]

This pull request is stale because it has been open for thirty days with no activity. Please update and respond to this comment if you're still interested in working on this.

[auderson]

This pull request is stale because it has been open for thirty days with no activity. Please update and respond to this comment if you're still interested in working on this.

Please do not mark as stale, this PR is ready for review.

[rhshadrach]

@auderson - can you merge main.

[auderson]

Failures seem unrelated: FAILED pandas/tests/io/test_spss.py::test_spss_metadata - AssertionError.

[rhshadrach]

This is looking good to me. A few requests.

[rhshadrach]

@lithomas1 - friendly ping.

[rhshadrach]

lgtm, going to wait on merging a bit longer in case there are other reviews.

[auderson]

Also I find there are some methods like RollingGroupby.apply which are also impacted by the changes to RollingAndExpandingMixin.apply. So I added more tests to cover these methods in tests/window/test_numba.

[mroeschke]

Thanks @auderson