当列包含集合时如何保存 pandas 数据框

Question

当尝试保存其中一列包含集合的 pandas 数据框时（请参见下面的示例），

import pandas as pd

df = pd.DataFrame({"col_set": [{"A", "B", "C"}, {"D", "E", "F"}]})
df.to_parquet("df_w_col_set.parquet")

会引发以下错误：

ArrowInvalid: ("Could not convert {'C', 'B', 'A'} with type set: did not recognize Python value type when inferring an Arrow data type", 'Conversion failed for column col_set with type object')

如何保存这种数据框并避免上述错误？

一些半相关的帖子提到提供西洋蓍草模式，但我不清楚在咨询时使用什么类型 pyarrow 数据类型。

代码使用 python 3.7.4、pandas==1.3.0 和 pyarrow==3.0.0 运行

主要寻找升级的解决方案不需要或真正最小化（以避免破坏其他依赖项）。

Answer 1

作为解决方法，您可以将您的 set 转换为字符串，并使用 ast.literal_eval 将您的字符串评估为 set：

import ast

df.astype({'col_set': str}).to_parquet('data.parquet')
df1 = pd.read_parquet('data.parquet') \
        .assign(col_set=lambda x: x['col_set'].map(ast.literal_eval))
print(df1)

# Output
     col_set
0  {C, B, A}
1  {F, E, D}

或者您可以将您的集合转换为tuple（或list）然后恢复为set：

df.assign(col_set=df['col_set'].map(tuple)).to_parquet('test.parquet')
df1 = pd.read_parquet('test.parquet') \
        .assign(col_set=lambda x: x['col_set'].map(set))
print(df1)

# Output
     col_set
0  {C, B, A}
1  {F, E, D}

您还可以使用pickle.dumps和pickle .loads 来序列化你的set：

import pickle

df.assign(col_set=df['col_set'].map(pickle.dumps)).to_parquet('test.parquet')
df1 = pd.read_parquet('test.parquet') \
        .assign(col_set=lambda x: x['col_set'].map(pickle.loads))
print(df1)

# Output
     col_set
0  {C, B, A}
1  {F, E, D}

事实上，您可以选择任何（非）序列化方法（JSON 除外，因为 set 不存在）。