DataFrameの型指定をしてサイズを削減する手順を自動化

以前の記事でDataFrameの各列の型を指定してサイズを小さく出来ると書いたが、それを自動化する。

anton0825.hatenablog.com

以下を使ってcsvを読み込むとサイズを自動的に削減してくれる。一度全部読み込んでから型変換しているので、 処理途中で使うメモリ量は増える。csvはこれを使ってDataFrameに変換した後、Feather形式で保存し以降はそれを使うようにするといい。

def read_csv(filename: str, verbose: bool = True) -> pd.DataFrame:
    # Seriesの型を最適なものにすることでメモリを削減する。
    # float16はFeather形式で保存出来ないため使用しない。
    df: pd.DataFrame = pd.read_csv(f"{INPUT_PATH}/{filename}")
    start_mem: float = df.memory_usage().sum() / 1024 ** 2
    for col in df.columns:
        _set_optimized_type(col, df)
    end_mem: float = df.memory_usage().sum() / 1024 ** 2
    _print_result(end_mem, start_mem, verbose)
    return df


def _print_result(end_mem: float, start_mem: float, verbose: bool):
    if verbose:
        reduce_percent: float = 100 * (start_mem - end_mem) / start_mem
        print('Mem. usage decreased to {:5.2f} Mb ({:.1f}% reduction)'.format(end_mem, reduce_percent))


def _set_optimized_type(col: str, df: pd.DataFrame):
    numerics: List[str] = ['int16', 'int32', 'int64', 'float32', 'float64']
    col_type: List[str] = df[col].dtypes
    if col_type in numerics:
        if str(col_type)[:3] == 'int':
            df[col] = df[col].astype(_to_optimized_int_type(df[col].min(), df[col].max()))
        else:
            df[col] = df[col].astype(_to_optimized_float_type(df[col].min(), df[col].max()))


def _to_optimized_int_type(c_min: int, c_max: int):
    if np.iinfo(np.int8).min < c_min and c_max < np.iinfo(np.int8).max:
        return np.int8
    elif np.iinfo(np.int16).min < c_min and c_max < np.iinfo(np.int16).max:
        return np.int16
    elif np.iinfo(np.int32).min < c_min and c_max < np.iinfo(np.int32).max:
        return np.int32
    elif np.iinfo(np.int64).min < c_min and c_max < np.iinfo(np.int64).max:
        return np.int64


def _to_optimized_float_type(c_min: float, c_max: float):
    if np.finfo(np.float32).min < c_min and c_max < np.finfo(np.float32).max:
        return np.float32
    else:
        return np.float64

参考：

qiita.com