歩いたら休め

If the implementation is easy to explain, it may be a good idea.

【Python】ジェネレーターをn個ずつに分割する実装

「巨大なテキストファイルをジェネレーターとして読み込み、100万行ごとに分割し、別々のファイルに保存する」という処理を書いてました。

数百行ごとに分けるのならリストにして分割するのですが、今回は分ける単位が1000万行ごとなので一度にメモリに載せたくありません。遅延評価を保ったまま、以下のようにファイルを分けるコードを書こうと考えていました。

要するにnumpy.array_splitの入力値も出力もgenerator版です。

numpy.array_split — NumPy v1.14 Manual

from pathlib import Path

path = Path("/tmp/example/")

# tableは巨大なテキストファイルを一行ずつ読み込んだgenerator

for i, rows in enumerate(split_generator(table, 1000000)):
    filepath = path / f"output{i}.txt"
    with filepath.open("w") as f:
        for row in rows:
            f.write(row)

しばらく悩んでいたのですが、一応、以下のようにitertoolsを駆使して実装することができました。

from itertools import islice, chain

def split_generator(iterator, n):
    while True:
        each = islice(iterator, n)
        first = next(each, None)
        if first is None:
            break
        yield chain([first], each)

ただ、ジェネレーターが先頭から評価されなかった場合にどうなるのかとか、そもそも本当に1行ずつメモリに載ってんのとか、ちゃんと見ていかないと安心して使えなさそうなコードになってしまいました。もっと安心して使えそうな実装を思いついた(知ってる)方はコメントください。

「この問題についての正しい回答はおそらくPythonではなく適切なデータ処理基盤を使うことだ」と思っているのですが、今回はデータ元が特殊(記事ではテキストファイルとしてしまいましたが)だったのでそれも難しいです。

エンジニアリングって難しいですね。