Automatically select GroupBy.apply algorithm based on if the UDF is jittable
#13113
Conversation
brandon-b-miller
added
feature request
|
hi all, this should be ready for review. |
| raise ValueError( | ||
| "Nulls not yet supported with groupby JIT engine" | ||
| ) | ||
| # TODO: don't check this twice under `engine='auto'` |
There was a problem hiding this comment.
Is there (will there be?) a bug for this? Or do you intend to fix it here?
There was a problem hiding this comment.
I thought this might be a perf hit but since each column caches its null count I think it's actually ok to just leave it.
| @@ -1198,7 +1199,7 @@ def _iterative_groupby_apply( | |||
| result.index = cudf.MultiIndex._from_data(index_data) | |||
| return result | |||
|
|
|||
| def apply(self, function, *args, engine="cudf"): | |||
| def apply(self, function, *args, engine="auto"): | |||
There was a problem hiding this comment.
docstring needs to be updated to discuss new "auto" option.
There was a problem hiding this comment.
Added some docs here.
| @@ -1252,7 +1249,7 @@ def apply(self, function, *args, engine="cudf"): | |||
| on the grouped chunk. | |||
| args : tuple | |||
| Optional positional arguments to pass to the function. | |||
| engine: {'cudf', 'jit'}, default 'cudf' | |||
| engine: {'cudf', 'jit'}, default 'auto' | |||
There was a problem hiding this comment.
For my knowledge: Is there a performance cost to the fallback? i.e. Does the JIT attempt have measurable overhead?
There was a problem hiding this comment.
Yes, it does have measurable overhead. One way of measuring it is with
import cudf
df = cudf.DataFrame({
'a':[0,1,1],
'b':[1,2,3]
})
def func(grp):
# binops can't be jitted without refcounting
return grp + grp
grouped = df.groupby('a')
import cProfile
cProfile.run('grouped.apply(func)', sort='cumtime')For this I get
1 0.000 0.000 0.213 0.213 groupby_utils.py:207(_can_be_jitted)
Meaning it's quite impactful. However if this becomes a problem users should be able to obtain the old behavior by just passing engine='cudf'.
There was a problem hiding this comment.
Full disclosure, we anticipated this and I was OK with it. I think the tradeoff is generally worthwhile. If we think it isn't then I think we'd just stop this work and remove 'auto' altogether.
Also I just noticed that 'auto' is not listed in the set of valid engine arguments here.
There was a problem hiding this comment.
Just to clarify — your snippet above is saying the JIT attempt costs ~200ms? That sounds right to me. I would think a cache could also be used here if needed to prevent multiple failed attempts from paying the overhead for the same function.
I am supportive of this change, because when it does pay off, it’s a big win. Just want to make sure we’re putting in the appropriate amount of engineering effort to mitigate the downside risk.
There was a problem hiding this comment.
Also I just noticed that 'auto' is not listed in the set of valid engine arguments here.
Fixed this.
|
/merge |
Closes #13103