Serialization / Deserialization of ExtensionArrays

[TomAugspurger]

(just getting an issue number to link to. Will update later)

What hooks do we want to provide for ExtensionArrays stored inside our containers?

Currently, to_csv works. Haven't really tried others

• pickle
• msgpack
• parquet
• HDF5
• json

xref https://github.com/pandas-dev/pandas/pull/20611/files

[jreback]

wouldn’t support msgpack generally

[jorisvandenbossche]

For pickling, ensuring your ExtensionArray picklable should be enough (the rest already works, the blocks pickle the values). For geopandas we have this:

https://github.com/geopandas/geopandas/blob/9cb93940bc2ac94c5d02651fe8a4fa50220827f2/geopandas/array.py#L135-L141

[jorisvandenbossche]

I am interested in discussing the parquet / arrow aspect of this issue (see eg also apache/arrow#4168, and I would also like to potentially use this in GeoPandas).

In general, I think it would be useful to have an interface to convert an ExtensionArray to an arrow array, and the other way around.
Eg our IntegerArray EA would map perfectly to the arrow integer arrays. Having such well defined conversion can help in supporting EAs in to_parquet(at least for the pyarrow engine), and might also be useful to convert in general to arrow data structures (eg fletcher might find this useful).

Personally, I think this responsibility should be on the ExtensionArray implementation itself (pyarrow should not need to know about all EAs, but only about a specific interface).

So eg we could add _to_arrow / _from_arrow methods to the EA interface.

def _to_arrow(self, type=None) -> pyarrow.Array:
    # default implementation (currently already works for basic cases)
    return pyarrow.array(self, type=type)

@classmethod
def _from_arrow(cls, array : pyarrow.Array) -> ExtensionArray:
    return ...

That would allow pyarrow to call those methods when receiving an ExtensionArray, and in principle (if they want that) even restore it when converting back to pandas if they store the extension dtype name in the pandas metadata (with something like pandas_dtype(stored_dtype_name).construct_array_type()._from_arrow(...))

[TomAugspurger]

Seems reasonable to reuse arrow where it has readers / writers.

Do we foresee problems with overloading to_arrow for both conversion to arrow and parquet IO? I don’t really see any problems right now.

[jorisvandenbossche]

Do we foresee problems with overloading to_arrow for both conversion to arrow and parquet IO? I don’t really see any problems right now.

One case I can think of is if you would have a custom data type that maps nicely to a certain arrow type (eg ragged arrays / multipolygons as lists of lists), but which is not supported in the specific IO option (eg the current arrow <-> parquet support for nested data is limited), so you might want to choose a suboptimal one for a certain IO option (eg binary blob instead).
But, 1) I am not sure we need to care about such case right now and 2) that might be solvable by being able to specify the resulting type (note the type argument I have in _to_arrow above, just like __array__(self, dtype=None) in numpy's coercing) in the conversion, and you could then specify the required resulting schema when doing the conversion to arrow (eg might want to add a schema keyword to to_parquet anyway, regardless of this issue).

[jorisvandenbossche]

To have it less pandas EA-specific, arrow could also look for a __arrow_array__ method instead of _to_arrow.