Pythonジェネレータでメモリ効率を高める方法

はじめに

業務で巨大なログファイルやマスターデータの大きいjsonファイルをプログラム上で読み込んだ時に大量のメモリが使用され、パフォーマンスが悪くなる事象が出ていました。
ネットで調べているとジェネレータを使用するとメモリの使用量が減らせることが分かりました。
実際、ジェネレータを使用しリファクタリングを行い検証を行ったところ性能が改善されました。
今回はジェネレータの使い方をまとめました。

ジェネレータとは

イテレータを反復して生成できる関数です。Pythonのlistのように大量のデータを一度にメモリに保持せずに、1度に1個ずつデータを生成しメモリに保持します。

具体例

ジェネレータを作成するにはreturnを使用せずにyeildを使用します

def sample_generator():
   yiled 1
   yiled 2
   yiled 3

gen = sample_generator()

ジェネレータはローカル変数を保持していて、関数を呼び出すと値が生成されます。

print(next(gen)) # 出力結果：1
# 処理Aを実行
print(next(gen)) # 出力結果：2
# 処理Bを実行
print(next(gen)) # 出力結果：3

send()関数を使用すると、ジェネレータから値を取得し、ジェネレータ関数に値を渡すことができます。

def generator_add():
    # 最初の値を受け取る
    x = yield
    print(f"最初に受け取った値: {x}")

    # 2つ目の値を受け取る
    y = yield
    print(f"次に受け取った値: {y}")

    # 受け取った値を足す
    yield x + y

# ジェネレータを生成
gen = generator_add()

# ジェネレータを初期化
next(gen)

# ジェネレータに最初の値を渡す
result = gen.send(3) # 出力結果→ 最初に受け取った値: 3

# ジェネレータに2つ目の値を渡す
result = gen.send(7) # 出力結果→ 次に受け取った値: 7

# 足し算(yield x + y)の結果を出力
print(result) # 出力結果→ 10

ジェネレータ式

ジェネレータ式はListに似ています。
Listは[]を使用しますが、ジェネレータ式は()を使用します。

0から9までの2乗したジェネレータを作成し、出力する処理です

gen = (x*x for x in range(10))
for i in gen:
    print(i)

# 0
# 1
# 4
# 9
# 16
# 25
# 36
# 49
# 64
# 81

👇ジェネレータ式(generator expression)についてのページです

ジェネレータの基本的な処理はまとめましたので
続いて、実際の業務を想定したサンプルコードをご紹介します

サンプルコード

仮に数万行に渡ってECの購入データのCSVファイルを扱うことにします。
各行には、ユーザIDと購入した商品の価格があります。
このCSVファイルからユーザID毎の合計購入金額を求めます。

性能的に悪いListを使った場合

use_list.py

user_orders = {}
# CSVファイルを読み込む
with open('order.csv') as f:
    for line in f:
        user_id, price = line.strip().split(',')
        # ユーザID毎に加算し、リストに詰める
        if user_id in user_orders:
            user_orders[user_id] += float(price)
        else:
            user_orders[user_id] = float(price)

ジェネレータを使う場合

use_generator.py

# ジェネレータ関数
def read_file(file_path):
    with open(file_path) as f:
        for line in f:
            yield line.strip() # stripで半角スペース、\nなどの空白文字を削除

user_orders = {}
# 1行ずつ読み込みながら処理する
for line in read_file('order.csv'):
    user_id, price = line.split(',')
    if user_id in user_orders:
        user_orders[user_id] += float(price)
    else:
        user_orders[user_id] = float(price)

おわりに

ジェネレータを使うと、1度にメモリ上にデータを保持せずに処理することができるためパフォーマンスの改善にオススメです。
大きなデータを扱っている方は参考にしてみたり、他の記事など調べてみてください！