如何在列包含集合时保存 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 datatypes。

代码为运行，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不存在）。

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

How to save a pandas dataframe when a column contains sets

python

pandas

dataframe

parquet

pyarrow