Pandasのメモリ削減方法を整理した

皆さんこんにちは
機械学習チーム YAMALEXチームの@tereka114です。最近、寒いので、鍋を中心に食べて生きています。
検証段階でも、規模の大きなデータを扱う機会が増えてきて、Pandasのメモリ消費量が厳しいと感じてきたので、その削減や効率化のテクニックまとめたいと思いました。
有名なものからマイナーなものまで、思いつく限り書いてみます。

そもそもなぜ、Pandasのメモリ削減技術が必要なのか

Pandasで扱うデータの多くのファイルはCSV,Parquet, JSON（JSONL）になります。
これらのファイルは扱いやすいこともあり、頻繁に利用されています。
しかし、特にログデータを扱う場合、10GBを超えるなどファイルの容量が膨らみ、所謂、普通のマシンで解析する場合、メモリに乗り切らずに実装が困難になります。

準備

まずは、データを準備します。
次のコードで実装します。1000万レコードほどのデータを作成します。

import pandas as pd
import numpy as np

N = 10000000
df = pd.DataFrame({
    "id": [i for i in range(N)],
    "user_id": np.array(["user_0", "user_1", "user_2", "user_3","user_4", "user_5"])[np.random.randint(0,6, N)],
    "category_id": np.random.randint(0, 10, N),
    "use": np.array([True, False])[np.random.randint(0,2, N)],
    "sales": np.random.randint(0, 1000000, N),
    "rate": np.random.rand(N),
})
df.to_csv("./sample_data.csv", index=False)

この結果を次の方法で読み込み、メモリを計測します。
計測結果は391MBの消費量ですので、ここからメモリを削減するといったことを試みます。

df = pd.read_csv("./sample_data.csv")
df.memory_usage().sum() / 1024**2 #  391.006591796875

Pandasのメモリ削減

先程準備したファイルを元にPandasのメモリを削減しながら処理をする方式を紹介します。

1. 型修正

Pandasで読み込んだままのメモリだと、int64やfloat64が最初に使われるので効率が悪いです。
そのため、最大、最小の値を利用して、int8やfloat32など、よりメモリを消費しない型に変換することも一つの手段です。

この削減は非常に有名な実装があるので、貼っておきます。この実装を適用することで133MBまで消費量が削減されます。
この実装の内部では各カラムの値の最大、最小を計算し、データの精度を落とさず、最もメモリを消費できる型変換をかけます。
ただし、booleanの自動変換はできないので、自分で、.astype(bool)などで変換しましょう。

def reduce_mem_usage(df):
    """ iterate through all the columns of a dataframe and modify the data type
        to reduce memory usage.        
    """
    start_mem = df.memory_usage().sum() / 1024**2
    print('Memory usage of dataframe is {:.2f} MB'.format(start_mem))
    
    for col in df.columns:
        col_type = df[col].dtype
        
        if col_type != object:
            c_min = df[col].min()
            c_max = df[col].max()
            if str(col_type)[:3] == 'int':
                if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:
                    df[col] = df[col].astype(np.int8)
                elif c_min > np.iinfo(np.int16).min and c_max < np.iinfo(np.int16).max:
                    df[col] = df[col].astype(np.int16)
                elif c_min > np.iinfo(np.int32).min and c_max < np.iinfo(np.int32).max:
                    df[col] = df[col].astype(np.int32)
                elif c_min > np.iinfo(np.int64).min and c_max < np.iinfo(np.int64).max:
                    df[col] = df[col].astype(np.int64)  
            else:
                if c_min > np.finfo(np.float16).min and c_max < np.finfo(np.float16).max:
                    df[col] = df[col].astype(np.float16)
                elif c_min > np.finfo(np.float32).min and c_max < np.finfo(np.float32).max:
                    df[col] = df[col].astype(np.float32)
                else:
                    df[col] = df[col].astype(np.float64)
        else:
            df[col] = df[col].astype('category')

    end_mem = df.memory_usage().sum() / 1024**2
    print('Memory usage after optimization is: {:.2f} MB'.format(end_mem))
    print('Decreased by {:.1f}%'.format(100 * (start_mem - end_mem) / start_mem))
    
    return df

df.memory_usage().sum() / 1024**2 #133.51

2. 逐次読み込み

マシンによってはそもそも巨大なデータは読み込めないことがあります。
Pandasには、巨大データでも読み込むために逐次ファイルから読み込む機能があります。
前述の型修正と組み合わせることで、メモリの削減をしつつ、全てのファイルを読み込めます。

reader = pd.read_csv("./sample_data.csv", chunksize=1000000)
concat_df = pd.concat([reduce_mem_usage(part_df) for part_df in reader])
concat_df.memory_usage().sum() / 1024**2 #133.51

3. 読み込み時の型指定

2の別解です。ファイルから分割して読めるといってもやはり結合は面倒です。
ファイル読み込み時にdtypeを指定する方法があります。予め型わかっている場合はこちらを利用した方が効率的です。

df = pd.read_csv("./sample_data.csv", dtype={
    "id": "int32",
    "user_id": "category",
    "category_id": "int8",
    "use": "bool",
    "sales": "int32",
    "rate": "float16"
})
df.memory_usage().sum() / 1024**2 # 123.97

4. 逐次読み込み＆集約

メモリに載りきらない膨大なデータはログや行動記録であることも多いです。
ログは、例えば、ユーザごとに集約すれば、ユーザ分のみのデータを作成できます。
実際の用途として、特徴量を作成する際に分割統治法の考え方を用いて集約すれば、メモリを節約できます。

user_idを元にsales, rateを合計集約するようなケースを考えましょう。
その場合、読み込んだチャンクごとに集約を行い、最後に部分ごとの合計を合算しても、結果はまとめて計算するのと変わらないので次の実装で計算ができます。
ただし、全ての場合（NG例：平均の計算）で可能ではないので、計算が正しくなるか検証してからにしましょう。

reader = pd.read_csv("./sample_data.csv", chunksize=1000000)
concat_df = pd.concat([part_df.groupby("user_id")["rate", "sales"].sum().reset_index() for part_df in reader])
concat_df.groupby("user_id")["rate", "sales"].sum()

5. 不要なものを読み込まない

そもそも全てのカラムが必要ないケースもあります。
例えば、user_idとsalesで、user_idごとのsalesの平均を計算する場合を考えます。
その場合、user_id, sales以外のカラムは不要になるので、そもそも読み込む必要はありません。

part_df = pd.read_csv("./sample_data.csv", usecols=['user_id', 'sales'])
part_df = reduce_mem_usage(part_df)
part_df.groupby("user_id")["sales"].mean()
concat_df.memory_usage().sum() / 1024**2 #48.684

6. 不要なカラム／DataFrameを消す

プログラムの最後まで全てのカラムが必要なケースは少ないと思います。
そのため、基本ではありますが、不要になったカラムやDataFrameは消しましょう。

# use列消去
del df["use"] 
df.memory_usage().sum() / 1024**2 # 113.44

# dfそのものを消す
del df
import gc
gc.collect()

番外編：そもそもPandasを利用しない

ここまで、Pandasのメモリ削減の話をしていましたが、そもそもPandasを利用しないといった方法があります。
最近だとPolarsと呼ばれるPandasとはAPIは異なりますが、所謂DataFrame系の構造操作ができるライブラリがあります。
Pandasと比較して、高速、メモリが省メモリ、並列処理など様々な恩恵があり、データによってはPolarsを利用するほうが良いかもしれません。

www.pola.rs

最後に

本日はPandasで巨大なデータを扱うためのメモリ削減の方法に関して紹介しました。
有名なものから、自分なりの工夫も紹介してみましたのでぜひ使って良いPandasライフをお過ごしください。
では、また年明けたら会いましょう。

Acroquest Technologyでは、キャリア採用を行っています。

• ディープラーニング等を使った自然言語／画像／音声／動画解析の研究開発
• Elasticsearch等を使ったデータ収集／分析／可視化
• マイクロサービス、DevOps、最新のOSSを利用する開発プロジェクト
• 書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長

　
少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。

www.wantedly.com