diff --git a/python/cudf/cudf/core/udf/_ops.py b/python/cudf/cudf/core/udf/_ops.py index 559a5bfad4f..6b0640b09ed 100644 --- a/python/cudf/cudf/core/udf/_ops.py +++ b/python/cudf/cudf/core/udf/_ops.py @@ -11,6 +11,13 @@ operator.floordiv, operator.mod, operator.pow, + operator.iadd, + operator.isub, + operator.imul, + operator.itruediv, + operator.floordiv, + operator.ipow, + operator.imod, ] bitwise_ops = [operator.and_, operator.or_, operator.xor] diff --git a/python/cudf/cudf/tests/test_udf_masked_ops.py b/python/cudf/cudf/tests/test_udf_masked_ops.py index 20245bd2a20..f1d110ba168 100644 --- a/python/cudf/cudf/tests/test_udf_masked_ops.py +++ b/python/cudf/cudf/tests/test_udf_masked_ops.py @@ -197,6 +197,10 @@ def func(row): operator.pow, operator.truediv, operator.floordiv, + operator.imod, + operator.ipow, + operator.itruediv, + operator.ifloordiv, }: # The following tests cases yield undefined behavior: # - truediv(x, False) because its dividing by zero @@ -219,7 +223,7 @@ def func(row): # Just a single column -> result will be all NA gdf = cudf.DataFrame({"data": data}) - if constant == 1 and op is operator.pow: + if constant == 1 and op in {operator.pow, operator.ipow}: # The following tests cases yield differing results from pandas: # - 1**NA # - True**NA @@ -237,7 +241,7 @@ def func(row): gdf = cudf.DataFrame({"data": data}) - if 1 in gdf["data"] and op is operator.pow: + if 1 in gdf["data"] and op in {operator.pow, operator.ipow}: # In pandas, 1**NA == 1. pytest.skip() run_masked_udf_test(func, gdf, check_dtype=False) @@ -483,7 +487,7 @@ def func(x): # Just a single column -> result will be all NA data = cudf.Series([1, 2, cudf.NA]) - if constant is cudf.NA and op is operator.pow: + if constant is cudf.NA and op in {operator.pow, operator.ipow}: # in pandas, 1**NA == 1. In cudf, 1**NA == 1. with pytest.xfail(): run_masked_udf_series(func, data, check_dtype=False) @@ -499,7 +503,11 @@ def func(x): # Just a single column -> result will be all NA data = cudf.Series([1, 2, cudf.NA]) - if constant is not cudf.NA and constant == 1 and op is operator.pow: + if ( + constant is not cudf.NA + and constant == 1 + and op in {operator.pow, operator.ipow} + ): # in pandas, 1**NA == 1. In cudf, 1**NA == 1. with pytest.xfail(): run_masked_udf_series(func, data, check_dtype=False)