python - 使用 SparseDtype['string', pd.NA] 从 pickle 加载数据帧会导致错误，具体取决于数据帧的大小

Question

我在腌制数据框时遇到了一些问题：我观察到的是，从 pickle 读取后的数据框由于StringArray requires a sequence of strings or pandas.NA. 进一步的测试导致我创建了一个较小的测试用例，但有趣的是数据框需要足够大，我的问题才能重现。

import pandas as pd
import numpy as np

# create a data frame with 5 rows
df_small = pd.DataFrame(
    {'A': list(range(5)),
     'B': ['v1', 'v1', 'v1', 'v2', pd.NA],
     'C': [pd.NA, 'value2', 'value1', 'value1', 'value1']})

# create a data frame with 25000 rows
df_large = pd.DataFrame(
    {'A': list(range(5)) * 5000,
     'B': ['v1', 'v1', 'v1', 'v2', pd.NA] * 5000,
     'C': [pd.NA, 'value2', 'value1', 'value1', 'value1'] * 5000})

# assign type, notice that the SparseDtype is used
df_small = df_small.astype({'A': 'int32', 'B': pd.SparseDtype('string', pd.NA), 'C': 'string'})
df_large = df_large.astype({'A': 'int32', 'B': pd.SparseDtype('string', pd.NA), 'C': 'string'})

#store to disk
df_small.to_pickle('issue_small.pkl')
df_large.to_pickle('issue_large.pkl')

# load from disk
df_small1 = pd.read_pickle((open('issue_small.pkl', 'rb')))
df_large1 = pd.read_pickle((open('issue_large.pkl', 'rb')))

# this one will print OK
print(df_small1.C)

# this one will throw a ValueError
print(df_large1.C)

几乎我可以尝试用 df_large1 做的任何事情都会抛出这个错误。

当类型的分配被注释掉时，问题就消失了。对我来说不幸的是，我拥有的大多数列的密度都低于 1%，并且总数据量大得令人望而却步，所以我或多或少地被困在使用稀疏数据对象上。

此外，无论是否使用 SparseDtype，都会df_small.isna()显示df_small1.isna()不同的结果 - 正如前者所预期的那样，但都是False后者。

关于我可以改变什么来规避这种行为的任何想法？

这是在 Python 3.8.1 和 pandas 1.0.0 上完成的——此时降级可能是一种选择。

标签： pythonpandasdataframepickle