Pythonを使って2MBのメモリで100万の数値をソートする¶

誰かからジョーク交じりに、100万個の32ビットの数値を2メガバイトのメモリでソートできるか？と聞かれたことがある。私はこれに挑戦してみたが、この中でI/Oのバッファリングについていくつか学ぶことができた。

明らかにこの問題はジョークである。バイナリエンコーディングだと仮定してもデータだけで4メガバイトになってしまうのである！しかし、これを可能にする実装方法がある。32ビットの数値が100万個納められているファイルがあるとして、どのようにすればメモリの使用量を最小にしてソートをすることができるだろうか？これは一種のマージソートになるはずである。メモリ内でソートされた小さいチャンク(かたまり)を一時ファイルに保存し、その後この一時ファイルをマージして最終的な出力エリアに出力する。

これは私の考えた解決策である。手短に(1分で)説明しよう。

 #!/usr/bin/env python3.0
 import sys, array, tempfile, heapq
 assert array.array('i').itemsize == 4

 def intsfromfile(f):
     while True:
         a = array.array('i')
         a.fromstring(f.read(4000))
         if not a:
             break
         for x in a:
             yield x

iters = []
while True:
    a = array.array('i')
    a.fromstring(sys.stdin.buffer.read(40000))
    if not a:
        break
    f = tempfile.TemporaryFile()
    array.array('i', sorted(a)).tofile(f)
    f.seek(0)
    iters.append(intsfromfile(f))

a = array.array('i')
for x in heapq.merge(*iters):
    a.append(x)
    if len(a) >= 1000:
        a.tofile(sys.stdout.buffer)
        del a[:]
if a:
    a.tofile(sys.stdout.buffer)

私のGoogleのデスクトップ(Guidoは現在Googleで働いている)では、実行に5.4秒かかった。3年前の古いPCでGoogle謹製のLinuxが稼動しており、Python3.0のpystoneのベンチマーク結果は34000である。また、入力には正確に100万個のランダムな32ビットの数値の入ったファイルを使用した。この結果はそんなに悪くない。同じ入力を使用して、すべてメモリ上でソートするという、直球な実装では2秒かかった。

#!/usr/bin/env python3.0
import sys, array
a = array.array('i', sys.stdin.buffer.read())
a = list(a)
a.sort()
a = array.array('i', a)
a.tofile(sys.stdout.buffer)

それではマージソートの方(ファイルを使用する方)に説明を戻そう。最初の３行は以下のようになっている。

import sys, array, tempfile, heapq
assert array.array('i').itemsize == 4

最初の行はPython 3.0を使用するということを言っている。二行目では必要なモジュールをインポートしている。三行目はタイプコードの’i’が32ビットの数値をあらわさない、64ビットのシステムを除外するためのものである。私は特にこのコードをいろいろな環境で動く汎用的なものにしようとは思っていない。

次に、小さなヘルパーを作る。これはジェネレータであり、ファイルから数値を読み込み、一度に1つずつ返すものである。

def intsfromfile(f):
    while True:
        a = array.array('i')
        a.fromstring(f.read(4000))
        if not a:
            break
        for x in a:
            yield x

ここの部分のパフォーマンスをチューニングする。このコードは一度に1000個の整数を読み込む。そして、その結果を分割して1つずつyieldで返す。最初に実装したコードはこのバッファリングを使用していないため、一度に4バイトずつファイルから読み込んで、整数に変換して、結果を返していた。しかし、これでは4倍も遅くなってしまうのである！(通常、細かいファイルアクセスは、一度にまとめて読むのに比べて時間がかかる) a.fromfile(f, 1000)というメソッドは使えないということは覚えておいて欲しい。このfromfile()メソッドは、ファイル上に十分な数のデータがそろっていないときには文句を言ってくるからである。私は、ファイル上にある整数の個数が何個であれ、自動で動作してくれるようにしたかったので使用しなかったのである。

次に来るのが入力のループである。ここでは繰り返し1万個の整数値が含まれるチャンクを入力ファイルから読み込み、メモリ上でソートし、一時ファイルに書き込んでいる。その後、上記のinsfromfile()関数を使用して作成した、一時ファイルを参照するイテレータを作成する。そして、この次のマージを行うフェーズで使用する、イテレータのリストに追加する。

iters = []
while True:
    a = array.array('i')
    a.fromstring(sys.stdin.buffer.read(40000))
    if not a:
        break
    f = tempfile.TemporaryFile()
    array.array('i', sorted(a)).tofile(f)
    f.seek(0)
    iters.append(intsfromfile(f))

100万個のデータが含まれるデータを使用すると、それぞれ1万個のデータを含む、100個の一時ファイルが作成される。

最後に、これらのソート済みの一時ファイルを一緒にマージする。heapqモジュールを使用すると、簡単にこの目的を達成できる。heapq.merge(iter1, iter2, ...)は、順番どおりの入力値をyieldで返すイテレータを作成する。この関数は、今回の場合のように、入力されたそれぞれイテレータは、ソート済みであるということが前提となる。

a = array.array('i')
for x in heapq.merge(*iters):
    a.append(x)
    if len(a) >= 1000:
        a.tofile(sys.stdout.buffer)
        del a[:]
if a:
    a.tofile(sys.stdout.buffer)

ここも、I/Oのバッファリングを利用して、効率的にチューニングできるもう一つのポイントである。ソートされた値を得られる端から一つずつファイルに追加していくと、だいたい２倍遅くなる。これらを総合すると、ファイルの読み込みと書き出しの両方にバッファを設けただけで、10倍のスピードアップが得られたことになります！

以上が今回のストーリーの金言である。

もう一つのheapqモジュールがすばらしいということも、今回の学びである。このモジュールは出力のフェーズで必要となった、イテレータをマージする機能を持っていた。そして、バイナリデータを管理するのにarrayモジュールを使用したことも忘れないで欲しい。

最後に、このサンプルを見て、Python3.0とPython2.5がそんなに大きな違いがないということがご理解いただけたと思う！