以前の記事でDataFrameの各列の型を指定してサイズを小さく出来ると書いたが、それを自動化する。
以下を使って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
参考: