Pandasの繰り返し処理を最適化してiterrowsの約1000倍速くする話

TL;DR

• iterrows()は型チェックが起こるためfor文で回すには遅い
• for文をどうしても使いたいのであれば、to_numpy()してから回す（iterrowsの約40倍高速）
• おすすめは関数化してnp.vectorizeした後に適用する（iterrowsの約50倍高速）
• 操作が単純で最速を目指すならnp.where等numpyだけで処理する（iterrowsの約1000倍高速）

手法	時間（ms）	iterrows比の高速化割合
numpyのみで処理	1.51	947
pandasのみで処理	7.72	185
関数をvectorizeして適用	30.7	47
to_numpy()でfor文を回す	35.2	41
df.valuesでfor文を回す	57.9	25
itertuples()でfor文を回す	92.7	15
関数をapplyで適用	315	5
While文	383	4
iterrows()でfor文を回す	1430	1

背景

pythonやpandasはfor文が遅いという認識が一般的ですが、遅い理由の多くは型チェックであることが多いので、特定の計算をする際にどの方法が速く、書きやすいのか複数手法で比べてみました。

実験概要

women's tennis data (n = 85288)に対して、rating_1の大きさとrating_2の大きさによって手法が変わる計算式を適用しています。なお計算された値にはなんの意味もありません。

df = pd.read_csv('2016-2017_orig.csv', sep=";") # women's tennis data 
r = []
for _, row in df.iterrows():
    if row["rating_1"] > row["rating_2"]:
        r.append(row["rating_1"]*2 + row["rating_2"])
    else:
        r.append(row["rating_1"] + row["rating_2"] * 2)

比較検討する手法は下記の9つです。なお今回は複数スレッド(pandarallel)やcython等は考慮していません。

1. iterrows()でfor文を回す
2. While文
3. itertuples()でfor文を回す
4. to_numpy()でfor文を回す
5. df.valuesでfor文を回す
6. 関数をapplyで適用
7. 関数をvectorizeして適用
8. pandasのみで処理
9. numpyのみで処理

実験結果

iterrows()でfor文を回す

r = []
for _, row in df.iterrows():
    if row["rating_1"] > row["rating_2"]:
        r.append(row["rating_1"]*2 + row["rating_2"])
    else:
        r.append(row["rating_1"] + row["rating_2"] * 2)

記録:1.43s

遅いがいまだによく使われている

While文

i = 1
l = len(df)
r= []
while i < l:
    r1=df["rating_1"][i]
    r2=df["rating_2"][i]
    if r1 > r2:
        r.append(r1*2 + r2)
    else:
        r.append(r1 + r2 * 2)
    i += 1

記録:383ms
型チェックの回数を減らすためにintでループしたが期待したほど速くならなかった

itertuples()でfor文を回す

r = []
for row in df.itertuples():
    if row[5] > row[6]:
        r.append(row[5]*2 + row[6])
    else:
        r.append(row[5] + row[6] * 2)

記録:92.7ms
itertuplesはnamedtupleが返却され、型チェックの回数が大幅に減るためそこそこ早い

to_numpy()でfor文を回す

r = []
for r1,r2 in zip(df["rating_1"].to_numpy(), df["rating_2"].to_numpy()):
    if r1 > r2:
        r.append(r1*2 + r2)
    else:
        r.append(r1 + r2 * 2)

記録:35.2ms
for文の最速。for文を書きたいならこれが一番いい。
事前に必要な列をto_numpy()でnumpyにしておきその後でzip->iterateすることにより型推論が行われないことや、値がすでに確定しているため、配列アクセスがfor中に不要などメリットがありそう。マジックコマンド%%prunで確認すると、型チェックも少なく効率的に処理されている。

df.valuesでfor文を回す

r = []
for row in df.values:
    if row[5] > row[6]:
        r.append(row[5]*2 + row[6])
    else:
        r.append(row[5] + row[6] * 2)

記録:57.9ms
そこそこ早いが、itertuples同様インデックスで値にアクセスしなければならないのでやや書きにくさが残る。

関数をapplyで適用

def add(val1, val2):
    if val1 > val2:
        return val1*2 + val2
    else:
        return val1 + val2 * 2
df.apply(lambda _df: add(_df["rating_1"], _df["rating_2"]), axis=1)

記録:315ms
これもよくある書き方だが、内側で型チェックが頻繁に起こるためこれも遅い。

関数をvectorizeして適用

def add(val1, val2):
    if val1 > val2:
        return val1*2 + val2
    else:
        return val1 + val2 * 2
v = np.vectorize(add)
v(df["rating_1"], df["rating_2"])

記録:30.7ms
これが個人的に最もおすすめ。関数化できるため保守性も高い
np.vectorizeは関数をnumpy配列でも適応できるようにするという利便性のために用いられているもので速さのためではないと公式ドキュメントに明記されているが、numpy配列を正しい型で適用すると、型チェックが必要なくなり高速になる。

pandasのみで処理

vals1 = df["rating_1"]*2 + df["rating_2"]
vals2 = df["rating_1"] + df["rating_2"]*2
df["r"] = vals1
df.loc[df.rating_1 > df.rating_2, "r"] = vals2

記録:7.72ms
pandas単体には条件分岐をする仕組みを見つけられなかったので両方の可能性を事前に計算しておき.locで一部値を更新している。
そこそこ速いが拡張性はない

numpyのみで処理

r1 = df["rating_1"].to_numpy() 
r2 = df["rating_2"].to_numpy()
r = np.where(r1 > r2,  r1 *2 + r2, r1 + r2*2)

記録:1.51ms
極力numpyのみで処理をすると大変高速になる。
簡単な条件分岐であればこれで十分だが、複雑な処理になると可読性が下がるのでどうしても速くしたい場合以外はあまりお勧めしない。

考察

• numpyやpandasの配列のまま処理すると非常に高速
• for文等ループを書くときはできるだけ、型チェックが起こらないように型情報を明示するような書き方をすると速くなる
• 複雑な処理はnumpyやpandasの演算だけで処理するのは難しいので、関数化して適用した方が良い。
• とりあえずvectorizedな関数で実装し、必要ならさらに高速化していく手法が個人的にお勧め
• pandarallelはオーバヘッドが大きな印象なので今回は使ってない
• cythonはおそらくnumpyのみに近い値が出せるはずだが、この程度の件数であればnumpyのみより速くなる印象がなかったため用いてない