python: 何これ？set operationを使ったらin operatorより爆速だった。

何をしていたか

あるデータベースから1,225,793個のウイルスゲノム配列をダウンロードして解析した。数か月たって、再度データをダウンロードしたところ、2,486,361個に増えていた。処理コードはすでに書いてあって、前回解析済みのものはスキップする仕組みになっている。コードのちょっとした手入れの最中に、思いがけない速度改善を得たのでそれについて書きます。

処理コードの概要

ウイルスゲノム配列はfastaファイルとして、また、前回までに解析済みのaccession idは、一行にひとつづつ書きこんだファイルに保存してあり、これらを入力ファイルとして用いる。

コードはこんな感じの作りです。

• 解析済みaccession idをファイルから読み込んでリスト(known_accessions)に格納する
• Bio.SeqIOを用いて、fasta fileから順次配列レコードを読み込んでいき、まずaccessionをとる
• そのaccessionが先のリストになければ次の処理に進む、あれば次のレコードへ

"accessionが先のリストになければ"、というところには単純にin operatorを使いました。

if acc not in known_accessions:
    # do sequence sanity check
    # re-construct id line
    # write seq to fasta file

ってな具合に。ちなみに、ファイルへの書き出し部分が速度のボトルネックだということに気が付いたので、ひとつづつ書き出すのではなく、ある程度ため込んでから（例えば5000配列）書き出すようにすることで、ずいぶん速度が上がりました。しかし、今日の要点は、そこではないのです。

あれ？

最初に気がついたのは、ファイルへの書き出しをさせない状態でのテストで、known_accessionsに短めの、あるいは空のリストを使うとやたら処理が速くなるということでした。え？もしかしてacc not in known_accessionsの評価で時間がかかっている？さすがにリストの長さが100万を超えると目に見えて影響が出る？ちなみにelementはすべて14文字のstrです。

こうしたら爆速になった

たまたま同じ日に、リスト間で共通するelementの抜き出しをしていて、set operationのスピードに感動していた。そこで、in operatorでやっている部分をset operationに変えたらどうなるか、を試してみた。if acc not in known_accessionsの部分に該当するsetを使ったバージョンは以下の通り。

if len(set([acc]) & set(known_accessions)) == 0:  # set([acc])は、{acc}でもいいですね
    # あれこれ

速い！しかも少し、じゃないですよ、めちゃくちゃ速い！実験用コードを書いて、10万のaccessionsをひとつづつ、2,486,361個のelementを含むリスト（setにしていますが)に対して試してみると、0:01:08.563127が0:00:00.445583、興奮してそのまま書いてしまいましたが、1分が0.45秒に！

ちなみに1万のaccessionsで試すと0.73秒が0.41秒で、それほどimpressiveには見えませんが、この数の繰り返しでは、コードの最初で行っている、ファイルからのaccessionの読み出しとリストへの格納が実行時間の大半を占めてしまうからだと思われます。

さて、もともとのコードでは7時間11分かかった本番の処理を、set operationでの評価に書き換えたコードで行ってみると、23分47秒で完了しました。たった一行の書き換えで、ここまでかわるなんて。あれこれ試して半日以上つぶれましたが、今後も繰り返す作業なので、もとはすぐに取れます。今後も精進を続けます。

まとめ

次のようなifブロックが繰り返し実行されるコードで、

if a in very_big_list:
    # do something

if a not in very_big_list:
    # do something

if文の評価に時間がかかっているのかも、と思ったら、

if len({a} & set(very_big_list)) > 0:
    # do something

if len({a} & set(very_big_list)) == 0:
    # do something

で改善されるか試してみましょう。