How to serialize custom classes as structs using pyarrow in dask dataframes?

Question

I have dask dataframe that has a column of type List[MyClass]. I want to save this dataframe to parquet files. Dask is using pyarrow as the backend, but it supports only primitive types.

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')

So when i try to save it i get this error:

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.

Obviously i have to convert MyClass to pyarrow compatible struct type, but i can't find a way how to do this. Pyarrow & dask have some serialization features (like this https://arrow.apache.org/docs/python/ipc.html#serializing-custom-data-types), but seems like that's not quite the thing i need.

Answer 1

a bit late, but maybe this link can help others.

It basically comes down to defining custom hand-made serialization functions. For example here's your class:

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

You write functions to convert to/from this class like:

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

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

Then initialize a context from these functions to later give to the Serialization/Deserialization methods:

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

Now you call the serialize/deserialize methods and pass them the context:

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