Skip to content

REF: Create numba helper function for jitting + generating cache key #33910

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
May 1, 2020
Merged

REF: Create numba helper function for jitting + generating cache key #33910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 5 additions & 16 deletions pandas/core/groupby/generic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,8 @@
from pandas.core.series import Series
from pandas.core.util.numba_ import (
NUMBA_FUNC_CACHE,
Copy link
Contributor

Choose a reason for hiding this comment

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

can you remove the CACHE from here as well?

Copy link
Member Author

Choose a reason for hiding this comment

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

There are some separate code paths that store the function in the cache independently

Copy link
Contributor

Choose a reason for hiding this comment

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

i c. ideally this is not visible at all in this code if at all possible (this is what I mean by using a Dispatcher, e.g. that you would be how you would interactive with the function & cache itself). but i guess that's for the future.

Copy link
Member Author

Choose a reason for hiding this comment

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

Ah I see. Okay I can look into that in a future PR when playing around more with a Dispatcher

check_kwargs_and_nopython,
get_jit_arguments,
jit_user_function,
generate_numba_func,
split_for_numba,
validate_udf,
)

from pandas.plotting import boxplot_frame_groupby
Expand Down Expand Up @@ -507,12 +504,8 @@ def _transform_general(
"""

if engine == "numba":
nopython, nogil, parallel = get_jit_arguments(engine_kwargs)
check_kwargs_and_nopython(kwargs, nopython)
validate_udf(func)
cache_key = (func, "groupby_transform")
numba_func = NUMBA_FUNC_CACHE.get(
cache_key, jit_user_function(func, nopython, nogil, parallel)
numba_func, cache_key = generate_numba_func(
func, engine_kwargs, kwargs, "groupby_transform"
)

klass = type(self._selected_obj)
Expand Down Expand Up @@ -1407,12 +1400,8 @@ def _transform_general(
obj = self._obj_with_exclusions
gen = self.grouper.get_iterator(obj, axis=self.axis)
if engine == "numba":
nopython, nogil, parallel = get_jit_arguments(engine_kwargs)
check_kwargs_and_nopython(kwargs, nopython)
validate_udf(func)
cache_key = (func, "groupby_transform")
numba_func = NUMBA_FUNC_CACHE.get(
cache_key, jit_user_function(func, nopython, nogil, parallel)
numba_func, cache_key = generate_numba_func(
func, engine_kwargs, kwargs, "groupby_transform"
)
else:
fast_path, slow_path = self._define_paths(func, *args, **kwargs)
Expand Down
13 changes: 3 additions & 10 deletions pandas/core/groupby/ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,11 +56,8 @@
)
from pandas.core.util.numba_ import (
NUMBA_FUNC_CACHE,
check_kwargs_and_nopython,
get_jit_arguments,
jit_user_function,
generate_numba_func,
split_for_numba,
validate_udf,
)


Expand Down Expand Up @@ -689,12 +686,8 @@ def _aggregate_series_pure_python(
):

if engine == "numba":
nopython, nogil, parallel = get_jit_arguments(engine_kwargs)
check_kwargs_and_nopython(kwargs, nopython)
validate_udf(func)
cache_key = (func, "groupby_agg")
numba_func = NUMBA_FUNC_CACHE.get(
cache_key, jit_user_function(func, nopython, nogil, parallel)
numba_func, cache_key = generate_numba_func(
func, engine_kwargs, kwargs, "groupby_agg"
)

group_index, _, ngroups = self.group_info
Expand Down
40 changes: 40 additions & 0 deletions pandas/core/util/numba_.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,3 +167,43 @@ def f(values, index, ...):
f"The first {min_number_args} arguments to {func.__name__} must be "
f"{expected_args}"
)


def generate_numba_func(
func: Callable,
engine_kwargs: Optional[Dict[str, bool]],
kwargs: dict,
cache_key_str: str,
) -> Tuple[Callable, Tuple[Callable, str]]:
"""
Return a JITed function and cache key for the NUMBA_FUNC_CACHE

This _may_ be specific to groupby (as it's only used there currently).

Parameters
----------
func : function
user defined function
engine_kwargs : dict or None
numba.jit arguments
kwargs : dict
kwargs for func
cache_key_str : str
string representing the second part of the cache key tuple

Returns
-------
(JITed function, cache key)

Raises
------
NumbaUtilError
"""
nopython, nogil, parallel = get_jit_arguments(engine_kwargs)
check_kwargs_and_nopython(kwargs, nopython)
validate_udf(func)
cache_key = (func, cache_key_str)
numba_func = NUMBA_FUNC_CACHE.get(
cache_key, jit_user_function(func, nopython, nogil, parallel)
)
return numba_func, cache_key