【python】pandasにおけるメモリ節約

はじめに

pythonのpandasにおけるメモリ節約の手法をいくつか紹介します。
一番言いたいことは「NaNを追加するとデータ型が自動的にfloat64に変更されてメモリが一気に倍増とかあるから気を付けて」ということです。その前準備としてデータ型の話から入ります。

列のデータ型に留意する
NaNの発生に注意
（おまけ1）使わないデータは読み込まない
（おまけ2）使わなくなったらメモリを解放する

列のデータ型に留意する

例えばint32は4バイト、int64は8バイトのメモリを確保します。レコード数、列数が増えれば大きな違いとなってコンピュータのメモリ容量を圧迫します。値の範囲がわかっていれば、できるだけバイト数の小さいデータ型を選択することをお勧めします。

実際にpandasのデータフレームを作成してメモリ使用量を見ていきます。その前準備として、データが入った2次元リストを作成します。

import numpy as np
import pandas as pd

# 数百メガバイトのメモリを使用します。
record_count = 10000000 # 1000万件
# メモリに余裕がない場合は、上の行ではなく以下の行を実行してください。
# record_count = 100000 # 10万件

# 3列、上で指定したレコード数の2次元リストを作成します。
data = [[i for i in range(3)] for j in range(record_count)]

# 確認のため、先頭から5レコードを表示します。
print(data[:5]) 
#実行結果
[[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]

このdataを使ってデータフレームを作成します。最初に列名のみ指定します。
「%%timeは」JupyterNotebookで実行時間を計測するために入れています。時間計測の目的がなければ不要です。

%%time
df = pd.DataFrame(data, columns = ['x', 'y', 'z'])
print(df.info()) # 基本情報表示

# 実行結果
RangeIndex: 10000000 entries, 0 to 9999999
Data columns (total 3 columns):
 #   Column  Dtype
---  ------  -----
 0   x       int64
 1   y       int64
 2   z       int64
dtypes: int64(3)
memory usage: 228.9 MB
None
Wall time: 5.06 s

データ型の指定がない場合は、上記の通り int64 になります。int32でも十分であれば不要なメモリを確保していることになります。

明示的にデータ型にint32を指定して指定してデータフレームを作成します。

%%time
df = pd.DataFrame(data, columns = ['x', 'y', 'z'], dtype = 'int32')
print(df.info())

#実行結果
RangeIndex: 10000000 entries, 0 to 9999999
Data columns (total 3 columns):
 #   Column  Dtype
---  ------  -----
 0   x       int32
 1   y       int32
 2   z       int32
dtypes: int32(3)
memory usage: 114.4 MB
None
Wall time: 5.03 s

メモリ使用量が 228.9 MB ⇒ 114.4MB と半分になりました。

dataをpd.DataFrameに渡す前に np.arrayに変換すると実行時間が短くなります。この場合、データ型はデフォルトでは4バイトになります。

%%time
df = pd.DataFrame(np.array(data), columns = ['x', 'y', 'z'])
print(df.info())

#実行結果
RangeIndex: 10000000 entries, 0 to 9999999
Data columns (total 3 columns):
 #   Column  Dtype
---  ------  -----
 0   x       int32
 1   y       int32
 2   z       int32
dtypes: int32(3)
memory usage: 114.4 MB
None
Wall time: 2.95 s

4バイト以外のデータ型を指定する場合はdtypeプロパティを使用します。しかし、3つそれぞれ別の型を指定することはできません。それぞれ別の型を指定する場合は、データフレームの作成後にastypeメソッドで以下のように指定します。

df[列名] = df[列名].astype('データ型')

NaNの発生に注意

これがこの記事で一番強調したいところです。列の中に1つでもNaNが入ると、それまでのデータ型は無視されて自動的にfloat64になってしまいます。これはメモリ節約の観点では大きなインパクトです。以下に具体例を示します。

上で作ったデータフレームに1レコード[0, 1, np.nan] を追加します。レコードの追加方法はここではpd.concatメソッドを使います。

#追加先(df1)のデータ型とメモリ使用量確認
print(df.info()) 
RangeIndex: 10000000 entries, 0 to 9999999
Data columns (total 3 columns):
 #   Column  Dtype
---  ------  -----
 0   x       int32
 1   y       int32
 2   z       int32
dtypes: int32(3)
memory usage: 114.4 MB

# 追加用レコード作成（1レコードからなるデータフレームdf2）
df2 = pd.DataFrame([[0,1,np.nan]], columns = ['x', 'y', 'z'], dtype = 'int32')

# df(1000万レコード)にdf2(1レコード)を追加
df = pd.concat([df, df2]).reset_index(drop = True)

# 追加後のデータフレームのデータ型とメモリ使用量確認
df.info()
RangeIndex: 10000001 entries, 0 to 10000000
Data columns (total 3 columns):
 #   Column  Dtype  
---  ------  -----  
 0   x       int32  
 1   y       int32  
 2   z       float64
dtypes: float64(1), int32(2)
memory usage: 152.6 MB

1000万レコードからたった1レコード増えただけなのに、メモリ使用量は大幅に増加してしまいました。列zのデータ型が4バイトから8バイトに変わったことが原因で、この該当列に限ればメモリ使用量は倍増しているわけです。今回NaNが入ったのが1列、1000万レコードなので1000万レコード×(8-4)バイト = 40メガバイト程度の増加ですが、1億レコードあったら400メガバイトです。はたして1つのNaNを保持するためにほぼ使わない大容量のメモリが必要でしょうか？

pandasでデータ分析をやっているとNaNが便利なこともあります。しかしpandasでデータ分析をやっている人ならば、1000万を超えるレコードのデータに出会ったり、1000列を超えるデータにさらに列を追加したりすることもあるでしょう。そのような時には本当にNaNである必要があるのか？-1や9999に置き換えられないか？などを検討すると良いかもしれません。（他のデータも普通にfloat64の列であれば変わりませんが）

ちなみに、NaNを-1に置き換えて良いという判断の場合、以下のように変換してメモリを節約できます。

df3.loc[pd.isnull(df3['z']), 'z'] = -1
df3['z'] = df3['z'].astype('int32')
df3.info()

#実行結果
RangeIndex: 10000001 entries, 0 to 10000000
Data columns (total 3 columns):
 #   Column  Dtype
---  ------  -----
 0   x       int32
 1   y       int32
 2   z       int32
dtypes: int32(3)
memory usage: 114.4 MB

メモリ使用量が元のサイズに戻りました。様々な処理をしているうちにNaNが発生していることがありますので、私はデータフレームのinfoメソッドをよく使います。

（おまけ1）使わないデータは読み込まない

入力するcsvに使わない列があることがわかっている場合、使う列のみを指定して読み込むことで無駄なメモリ使用を回避できます。dtypeパラメータに辞書を渡すことでデータ型の指定ができます。dtypeを指定しないと、intはint64に、floatはfloat64になります。

df = pd.read_csv('data.csv', usecols = ['x', 'y'], dtype = {'x':'int32', 'y':'int64'})

「使うか使わないかは見てみないとわからない。そして、データを全部読み込んで見るのは時間がもったいないし、最初のレコードで判断できる」状況であれば、nrows = 読み込むレコード数で任意のレコード数を読み込むことができます。

df = pd.read_csv('data.csv', nrows = 5)

データの範囲から自動的に必要なバイト数を鑑みて
int16/int32/int64、float16/float32/float64を選択してくれる関数がkaggle内で発表されています。以下に解説付きの素晴らしい記事を紹介します。
めっちゃ使えるpandasのメモリサイズをグッと抑える汎用的な関数

（おまけ2）使わなくなったらメモリを解放する

pandasに限らず、使わなくなったオブジェクトは削除してガベージコレクションを行い、メモリを解放しましょう。

import gc
del df 
gc.collect()

2020/12/5
誤字脱字を修正しました。