Add tokenize NA functionality (dask#10640)

mrocklin · Nov 22, 2023 · da7dc67 · da7dc67
1 parent 204be6d
commit da7dc67
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/dask/base.py b/dask/base.py
@@ -1194,6 +1194,10 @@ def normalize_categorical_dtype(dtype):
     def normalize_period_dtype(dtype):
         return normalize_token(dtype.name)
 
+    @normalize_token.register(type(pd.NA))
+    def normalize_na(dtype):
+        return pd.NA
+
 
 @normalize_token.register_lazy("numpy")
 def register_numpy():

diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
@@ -354,6 +354,11 @@ def test_tokenize_pandas_extension_array():
         assert tokenize(arr) == tokenize(arr)
 
 
+@pytest.mark.skipif("not dd")
+def test_tokenize_na():
+    assert tokenize(pd.NA) == tokenize(pd.NA)
+
+
 @pytest.mark.skipif("not pd")
 def test_tokenize_pandas_index():
     idx = pd.Index(["a", "b"])