diff --git a/python/cudf/cudf/tests/test_udf_masked_ops.py b/python/cudf/cudf/tests/test_udf_masked_ops.py index dc126546f15..c9c2c440632 100644 --- a/python/cudf/cudf/tests/test_udf_masked_ops.py +++ b/python/cudf/cudf/tests/test_udf_masked_ops.py @@ -593,3 +593,22 @@ def func(row, c, k): return y run_masked_udf_test(func, data, args=(1, 2), check_dtype=False) + + +def test_masked_udf_caching(): + # Make sure similar functions that differ + # by simple things like constants actually + # recompile + + data = cudf.Series([1, 2, 3]) + expect = data ** 2 + got = data.applymap(lambda x: x ** 2) + + assert_eq(expect, got, check_dtype=False) + + # update the constant value being used and make sure + # it does not result in a cache hit + + expect = data ** 3 + got = data.applymap(lambda x: x ** 3) + assert_eq(expect, got, check_dtype=False) diff --git a/python/cudf/cudf/utils/cudautils.py b/python/cudf/cudf/utils/cudautils.py index 5fa091a0081..f0533dcaa72 100755 --- a/python/cudf/cudf/utils/cudautils.py +++ b/python/cudf/cudf/utils/cudautils.py @@ -216,12 +216,14 @@ def make_cache_key(udf, sig): recompiling the same function for the same set of types """ codebytes = udf.__code__.co_code + constants = udf.__code__.co_consts if udf.__closure__ is not None: cvars = tuple([x.cell_contents for x in udf.__closure__]) cvarbytes = dumps(cvars) else: cvarbytes = b"" - return codebytes, cvarbytes, sig + + return constants, codebytes, cvarbytes, sig def compile_udf(udf, type_signature):