时间戳

首页 > 解决方案 > 使用 pyarrow 将结构列写入镶木地板

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

问题描述

我有以下数据框和架构：

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 中的结构类型相对应的正确类型的方法。有没有办法做到这一点？

标签： pythonpandasdataframeparquetpyarrow

解决方案

对于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)

从箭头转换回熊猫时，struct 表示为 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} |

推荐阅读