_util.py

from __future__ import annotations

from typing import Callable

from pandas.compat._optional import import_optional_dependency

import pandas as pd


def _arrow_dtype_mapping() -> dict:
    pa = import_optional_dependency("pyarrow")
    return {
        pa.int8(): pd.Int8Dtype(),
        pa.int16(): pd.Int16Dtype(),
        pa.int32(): pd.Int32Dtype(),
        pa.int64(): pd.Int64Dtype(),
        pa.uint8(): pd.UInt8Dtype(),
        pa.uint16(): pd.UInt16Dtype(),
        pa.uint32(): pd.UInt32Dtype(),
        pa.uint64(): pd.UInt64Dtype(),
        pa.bool_(): pd.BooleanDtype(),
        pa.string(): pd.StringDtype(),
        pa.float32(): pd.Float32Dtype(),
        pa.float64(): pd.Float64Dtype(),
    }


def arrow_string_types_mapper() -> Callable:
    pa = import_optional_dependency("pyarrow")

    return {
        pa.string(): pd.StringDtype(storage="pyarrow_numpy"),
        pa.large_string(): pd.StringDtype(storage="pyarrow_numpy"),
    }.get