Skip to content

Commit

Permalink
fix the bug
Browse files Browse the repository at this point in the history
  • Loading branch information
brandon-b-miller committed Dec 2, 2021
1 parent 582cc6e commit 91daf5a
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
19 changes: 19 additions & 0 deletions python/cudf/cudf/tests/test_udf_masked_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
4 changes: 3 additions & 1 deletion python/cudf/cudf/utils/cudautils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 91daf5a

Please sign in to comment.