ENH: Allow numba aggregations to return non-float64 results #53444
Conversation
lithomas1
added
Groupby
Dtype Conversions
lithomas1
changed the title
| return column_looper | ||
|
|
||
|
|
||
| default_dtype_mapping: dict[np.dtype, Any] = { |
There was a problem hiding this comment.
Curious, could we not just define signatures for numba.jit to use when running the function?
There was a problem hiding this comment.
We allocate arrays inside the function and need to pass a dtype there as well.
Not sure how to access the signature from inside the func.
pandas/core/groupby/groupby.py
Outdated
| result = aggregator(sorted_data, starts, ends, 0, *aggregator_args) | ||
| result = sorted_df._mgr.apply( | ||
| aggregator, start=starts, end=ends, **aggregator_kwargs | ||
| ) |
There was a problem hiding this comment.
This is changing *arggregator_args -> **aggregator_kwargs, but then within aggregator it is being used as *aggregator_kwargs. This is only used internally right? I'm just wondering if we can make this less fragile somehow (changing of order kwargs might produce a bug, right?), but I'm not seeing a way.
There was a problem hiding this comment.
Yeah, the current method is really sketchy but it should be OK, since UDFs take another path.
(only args/kwargs that go through here are stuff like ddof for std/var).
The reason it's like this is since BlockManager.apply only takes kwargs.
Is it fine to change that?
| @@ -646,10 +646,27 @@ def _numba_apply( | |||
| step=self.step, | |||
| ) | |||
| self._check_window_bounds(start, end, len(values)) | |||
| # For now, map everything to float to match the Cython impl | |||
| # even though it is wrong | |||
| # TODO: Could preserve correct dtypes in future | |||
There was a problem hiding this comment.
#53214, I'll add it to the comment.
|
Could you also add a whatsnew note for 2.1? |
|
There seems to be some flakiness with the benchmarks I added. I'll let this sit for a couple of days then, but other than that it should be good to go. EDIT: Root caused, it was a timeout in the benchs. |
| # because it re-uses the Window min/max kernel | ||
| # so it will time out ASVs | ||
| # "min", | ||
| # "max", |
There was a problem hiding this comment.
Disabled min/max because it's reaaaallllly sloooooow.
It takes 20s (as opposed to milliseconds for the other kernels) to run, and can time out the ASVs sometimes(causing flakiness).
Best guess is that the list operations are slowing it down. Snakeviz tells me most (99% of the time) is spent in the numba kernel, and I can't profile into there.
I'm planning on splitting groupby stuff from the Window numba kernels in the future, so hopefully this doesn't stay commented for long.
…ev#53444) * ENH: non float64 result support in numba groupby * refactor & simplify * fix CI * maybe green? * skip unsupported ops in other bench as well * updates from code review * remove commented code * update whatsnew * debug benchmarks * Skip min/max benchmarks
…ev#53444) * ENH: non float64 result support in numba groupby * refactor & simplify * fix CI * maybe green? * skip unsupported ops in other bench as well * updates from code review * remove commented code * update whatsnew * debug benchmarks * Skip min/max benchmarks
doc/source/whatsnew/vX.X.X.rstfile if fixing a bug or adding a new feature.