python - 如何使用 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 保存到 parquet，让它解释哪个列需要哪个 Extension .将来我可能会更改列，并且我有几种不同的类型需要这种处理。

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

我试图创建一个 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

如果我应该使用 ExtensionType，我会迷失在 pyarrow 的文档中, serialization或其他东西来编写这些函数。任何指针将不胜感激。

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

Answer 1

我认为这可能是因为 'ObjectId' 不是 python 中定义的关键字，因此它在类型转换中抛出了这个异常。
我尝试了您提供的示例，并通过在数据帧创建期间将 oid 值转换为字符串类型进行了尝试，并且成功了。
检查以下步骤:

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

df.to_parquet('parquet_file.parquet')

df1 = pd.read_parquet('parquet_file.parquet',engine='pyarrow')
df1

输出:

    name    oid
0   alice   ObjectId('5e9992543bfddb58073803e7')
1   bob ObjectId('5e9992543bfddb58073803e8')