gzip, bz2, lzma の binary mode でのイテレーションを高速化しました。

次の2つのプログラムは、どちらも test.gz というファイルを開いて、片方はバイナリモードで、もう片方はテキストモードでイテレートします。Pythonの使用上ファイルをイテレートしたら行単位で内容を読み出せます。

$ cat dec_gzip.py
import gzip
with gzip.open("test.gz", 'rb') as f:  # バイナリモード
    for l in f:
        pass

$ cat dec_gzip_t.py
import gzip
with gzip.open("test.gz", 'rt', encoding="utf-8") as f:  # テキストモード
    for l in f:
        pass

この2つのプログラム、どちらが速いと思いますか？僕はバイナリモードの方だと思っていました。 TextIOWrapper() の処理がまるまるオーバーヘッドになるはずだからです。

$ python3 -V
Python 3.9.1

$ time python3 dec_gzip.py

real    0m0.528s
user    0m0.508s
sys     0m0.021s

$ time python3 dec_gzip_t.py

real    0m0.285s
user    0m0.280s
sys     0m0.007s

なんと、バイナリモードの方が2倍くらい遅いのです。

バイナリモードで gzip.open() すると GzipFile というオブジェクトが返されます。この __iter__ メソッドは親クラスの IOBase.__iter__ で、Cで実装されているのですが、振る舞いとしては毎回 self.readline() を呼び出してそれを返します。そして GzipFile.readline() は Python で実装されているのです。

    def readline(self, size=-1):
        self._check_not_closed()
        return self._buffer.readline(size)

一方テキストモードで開いた場合は、 GzipFile を TextIOWrapper でラップしたオブジェクトが返されます。 TextIOWrapper.__iter__ も IOBase.__iter__ なのですが、 TextIOWrapper.readline() が C で実装されているのと、TextIOWrapperがある程度バッファリングしていて毎回 GzipFile.read() を呼び出さなくて良いので速いのです。

bz2 モジュールの場合はさらにこの差が広がります。 gzip と lzma モジュールはスレッドセーフではないのですが、 bz2 モジュールの BZ2File だけは複数スレッドからの並行な read/write が許可されていたので、行ごとに毎回呼び出されるPythonで実装された readline メソッドの中で with self._lock しているのです。

$ time python3 dec_bz2.py

real    0m1.657s
user    0m1.655s
sys     0m0.004s

$ time python3 dec_bz2_t.py

real    0m0.868s
user    0m0.851s
sys     0m0.016s

bz2の方がgzipよりも圧倒的にデコード負荷が高いために比率で見たら同じ程度の速度低下で済んでいますが、差（オーバーヘッドの大きさ）を考えると、同じ100万行を処理するのにかかるオーバーヘッドが0.22sec程度から0.8sec程度に増えています。

これを Python 3.10 で次のように修正しました。

bz2 はスレッドセーフをやめる。 (issue43785)
GzipFile, BZ2File, LZMAFile で __iter__ をオーバーライドし、内部で使っているバッファ (io.BufferedReader) の __iter__ に移譲することで、毎回Python実装された readline() が呼ばれないようにする。 (Add iter to GzipFile, BZ2File, and LZMAFile #25353)

結果、バイナリモードでのイテレートがPython 3.9より倍以上速くなり、テキストモードよりも少し速くなりました。

$ time ./python dec_gzip.py

real    0m0.198s   # old: 0.528s, text: 0.285s
user    0m0.195s
sys     0m0.005s

$ time ./python dec_bz2.py

real    0m0.771s  # old: 1.657s, text: 0.868s
user    0m0.758s
sys     0m0.015s

ちなみに、Python 3.9 でバイナリモードで行単位の処理をしたい場合は、プライベートメンバーに触ってしまいますが次のようにすれば同じ高速化ができます。

with gzip.open(filename) as f:
    for l in f._buffer:
        pass