Saving pint pandas dataframe with types

gistfile1.txt

from typing import Dict

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import json


def session_data_to_parquet(fname: str, header: Dict, data: pd.DataFrame) -> None:
    index_col = data.index.name
    new_data = data.reset_index()

    pa_fields = []
    for col, dtype in new_data.dtypes.items():
        pa_fields.append(
            pa.field(col, type=pa.float64(), metadata={
                "dtype": str(dtype)
            })
        )

    new_schema = pa.schema(
        pa_fields,
        metadata={
            "daq_tools": json.dumps({
                "index": index_col,
                "header": header
            })
        }
    )

    table = pa.Table.from_pandas(
        new_data.astype(float),
        new_schema
    )
    print(index_col, index_col in table.schema.names)
    pq.write_table(table, fname)


def session_data_from_parquet(fname: str) -> [Dict, pd.DataFrame]:
    table = pq.read_table(fname)
    meta = json.loads(table.schema.metadata[b"daq_tools"].decode("utf8"))
    index_col = meta["index"]
    dtypes = {}
    for name in table.schema.names:
        dtypes[name] = table.schema.field(name).metadata[b"dtype"].decode("utf8")
    data = table.to_pandas().astype(dtypes).set_index(index_col)
    return meta["header"], data