Pandas.DataFrameのメモリサイズを削減する（最大で8分の1） [Python]

はじめに

Pandasで巨大なデータを扱うと、貧弱なPCではすぐメモリエラーになるのではないでしょうか。
これまで結構苦労したので、Pandasでメモリ消費を抑えるコツを挙げておきます。
DataFrameについて書きますが、Seriesも同様です。Panelは触ったことないですが、きっと同様でしょう。多分。

使用した環境
Python 3.6
Pandas 0.20.3

メモリが必要以上に増大してしまうケース

いろんな場合がありますが、以下のケースは、よくあるかつコードで対処可能なものだと思います。

【ケース1】 DataFrame構築時にカラムの型(dtype)を指定していない
【ケース2】 カラムを追加する・カラム全体へ値を代入する
【ケース3】 DataFrameに対する処理の戻り値を、他の変数で受け取る

詳しく見ていきます。

【ケース1】 DataFrame構築時にカラムの型(dtype)を指定していない

DataFrame構築時にカラムの型(dtype)を指定していないと、整数はint64 、小数はfloat64 が勝手に割り当てられます。どんな値も扱えるように、とにかく大きなサイズの型になっています。

input.csv

id,transaction_date,group,value1,value2
1001,200612,1,23,3.1
1002,200630,2,83,9.8
1003,200703,3,24,8.7
1004,200707,8,31,4.8

df = pd.read_csv('input.csv')
print(df.dtypes)
print(df.memory_usage(index=False))

# 出力。
id                    int64
transaction_date      int64
group                 int64
value1                int64
value2              float64
dtype: object

id                  32
transaction_date    32
group               32
value1              32
value2              32
dtype: int64                       # 合計160バイト

※データサイズは、適宜 df.memory_usage(index=True).sum() / 1024 ** 2 (KBの場合)の様に、 1,024ベースで計算してください。
全部 int64, float64になってしまってますね。

参考：代表的な型のサイズ

型	サイズ	値の範囲
int8	1バイト	-128 to 127
int16	2バイト	-32768 to 32767
int32	4バイト	-2147483648 to 2147483647
int64(デフォルト)	8バイト	-9223372036854775808 to 9223372036854775807
float16	2バイト	半精度浮動小数 sign bit, 5 bits exponent, 10 bits mantissa
float32	4バイト	単精度浮動小数 sign bit, 8 bits exponent, 23 bits mantissa
float64(デフォルト)	8バイト	倍精度浮動小数 sign bit, 11 bits exponent, 52 bits mantissa
https://docs.scipy.org/doc/numpy-1.13.0/user/basics.types.html

対策

DataFrame構築時に、dtypeオプションでカラムの型を指定する

# 変換したいカラムの型を、ディクショナリで指定する
dtyp = {'id': 'int16', 'transaction_date': 'int32', 'group': 'int8', 
        'value1': 'int8', 'value2': 'float16'}
df = pd.read_csv('input.csv', dtype=dtyp)
print(df.dtypes)
print(df.memory_usage(index=False))

# 出力
id                    int16
transaction_date      int32
group                  int8
value1                 int8
value2              float16
dtype: object

id                   8
transaction_date    16
group                4
value1               4
value2               8
dtype: int64                    # 合計48バイト。3分の1以下になった。

pd.read_csv()では、dtypeオプションで「カラム名: 型」のディクショナリを渡せば、カラムごとに型を指定できます。変換したいカラムだけ指定すれば良いです。
pd.DataFrame()では、dtypeはディクショナリでの個別指定は出来ず、dtype='int8' の様に一律同じ型の指定しか出来ません。
1〜2桁しか使わない整数項目などでは、int64からint8に変換してやると、サイズが一気に８分の１になります。

DataFrame構築後に、個別に型変換する

df['value1'] = df['value1'].astype('int8')

カテゴリ型

category型というのもあります。コード値のような、特定の離散値を取る性質のデータに使います。これは以下のようなメリットがあり、積極的に活用したい型です。

• コードのユニーク数（例えば性別なら、2とか3）が少ないなら、メモリ使用量はかなり減る。
• 処理速度が改善される。特にGroupBy は、Pandasの内部処理として整数型の配列で扱うため、整数型ではない列の処理は効果が大きいみたい。
• catアクセサを使って、ラベルエンコーディング（0から始まる整数値に変換）したり、それを元の値に戻したりすることが簡単に出来る。(参考 Pandas公式ドキュメント）
• LightGBMでは、pandasのカテゴリ型を認識してくれるので、 categorical_feature を自分でわざわざ設定したり、One-hotエンコーディングをしたりする必要がない。

LightGBM公式ドキュメントより LightGBMのcategorical_feature について

categorical_feature (list of strings or int, or 'auto', optional (default="auto")) – Categorical features. If list of int, interpreted as indices. If list of strings, interpreted as feature names (need to specify feature_name as well). If ‘auto’ and data is pandas DataFrame, pandas categorical columns are used.

型指定する際の注意

精度を超える値で更新すると、エラーにもならずに変な値になるので、後で演算・更新をする数値項目なのか、コード値のようなカテゴリ項目なのかは、ちゃんと考えて使いましょう。

【ケース2】 カラムを追加する・カラム全体へ値を代入する

追加された数値カラムは、勝手に int64, float64 になります。
カラム全体に対して新たな値を代入した場合も、すでに型指定をしていても、またint64, float64 になります。

# 部分代入では型は変わらない
df.loc[:50, 'value1']  = 99
# 全体への代入は、カラムの追加と同じく勝手に `int64`, `float64` になる。
df.loc[:, 'value1']  = 99
df.loc['value2']  = 1.1

対策

後から型変換する

df['value1'] = df['value1'].astype('int8')

【ケース3】 DataFrameに対する処理の戻り値を、他の変数で受け取る

Pandasはコピーを返すのが基本なので、以下の場合、df1 はそのままで、別のオブジェクトdf2 が出来ます。無駄なオブジェクトを放置していないか、注意しましょう。

# こんなのとか。
df2 = df1.fillna(0)
# スライスも。
df2 = df1.iloc[:, 1:]

対策

df1を後で使用せず、単に更新したいだけならこうすれば良いです。
元のオブジェクトを上書きする

df1 = df1.iloc[:, 1:]

inplace=True として、直接元のオブジェクトを更新する
戻り値はありません。

df1.fillna(0, inplace=True)

ちなみに、pandasメソッド内部の処理過程で、処理対象のDataFrameと同程度のメモリが一時的に別途消費されると考えたほうが良いと思います。inplaceの場合でも。（メソッドによって程度は違うと思いますが）

df1をとっておきたい場合でも、
元のオブジェクトが不要になったら、その時点で削除する

del df
# 複数まとめて実行するならこう。
del df1, df2, .....

さらに念の為に、delの後にセットでgc.collect() でガベージコレクションを強制実行する人もいますね。
参考: ガベージコレクタインターフェース

さいごに

DataFrameの型をまとめて最適化するモジュールを作りました。DataFrameを何も考えずに放り込むだけなので、らくちんです。良かったらご利用ください。pickleファイル出力の前に実行すると、出力ファイルのサイズを減らせます。
ただ、前述の通り、精度を超える値で更新する可能性がある場合 は要注意です！
https://github.com/nannoki/myutil/blob/master/mem_shrinker.py