Merge branch 'groupby-apply-caching' of github.com:brandon-b-miller/c…

…udf into groupby-apply-caching
rapidsai · Mar 17, 2023 · 197bce6 · 197bce6
2 parents 4858af2 + 6866975
commit 197bce6
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/python/cudf/cudf/core/udf/groupby_utils.py b/python/cudf/cudf/core/udf/groupby_utils.py
@@ -150,13 +150,11 @@ def jit_groupby_apply(offsets, grouped_values, function, *args):
     ngroups = len(offsets) - 1
 
     cache_key = _generate_cache_key(grouped_values, function)
-    if precompiled.get(cache_key) is not None:
-        kernel, return_type = precompiled[cache_key]
-    else:
-        kernel, return_type = _get_groupby_apply_kernel(
+    if cache_key not in precompiled:
+        precompiled[cache_key] = _get_groupby_apply_kernel(
             grouped_values, function, args
         )
-        precompiled[cache_key] = (kernel, return_type)
+    kernel, return_type = precompiled[cache_key]
 
     return_type = numpy_support.as_dtype(return_type)
     output = cudf.core.column.column_empty(ngroups, dtype=return_type)