Pandas の chunksize

pandas ライブラリはデータを行と列からなる DataFrame として整理し、read_csv のようなリーダーは通常、ファイル全体を 1 つの DataFrame に読み込みます。ファイルが大きすぎて無理なく読み込めない場合、たとえば数百万行に及ぶ MovieLens の評価データ のエクスポートのようなデータセットでは、chunksize パラメーターを使うと、行のバッチに分けて読み込み、処理できます。

pandas での chunksize の仕組み

対応している pandas リーダーに chunksize=n を渡すと、その戻り値が変わります。DataFrame の代わりに、read_csv() は TextFileReader を返します。これはイテレータのようなオブジェクトで、バッチごとに 1 つの DataFrame を生成します。各バッチには最大 n 行が含まれ、行数がチャンクサイズのちょうど倍数でない場合は、最後のバッチの行数がそれより少なくなります。

このパターンを応用する際には、2 つの点が重要です。1 つ目は、chunksize が数えるのはバイト数ではなく行数だということです。すべてのチャンクがファイルの全列を保持し、変わるのは行数だけです。2 つ目は、このパラメーター自体では解析は速くならないということです。変わるのはメモリが割り当てられるタイミングなので、効果があるかどうかは、プログラムが各チャンクから何を保持するかに依存します。

CSV ファイルをチャンクで読み込む

ファイルをバッチで読み込むには、read_csv に chunksize を渡し、その結果をループで処理します。次の例では、小さな評価データ風の CSV を一時ファイルに書き出してから 3 行ずつ読み戻し、生成される各 DataFrame の shape を出力して、バッチ分割の様子を確認できるようにしています。

この例を実行してみましょう。

このセクションのコードと表示されるラベルや出力は英語のままです。

import os
import tempfile

import pandas as pd

rows = ["userId,movieId,rating"]
for i in range(7):
    rows.append(f"{i + 1},{100 + i},{1.0 + (i % 5) * 0.5}")

path = os.path.join(tempfile.mkdtemp(), "ratings.csv")
with open(path, "w") as f:
    f.write("\n".join(rows) + "\n")

reader = pd.read_csv(path, chunksize=3)
for number, chunk in enumerate(reader, start=1):
    print(f"chunk {number}: shape={chunk.shape}")

chunk 1: shape=(3, 3)
chunk 2: shape=(3, 3)
chunk 3: shape=(1, 3)

このファイルには 7 行のデータがあり、chunksize=3 なので、TextFileReader は完全な 2 つのチャンクに続いて 1 行の残りを生成します。どのチャンクでも列数は 3 のままで、バッチ分割が行にのみ適用されることが確認できます。

各チャンクを処理する

有用なパターンは、チャンク自体を集めるのではなく、チャンクごとに小さな結果、たとえば件数、合計、フィルタリングした部分集合などを累積していくことです。次の例では同じファイルを集計し、行数と評価値の合計を累積しながら処理するため、ファイル全体を一度にメモリに置く必要がありません。

この例を実行してみましょう。

このセクションのコードと表示されるラベルや出力は英語のままです。

import os
import tempfile

import pandas as pd

rows = ["userId,movieId,rating"]
for i in range(7):
    rows.append(f"{i + 1},{100 + i},{1.0 + (i % 5) * 0.5}")

path = os.path.join(tempfile.mkdtemp(), "ratings.csv")
with open(path, "w") as f:
    f.write("\n".join(rows) + "\n")

total_rows = 0
rating_sum = 0.0
for chunk in pd.read_csv(path, chunksize=3):
    total_rows += len(chunk)
    rating_sum += chunk["rating"].sum()

print(f"rows read: {total_rows}")
print(f"average rating: {rating_sum / total_rows:.2f}")

rows read: 7
average rating: 1.79

ループ内では、フィルタリング、列の選択、groupby による集約など、あらゆる DataFrame 操作が使えます。そのため、各チャンクを累積値に畳み込む前に変換できます。

チャンクサイズとメモリのトレードオフを理解する

チャンクで読み込むとピーク時のメモリ使用量を減らせますが、それは自動的には起こりません。ピーク使用量は、おおよそ最大チャンクのサイズにプログラムが保持する分を加えたものです。すべてのチャンクをリストに追加してあとで連結すると、結局データセット全体を保持することになり、通常の一括読み込みより多くのメモリを使う可能性さえあります。

値は 2 つの条件のバランスを取って選びます。chunksize を大きくするとループの反復回数が減り、反復ごとのオーバーヘッドも小さくなりますが、各チャンクが占有するメモリは増えます。小さくするとピークは下がりますが、そのぶんループの通過回数が増えます。また、最後のチャンクは通常短くなるため、固定の行数を前提としないループコードを書いてください。

他のデータソースでチャンク読み込みを使う

chunksize は read_csv 専用ではありません。このパラメーターを受け付ける他の pandas リーダーには、pandas.read_json(lines=True と組み合わせた場合)、pandas.read_sql、pandas.read_sql_table、pandas.read_stata、pandas.read_sas などがあります。pandas.read_hdf のように、代わりに同等の iterator モードを提供するリーダーもあります。対応状況や受け付けるオプションはリーダーと pandas のバージョンによって異なるため、使う予定の関数について公式ドキュメントを確認してください。

チャンク読み込みが使える場所なら、パターンはどこでも同じです。バッチを要求し、それをループで処理し、各バッチを小さな結果に畳み込みます。