问dask和熊猫数据中嵌套的numpy数组

EN

# these monkey-patches into the pandas and dask I/O API allow us to save multi-dimensional numpy
# arrays# in parquet format by serializing them into byte arrays

from dask import dataframe as dd
import pandas as pd
from io import BytesIO

def _patched_pd_read_parquet(*args, **kwargs):
    return _orig_pd_read_parquet(*args, **kwargs).applymap(
        lambda val: np.load(BytesIO(val)) if isinstance(val, bytes) else val)
_orig_pd_read_parquet = pd.io.parquet.PyArrowImpl.read
pd.io.parquet.PyArrowImpl.read = _patched_pd_read_parquet

def _serialize_ndarray(arr: np.ndarray) -> bytes:
    if isinstance(arr, np.ndarray):
        with BytesIO() as buf:
            np.save(buf, arr)
            return buf.getvalue()
    return arr

def _deserialize_ndarray(val: bytes) -> np.ndarray:
    return np.load(BytesIO(val)) if isinstance(val, bytes) else val

def _patched_pd_write_parquet(self, df: pd.DataFrame, *args, **kwargs):
    return _orig_pd_write_parquet(self, df.applymap(_serialize_ndarray), *args, **kwargs)
_orig_pd_write_parquet = pd.io.parquet.PyArrowImpl.write
pd.io.parquet.PyArrowImpl.write = _patched_pd_write_parquet

def _patched_dask_read_pyarrow_parquet_piece(*args, **kwargs):
    return _orig_dask_read_pyarrow_parquet_piece(*args, **kwargs).applymap(_deserialize_ndarray)
_orig_dask_read_pyarrow_parquet_piece = dd.io.parquet._read_pyarrow_parquet_piece
dd.io.parquet._read_pyarrow_parquet_piece = _patched_dask_read_pyarrow_parquet_piece

def _patched_dd_write_partition_pyarrow(df: pd.DataFrame, *args, **kwargs):
    return _orig_dd_write_partition_pyarrow(df.applymap(_serialize_ndarray), *args, **kwargs)
_orig_dd_write_partition_pyarrow = dd.io.parquet._write_partition_pyarrow
dd.io.parquet._write_partition_pyarrow = _patched_dd_write_partition_pyarrow