From f63a9582e3214212da1bd87566d1effee06b9b77 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 9 May 2022 09:41:31 -0700
Subject: [PATCH] address reviews, fix tests

---
 python/cudf/cudf/__init__.py                  | 5 ++---
 python/cudf/cudf/core/udf/typing.py           | 4 ++--
 python/cudf/cudf/tests/test_udf_masked_ops.py | 2 +-
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/python/cudf/cudf/__init__.py b/python/cudf/cudf/__init__.py
index 2741058cd37..27aaa07f940 100644
--- 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
index ed5fc1d6d23..83aa12ddf93 100644
--- 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
index 438f46d4266..73bf5160347 100644
--- 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,