From 68669759b3ba403d3ec199e098aea2c13f124f6f Mon Sep 17 00:00:00 2001 From: brandon-b-miller <53796099+brandon-b-miller@users.noreply.github.com> Date: Wed, 15 Mar 2023 13:50:42 -0500 Subject: [PATCH] Update python/cudf/cudf/core/udf/groupby_utils.py Co-authored-by: Bradley Dice --- python/cudf/cudf/core/udf/groupby_utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/python/cudf/cudf/core/udf/groupby_utils.py b/python/cudf/cudf/core/udf/groupby_utils.py index 67a90a920e9..669db76db9a 100644 --- 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)