如何使用 pyarrow 和 parquet 保存具有自定义类型的 Pandas DataFrame

Question

我想将 Pandas DataFrame 保存到 parquet，但我有一些不受支持的类型（例如 bson ObjectIds）。

在整个示例中，我们使用：

import pandas as pd
import pyarrow as pa

这是一个显示这种情况的最小示例：

df = pd.DataFrame(
    [
        {'name': 'alice', 'oid': ObjectId('5e9992543bfddb58073803e7')},
        {'name': 'bob',   'oid': ObjectId('5e9992543bfddb58073803e8')},
    ]
)

df.to_parquet('some_path')

我们得到：

ArrowInvalid: ('Could not convert 5e9992543bfddb58073803e7 with type ObjectId: did not recognize Python value type when inferring an Arrow data type', 'Conversion failed for column oid with type object')

我试图按照这个参考：https : //arrow.apache.org/docs/python/extending_types.html

因此我编写了以下类型扩展：

class ObjectIdType(pa.ExtensionType):

    def __init__(self):
        pa.ExtensionType.__init__(self, pa.binary(12), "my_package.objectid")

    def __arrow_ext_serialize__(self):
        # since we don't have a parametrized type, we don't need extra
        # metadata to be deserialized
        return b''

    @classmethod
    def __arrow_ext_deserialize__(self, storage_type, serialized):
        # return an instance of this subclass given the serialized
        # metadata.
        return ObjectId()

并且能够为我的oid专栏获得一个有效的 pyarray ：

values = df['oid']
storage_array = pa.array(values.map(lambda oid: oid.binary), type=pa.binary(12))
pa.ExtensionArray.from_storage(objectid_type, storage_array)

现在我被卡住了，在互联网上找不到任何好的解决方案，是如何将我的 df 保存到镶木地板，让它解释哪个列需要哪个Extension. 将来我可能会更改列，并且我有几种不同的类型需要这种处理。

如何简单地从数据帧创建镶木地板文件并在透明转换类型的同时恢复它们？

我试图创建一个pyarrow.Table对象，并在预处理后向它附加列，但它不能像使用table.append_column二进制列而不是那样工作，而且pyarrow.Arrays整个isinstance事情看起来像是一个糟糕的解决方案。

table = pa.Table.from_pandas(pd.DataFrame())
for col, values in test_df.iteritems():

    if isinstance(values.iloc[0], ObjectId):
        arr = pa.array(
            values.map(lambda oid: oid.binary), type=pa.binary(12)
        )

    elif isinstance(values.iloc[0], ...):
        ...

    else:
        arr = pa.array(values)

    table.append_column(arr, col)  # FAILS (wrong type)

理想解的伪代码：

parquetize(df, path, my_custom_types_conversions)
# ...
new_df = unparquetize(path, my_custom_types_conversions)

assert df.equals(new_df)  # types have been correctly restored

我迷失在pyarrow的，如果我要使用的文档ExtensionType，serialization或其他东西来写这些功能。任何指针将不胜感激。

旁注，我根本不需要parquet，主要问题是能够使用自定义类型quickly和space efficiently. 我尝试了基于 jsonifying 和 gziping 数据帧的解决方案，但速度太慢。

Answer 1

您可以编写一个方法来读取列名称和类型，并输出一个新的 DF，其中列转换为兼容类型，使用 switch-case 模式来选择将列转换为什么类型（或是否保留其原样）。