address reviews, fix tests

rapidsai · May 9, 2022 · f63a958 · f63a958
1 parent 5358172
commit f63a958
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/python/cudf/cudf/__init__.py b/python/cudf/cudf/__init__.py
@@ -17,9 +17,8 @@
     register_index_accessor,
     register_series_accessor,
 )
-from cudf.core.scalar import (
-    Scalar,
-)
+from cudf.core.scalar import Scalar
+
 from cudf.core.index import (
     BaseIndex,
     CategoricalIndex,

diff --git a/python/cudf/cudf/core/udf/typing.py b/python/cudf/cudf/core/udf/typing.py
@@ -17,7 +17,7 @@
 )
 from numba.core.typing.typeof import typeof
 from numba.cuda.cudadecl import registry as cuda_decl_registry
-from pandas._libs.missing import NAType as _NAType
+from cudf.core.missing import NA
 
 from cudf.core.udf import api
 from cudf.core.udf._ops import (
@@ -214,7 +214,7 @@ def unify(self, context, other):
 na_type = NAType()
 
 
-@typeof_impl.register(_NAType)
+@typeof_impl.register(type(NA))
 def typeof_na(val, c):
     """
     Tie instances of _NAType (cudf.NA) to our NAType.

diff --git a/python/cudf/cudf/tests/test_udf_masked_ops.py b/python/cudf/cudf/tests/test_udf_masked_ops.py
@@ -7,7 +7,7 @@
 from numba import cuda
 
 import cudf
-from cudf.core.scalar import NA
+from cudf.core.missing import NA
 from cudf.core.udf._ops import (
     arith_ops,
     bitwise_ops,