Pythonでジェネレータを利用して大規模CSVデータを効率的に読み込む方法

最近、20GBを超える複数のCSVファイルを取り扱う機会が増えてきました。しかし、Pythonでそれらのファイルをそのまま読み込むとメモリ負荷が非常に大きく（最大30GB程度）、処理時間も20分以上かかってしまう問題に直面していました。
そんな中、ジェネレータを利用することでメモリ使用量を7GBまで抑えることに成功し、処理時間も短縮できたので、その際に行ったことについて紹介していきます。

CSVファイルの初期読み込み

まず最初に行ったのは、以下のようなコードでCSVファイル全体を一度に読み込む方法です。

csv_file = "ファイルパス"

# 全体を一気に読み込む
df = pd.read_csv(csv_file)
df.head()

しかし、この方法では巨大なファイルを全てメモリ上にロードするため、メモリ消費が非常に大きく、特にメモリ容量が限られている環境では現実的ではありませんでした。

ジェネレータとは？

ジェネレータの基本

ジェネレータは、Pythonにおける特殊なイテレータです。通常の関数とは異なり、yield文を使って関数の実行を一時停止し、必要なときに再開する仕組みを持っています。これにより、関数が全ての結果を一度に返すのではなく、1つずつ結果を「生成」しながら返します。この特性により、大量のデータを一度に処理するのではなく、リアルタイムで処理ができ、メモリの使用効率が劇的に向上します。

通常の関数との違い

通常の関数は全ての計算が完了した後に結果を一度に返しますが、ジェネレータはyieldを使うことで、計算の途中で一時停止し、その時点までの結果を返します。これにより、すべてのデータをメモリに保持する必要がなく、逐次的なデータ処理が可能になります。
この特性は、特に大規模データのストリーミング処理や、ファイルの逐次読み込みといったリアルタイムデータ処理に非常に有効です。

ジェネレータを使った大規模データの効率的な読み込み

メモリ使用量の削減

従来のデータ処理では、全データを一度にメモリに読み込むため、メモリを大量に消費します。しかし、ジェネレータを使うことで、必要なデータをその都度生成し、1行ずつ処理できるため、メモリ使用量を大幅に削減できます。特に、数GBを超える大規模CSVファイルを扱う際には、この方法が非常に有効です。

実際のコード例

以下のコードでは、200,000行ずつCSVファイルを分割して読み込み、必要な処理を逐次行っています。ジェネレータを使ってファイルを分割読み込みすることで、メモリ使用量を抑えつつ処理することが可能です。

import pandas as pd

csv_file = "ファイルパス"

# chunksizeを指定してジェネレータで分割読み込み
chunksize = 200000  # 一度に読み込む行数を設定
df_iter = pd.read_csv(csv_file, chunksize=chunksize)

chunks = []
for i, chunk in enumerate(df_iter):
    print(f"Chunk {i + 1}")
    print(chunk.head())  # 最初の5行を表示
    # 必要なデータ処理をここで行う
    chunks.append(chunk)

# リスト内のデータフレームを結合
df = pd.concat(chunks, ignore_index=True)

# データフレームの内容を表示
print(df.head())

結果

この方法を使った結果、メモリ消費を7GB程度に抑えることができ、さらに処理時間も20分から7分へと大幅に短縮することができました。

まとめ

ジェネレータを使うことで、大規模データの処理が非常に効率的になりました。メモリ消費の削減に加え、処理時間の大幅な短縮も実現できました。今後、さらに高度なジェネレータの使い方や、分散データ処理についても学んでいき、また新たな知見が得られた際には共有していきたいと思います。