Repository: pandas-dev/pandas
Pull request: 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 https://github.com/pandas-dev/pandas/issues/31847

Problem statement:
BUG: pd.NA doesn't pickle/unpickle faithfully
#### Code Sample, a copy-pastable example if possible

```python

In [5]: df['Gold Categories'].count()
Out[5]: 135218

In [6]: df['Gold Categories'].isna().sum()
Out[6]: 0

In [7]: df['Gold Categories'].iloc[256]
Out[7]: <NA>

In [8]: pd.isna(df['Gold Categories'].iloc[256])
Out[8]: False

In [9]: type(df['Gold Categories'].iloc[256])
Out[9]: pandas._libs.missing.NAType

In [10]: pd.__version__
Out[10]: '1.0.1'


```

#### Output of ``pd.show_versions()``

<details>


INSTALLED VERSIONS
------------------
commit           : None
python           : 3.7.5.final.0
python-bits      : 64
OS               : Linux
OS-release       : 5.3.16-200.fc30.x86_64
machine          : x86_64
processor        : x86_64
byteorder        : little
LC_ALL           : None
LANG             : nb_NO.UTF-8
LOCALE           : nb_NO.UTF-8

pandas           : 1.0.1
numpy            : 1.17.3
pytz             : 2019.3
dateutil         : 2.8.0
pip              : 19.3.1
setuptools       : 41.6.0.post20191030
Cython           : 0.29.13
pytest           : 5.2.2
hypothesis       : None
sphinx           : 2.2.1
blosc            : None
feather          : None
xlsxwriter       : 1.2.2
lxml.etree       : 4.4.1
html5lib         : 1.0.1
pymysql          : None
psycopg2         : 2.8.4 (dt dec pq3 ext lo64)
jinja2           : 2.10.3
IPython          : 7.9.0
pandas_datareader: None
bs4              : 4.8.1
bottleneck       : 1.2.1
fastparquet      : None
gcsfs            : None
lxml.etree       : 4.4.1
matplotlib       : 2.2.3
numexpr          : 2.7.0
odfpy            : None
openpyxl         : 3.0.0
pandas_gbq       : None
pyarrow          : 0.15.1
pytables         : None
pytest           : 5.2.2
pyxlsb           : None
s3fs             : None
scipy            : 1.3.1
sqlalchemy       : 1.3.10
tables           : 3.5.2
tabulate         : 0.8.5
xarray           : None
xlrd             : 1.2.0
xlwt             : 1.3.0
xlsxwriter       : 1.2.2
numba            : 0.46.0


</details>

Hints from issue discussion:
Could you include a reproducible example? I could not reproduce this on master:
```
>>> pd.DataFrame({'Gold categories': [pd.NA]})['Gold categories'].iloc[0]                                                                                                                                   
<NA>

>>> pd.isna(pd.DataFrame({'Gold categories': [pd.NA]})['Gold categories'].iloc[0])                                                                                                                          
True
```
@MarcoGorelli  I can upload a sample, if that will suffice, but I do not know how to reproduce
@tsoernes if you run the two lines I posted above, do you get the same output?
@MarcoGorelli  Yes
Here is a sample of that column with 2 rows. It is a zipped pickle file (Github only allows zips).
```

In [167]: na_test = read_pickle('/tmp/na_test.pickle')
Loaded 2 entries (a Series) from /tmp/na_test.pickle (2020-02-11 21:51)

In [168]: na_test.isna().sum()
Out[173]: 0

In [174]: na_test

Out[176]: 
268    [Fintech]
269         <NA>
Name: Gold Categories, dtype: object
```

@tsoernes I'm afraid we can't accept raw pickle files in bug reports, [as they could be unsafe](https://docs.python.org/3/library/pickle.html). Please remove the attachment from your message :)

Could you please paste the output of `na_test.to_dict()`, so we can copy-and-paste it and reproduce the issue?



```
na_test.to_dict()
Out[195]: {268: ['Fintech'], 269: <NA>}
```
Thanks @tsoernes 

TBH I still can't reproduce the issue though
```
>>> pd.Series({268: ['Fintech'], 269: pd.NA}).isna()                                                        
268    False
269     True
dtype: bool
```
I can't either, when going via a dictionary.

On Tue, Feb 11, 2020 at 11:57 PM Marco Gorelli <notifications@github.com>
wrote:

> Thanks @tsoernes <https://github.com/tsoernes>
>
> TBH I still can't reproduce the issue though
>
> >>> pd.Series({268: ['Fintech'], 269: pd.NA}).isna()
> 268    False
> 269     True
> dtype: bool
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pandas-dev/pandas/issues/31847?email_source=notifications&email_token=ABTX3RBRDVT7VZ5DSRTPSYTRCMNNRA5CNFSM4KSMLNH2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOELOHURQ#issuecomment-584874566>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ABTX3RADPUPEHTYF7YCSB3LRCMNNRANCNFSM4KSMLNHQ>
> .
>

When pickling/unpickling, I can reproduce this:

```
In [40]: s = pd.Series({268: ['Fintech'], 269: pd.NA})                                                                                                                                                             

In [41]: s.isna()                                                                                                                                                                                                  
Out[41]: 
268    False
269     True
dtype: bool

In [42]: s.to_pickle('test_na_pickle.pkl')                                                                                                                                                                         

In [43]: s2 = pd.read_pickle('test_na_pickle.pkl')                                                                                                                                                                 

In [44]: s2.isna()                                                                                                                                                                                                 
Out[44]: 
268    False
269    False
dtype: bool

In [45]: type(s2.values[1])                                                                                                                                                                                        
Out[45]: pandas._libs.missing.NAType

In [46]: s2.values[1] is pd.NA                                                                                                                                                                                     
Out[46]: False
```

So apparently, when unpickling, it doesn't return the same singleton.

Commit message: more

Patch to review:
diff --git a/doc/source/whatsnew/v1.0.2.rst b/doc/source/whatsnew/v1.0.2.rst
index c9031ac1ae9fe..0201f29fa87e4 100644
--- a/doc/source/whatsnew/v1.0.2.rst
+++ b/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`).
 
 
diff --git a/pandas/_libs/missing.pyx b/pandas/_libs/missing.pyx
index 4d17a6f883c1c..c54cb652d7b21 100644
--- a/pandas/_libs/missing.pyx
+++ b/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__")
diff --git a/pandas/tests/scalar/test_na_scalar.py b/pandas/tests/scalar/test_na_scalar.py
index dcb9d66708724..6dff0eb8ae91f 100644
--- a/pandas/tests/scalar/test_na_scalar.py
+++ b/pandas/tests/scalar/test_na_scalar.py
@@ -1,3 +1,5 @@
+import pickle
+
 import numpy as np
 import pytest
 
@@ -267,3 +269,29 @@ 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
+
+
+def test_pickle_roundtrip_pandas():
+    with tm.ensure_clean("data.pkl") as f:
+        with open(f):
+            pd.to_pickle(pd.NA, f)
+        with open(f):
+            result = pd.read_pickle(f)
+    assert result is pd.NA
+
+
+def test_pickle_roundtrip_series():
+    s = pd.Series(pd.array([1, 2, pd.NA]))
+    with tm.ensure_clean("data.pkl") as f:
+        with open(f):
+            pd.to_pickle(s, f)
+        with open(f):
+            result = pd.read_pickle(f)
+
+    tm.assert_series_equal(result, s)


Code review comments:
diff_hunk: @@ -267,3 +269,29 @@ 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
+
+
+def test_pickle_roundtrip_pandas():
+    with tm.ensure_clean("data.pkl") as f:
path: pandas/tests/scalar/test_na_scalar.py
text: @user1:
there are test functions to do this in pandas._testing