BUG: Pickle NA objects

According to
https://docs.python.org/3/library/pickle.html#object.__reduce__,

> If a string is returned, the string should be interpreted as the name
> of a global variable. It should be the object’s local name relative to
> its module; the pickle module searches the module namespace to determine
> the object’s module. This behaviour is typically useful for singletons.

Closes pandas-dev#31847

Author: TomAugspurger

doc/source/whatsnew/v1.0.2.rst

@@ -37,6 +37,7 @@ Bug fixes
 **I/O**
 
 - Using ``pd.NA`` with :meth:`DataFrame.to_json` now correctly outputs a null value instead of an empty object (:issue:`31615`)
+- Fixed pickling of ``pandas.NA``. Previously a new object was returned, which broke computations relying on ``NA`` being a singleton (:issue:`31847`)
 - Fixed bug in parquet roundtrip with nullable unsigned integer dtypes (:issue:`31896`).
 
 

pandas/_libs/missing.pyx

@@ -364,6 +364,9 @@ class NAType(C_NAType):
         exponent = 31 if is_32bit else 61
         return 2 ** exponent - 1
 
+    def __reduce__(self):
+        return "NA"
+
     # Binary arithmetic and comparison ops -> propagate
 
     __add__ = _create_binary_propagating_op("__add__")

pandas/tests/scalar/test_na_scalar.py

@@ -1,3 +1,5 @@
+import pickle
+
 import numpy as np
 import pytest
 
@@ -267,3 +269,9 @@ def test_integer_hash_collision_set():
     assert len(result) == 2
     assert NA in result
     assert hash(NA) in result
+
+
+def test_pickle_roundtrip():
+    # https://github.com/pandas-dev/pandas/issues/31847
+    result = pickle.loads(pickle.dumps(pd.NA))
+    assert result is pd.NA