Как сохранить pandas DataFrame с пользовательскими типами, используя pyarrow и parquet - PullRequest
Новая
       85

Как сохранить pandas DataFrame с пользовательскими типами, используя pyarrow и parquet

0 голосов
/ 17 апреля 2020

Я хочу сохранить pandas DataFrame для паркета, но у меня есть некоторые неподдерживаемые типы (например, 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()

И смог получить рабочий pyarray для моего oid столбца: 

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)

Теперь, где я нахожусь застрял, и не может найти никакого хорошего решения для inte rnet, как сохранить мой df на паркет, позволяя ему интерпретировать, какой столбец нуждается в Extension. Я мог бы изменить столбцы в будущем, и у меня есть несколько различных типов, которые нуждаются в этой обработке.

Как я могу просто создать файл паркета из кадров данных и восстановить их при прозрачной конвертации типов?

Я попытался создать объект 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

Я теряюсь при включении пиарроу c, если я должен используйте ExtensionType, serialization или другие вещи для написания этих функций. Любой указатель был бы признателен.

Примечание: мне не нужно ни в коем случае parquet, основная проблема заключается в возможности сохранять и восстанавливать кадры данных с пользовательскими типами quickly и space efficiently. Я попробовал решение, основанное на jsonification и gziping кадре данных, но оно было слишком медленным.

...