python - 如何在 dask 数据帧中使用 pyarrow 将自定义类序列化为结构？

Question

我有一个 dask 数据框，其中有一列类型为 List[MyClass]。我想将此数据框保存到 Parquet 文件中。 Dask 使用 pyarrow 作为后端，但它仅支持原始类型。

import pandas as pd
import dask.dataframe as dd


class MyClass:

    def __init__(self, a):
        self.a = a


def transform(v):
    return [MyClass(v)]


a = [[1], [2], [3]]
pdf = pd.DataFrame.from_dict(a)
ddf = dd.from_pandas(pdf, npartitions=1)
result = ddf.assign(mycol=ddf[0].apply(transform))
result.to_parquet('my_parquet.parquet')

所以当我尝试保存它时，我得到了这个错误:

ArrowInvalid:为 Python 对象集合推断 Arrow 数据类型时出错。获得了 MyClass 类型的 Python 对象，但只能处理这些类型:bool、float、integer、date、datetime、bytes、unicode、decimal。

显然，我必须将 MyClass 转换为与 pyarrow 兼容的结构类型，但我找不到执行此操作的方法。 Pyarrow 和 dask 具有一些序列化功能(如 https://arrow.apache.org/docs/python/ipc.html#serializing-custom-data-types )，但似乎这不是我需要的东西。

Answer 1

有点晚了，但也许这个link可以帮助别人。

基本上归结为定义自定义的手工序列化函数。例如，这是您的类(class):

class MyData:
    def __init__(self, name, data):
        self.name = name
        self.data = data

你编写函数来转换到/从这个类，如:

def _serialize_MyData(val):
    return {'name': val.name, 'data': val.data}

def _deserialize_MyData(data):
    return MyData(data['name'], data['data']

然后从这些函数中初始化一个上下文，以便稍后提供给序列化/反序列化方法:

context = pa.SerializationContext()
context.register_type(MyData, 'MyData',
                      custom_serializer=_serialize_MyData,
                      custom_deserializer=_deserialize_MyData)

现在调用serialize/deserialize 方法并将上下文传递给它们:

buf = pa.serialize(val, context=context).to_buffer()
restored_val = pa.deserialize(buf, context=context)