使用 pyarrow 将结构列写入镶木地板

【问题标题】：Write struct columns to parquet with pyarrow使用 pyarrow 将结构列写入镶木地板
【发布时间】：2021-08-27 11:36:10
【问题描述】：

我有以下数据框和架构：

df = pd.DataFrame([[1,2,3], [4,5,6], [7,8,9]], columns=['a', 'b', 'c'])
SCHEMA = pa.schema([("a_and_b", pa.struct([('a', pa.int64()), ('b', pa.int64())])), ('c', pa.int64())])

然后我想从 df 创建一个 pyarrow 表并使用此模式将其保存到镶木地板。但是，我找不到在 pandas 中创建与 pyarrow 中的结构类型相对应的正确类型的方法。有没有办法做到这一点？

【问题讨论】：

标签： python pandas dataframe parquet pyarrow

【解决方案1】：

对于 pa.struct 从 pandas 转换，您可以使用元组（例如：[(1, 4), (2, 5), (3, 6)]）：

df_with_tuples = pd.DataFrame({
    "a_and_b": zip(df["a"], df["b"]),
    "c":  df["c"]
})
pa.Table.from_pandas(df_with_tuples, SCHEMA)

或 dict[{'a': 1, 'b': 2}, {'a': 4, 'b': 5}, {'a': 7, 'b': 8}]:

df_with_dict = pd.DataFrame({
    "a_and_b": df.apply(lambda x: {"a": x["a"], "b": x["b"] }, axis=1),
    "c":  df["c"]
})
pa.Table.from_pandas(df_with_dict , SCHEMA)

从箭头转换回熊猫时，结构表示为dict：

pa.Table.from_pandas(df_with_dict , SCHEMA).to_pandas()['a_and_b']
| a_and_b          |
|:-----------------|
| {'a': 1, 'b': 2} |
| {'a': 4, 'b': 5} |
| {'a': 7, 'b': 8} |

【讨论】：