Как сериализовать пользовательские классы как структуры, используя pyarrow в dask dataframes?

Question

У меня есть dask dataframe, в котором есть столбец типа List [MyClass]. Я хочу сохранить этот фрейм данных в файлы паркета. 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: Error inferring Arrow data type for collection of Python objects. Got Python object of type MyClass but can only handle these types: 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

немного поздно, но, возможно, эта ссылка может помочь другим.

В основном все сводится к определению пользовательских функций сериализации, сделанных вручную. Например вот ваш класс:

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)