Introduction
ここに以下のような先人の方が作成したPython関数があります。あるとき、これを高速化してほしいとの依頼がありました。さてあなたならどうするでしょうか。ちなみにこの関数はforループによって複数回呼び出されており、その実行時間は135s程度でした。この待ち時間を短縮せよ、というのが今回のミッションです。
def func(hogelist, dblines):
annotation = ''
flag = 0
for name in hogelist:
for line in dblines:
dbname = line.split(sep='\t')[0]
dbvalue = line.split(sep='\t')[7].replace(' ', '_')
if name == dbname:
if dbvalue != '':
if flag > 0:
annotation = annotation + ','
annotation = annotation + dbvalue
flag = 1
result = ','.join(list(set(annotation.split(sep=','))))
if result == '':
result = '-'
return(result)
この関数を呼び出しているmainはこんな感じです。
'''
色々な処理
'''
with open('db.tsv', 'r') as f:
dblines = f.readlines()
for hogelist in hogelists:
result = func(hogelist, dblines)
'''
色々な処理
'''
処理概要
dblinesはとあるデータベースから読み込んだものです。
この関数はhogelistのある要素をnameと呼称し、それがdblinesのある一行lineをsplitしたもののうちの0番目の要素が一致するかを判定します。
もし一致すれば、splitしたうちの7番目の要素を、そのnameに付随する情報として、annotationに付与します。
最後に、annotationのうち、重複するものがあればそれをsetを用いて除去したものをresultとして返します。もし、元からresultが空文字であれば、データベースに該当なしという意味を込めてハイフンを代入しています。
高速化
データベースをリスト型から辞書型へ変更する
以下のように関数を書き換えてみます。
def func(hogelist, dblines):
annotation = ''
# set dbdict{dbname:dbvalue}
dbdict = {}
for line in dblines:
dbname = line.split(sep='\t')[0]
dbvalue = line.split(sep='\t')[7].replace(' ', '_')
dbdict.setdefault(dbname, []).append(dbvalue)
# set annotation
for name in hogelist:
for key in dbdict:
if name == key:
annotation = dbdict[key]
result = ','.join(list(set(annotation)))
if result == '':
result = '-'
return(result)
lineをsplitしたときに、0番目の要素は同じですが、7番目の要素が違うものがある場合があります。それを辞書型のsetdefaultメソッドとappendを用いて、登録してあげています。この方法はこちらのページから勉強させていただきました。このアルゴリズム改善のおかげで、実行時間は119sまで短縮されました。
データベースのやりくりをforループの外側に持っていく
すでにお気付きかと思いますが、先の辞書dbdictは作成された後、この関数内では値の更新などは一切されていません。なので、この関数が複数回呼び出されるごとに辞書を一から作り直すのではなく、辞書を作成する外部関数を新たに作成し、その結果を参照する形にしましょう。
以下のように新しい関数を作り、既存の関数も書き換えます。
# make dictionary
def make_dbdict(dblines):
dbdict = {}
for line in dblines:
dbname = line.split(sep='\t')[0]
dbvalue = line.split(sep='\t')[7].replace(' ', '_')
dbdict.setdefault(dbname, []).append(dbvalue)
return(dbdict)
def func(hogelist, dbdict):
annotation = ''
# set annotation
for name in hogelist:
for key in dbdict:
if name == key:
dbdict[key] = set(dbdict[key])
dbdict[key].discard('')
annotation = dbdict[key]
result = ','.join(list(annotation))
if result == '':
result = '-'
return(result)
そしてmainも以下のように書き変わります。
'''
色々な処理
'''
with open('db.tsv', 'r') as f:
dblines = f.readlines()
dbdict = make_dbdict(dblines)
for hogelist in hogelists:
result = func(hogelist, dbdict)
'''
色々な処理
'''
これでhogelistsの要素数の数だけ繰り返されていた辞書作成処理が1回になりました。このアルゴリズムの見直しで、この実行時間が0.54sまで縮みました。
@ lru_cacheデコレーターで処理をさらに高速化
アルゴリズムの見直しはここまで。ここから先は関数の戻り値をキャッシュとして記憶させることで実行時間を短縮させるという、少しディープなやり方をします。そのためには以下のモジュールをインポートします。
from functools import lru_cache
このモジュールの機能の詳細はこちらの記事やこちらのブログを参照してください。
時間を少しでも短縮したい関数の前に
@lru_cache(maxsize=1024)
と、main関数の最後にメモリ解放のおまじない
func.cache_clear()
をつけます。するとこれだけで実行時間が0.37sまで短縮されました。
結論
今回は、すでに他の方が作成したPythonスクリプトを高速化するというミッションについてご説明しました。アルゴリズムの見直しから、Pythonの便利モジュールまで、あらゆる手を尽くしての高速化するという今回のミッション。結果として300-400倍高速化することに成功しました。
この記事が皆様の参考になれば幸いです。