Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 23 commits into
base: main
Choose a base branch
from

Conversation

auderson
Copy link
Contributor

@auderson auderson commented Jun 13, 2024

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

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

@auderson auderson changed the title Enh numba invalid kwargs checking ENH numba invalid kwargs checking Jun 13, 2024
@auderson auderson changed the title ENH numba invalid kwargs checking ENH: numba apply supports positional arguments passed as **kwargs Jun 13, 2024
@auderson auderson force-pushed the enh_numba_apply_support_kwargs branch from 8168d9b to c211119 Compare June 13, 2024 11:49
@rhshadrach rhshadrach added Enhancement Groupby Apply Apply, Aggregate, Transform, Map numba numba-accelerated operations labels Jun 15, 2024
Copy link
Contributor

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.

@github-actions github-actions bot added the Stale label Jul 16, 2024
@auderson
Copy link
Contributor Author

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
Copy link
Member

@auderson - can you merge main.

@rhshadrach rhshadrach removed the Stale label Oct 21, 2024
@auderson
Copy link
Contributor Author

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

Copy link
Member

@rhshadrach rhshadrach left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is looking good to me. A few requests.

pandas/core/apply.py Outdated Show resolved Hide resolved
pandas/core/util/numba_.py Outdated Show resolved Hide resolved
@rhshadrach
Copy link
Member

@lithomas1 - friendly ping.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Apply Apply, Aggregate, Transform, Map Enhancement Groupby numba numba-accelerated operations
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants