pandasの高速化はiterrows解消が9割

はじめに

こんにちは、株式会社オークンでデータサイエンティストをやっています、nobuです。

pandas、とにかく遅いですよね。
pandasの高速化については、いくつかのアプローチがあります

• 並行処理、並列処理を使う
• 他ライブラリを使う
  • polars
  • cudf

しかしそれらに手を出す前に、iterrows(ループ処理)を使っていたりしませんか？
知っている方には今更な話ですが、iterrowsは絶対に使ってはいけません！

iterrowsを解消し一括処理するだけで、かなりの速度改善ができます。
また、iterrowsを使ったままだと、上に挙げた他のアプローチをとってもあまり効果を得られなかったりします。

iterrowsは、ほとんどのケースについて、pandasのメソッドを駆使することで一括処理への変換が可能です。
今回は、iterrowsを使う代表的なケースについて、一括処理への変換例をいくつかご紹介できればと思います。
本記事を参照することで、手軽に高速化を実現しましょう！

事前準備

今回は以下のDataFrameを使っていきます。

import pandas as pd
import numpy as np

categories = ['category1', 'category2', 'category3', 'category4']
np.random.seed(1)
df = pd.DataFrame({
    'category': np.random.choice(categories, 30000),
    'value': np.random.randint(1, 100, 30000)
})

print(df.head())

    category  value
0  category1     21
1  category4     27
2  category2     92
3  category1     43
4  category3     13

適当に作った3万行のDataFrameです。
Category列（Category1~Category4の中からランダム）と、
Value列(1から99までのランダムな整数)を持っています。

では早速ループの解消例を見ていきます。

Case1. ループ中に条件一致した行に値を代入

以下のようなループはありがちです。

df['is_high'] = False
for i, row in df.iterrows():
    if row['value'] > 80:
        df.loc[i, 'is_high'] = True
    
print(df.head())

    category  value  is_high
0  Category1     21    False
1  Category4     27    False
2  Category2     92     True
3  Category1     43    False
4  Category3     13    False

実行時間は8.40sでした!

これを一括処理にする場合は以下のように書けます。

例1：numpyのwhereを使用 (実行時間: 1.73ms (ループより4800倍高速))

import numpy as np
df['is_high'] = np.where(df['value'] > 80, True, False)

例2：locを使用 (実行時間: 16.3ms (ループより500倍高速))

df['is_high'] = False
df.loc[df['value'] > 80, 'is_high'] = True

例3：applyを使用 (実行時間: 16.4ms (ループより500倍高速))

def judge_is_high(x):
    if x > 80:
        return True
    else:
        return False

df['is_high'] = df['value'].apply(judge_is_high)
 

Case2. ループしつつ、計算

df['calc'] = 0
for index, row in df.iterrows():
    # 50以下の場合は2乗し、それ以外の場合はルートをとる
    if row['value'] <= 50:
        calc = row['value'] ** 2
    else:
        calc = np.sqrt(row['value'])
    df.loc[index, 'calc'] = calc

23.68sかかりました

これを一括処理にするには、apply関数を使うのが簡単です。

def calc(x):
    # 50以下の場合は2乗し、それ以外の場合はルートをとる
    if x <= 50:
        calc = x ** 2
    else:
        calc = np.sqrt(x)

df['calc'] = df['value'].apply(calc)

実行時間は、0.12sと、190倍、高速になりました。

Chat GPTに聞く

本記事の存在意義がなくなってしまいますが(爆)、
他のケースについてはChatGPTにどう解消できるかと聞くと適切なpandasのメソッドを教えてくれます。
ループを一括処理に変換する上での一番のハードルは、pandasのメソッドを思いつくことですが、ChatGPTを活用することで、乗り越えられます。

まとめ

ループ処理を解消すると、桁違いに早くなることを体感頂けたかと思います。
他の高速化の手法も色々ありますが、はじめはとにかくiterrowsの解消だけ取り組むのがよいかと思っています！

さいごに

最後までご覧いただき、ありがとうございました！
株式会社オークンでは、ヘッジファンド向けのトレードシステムを先日リリースいたしました。
現在、さらなる事業拡大に向け、エンジニアを積極的に採用しています。
本プロジェクトにご興味をお持ちいただけましたら、こちらの記事もぜひご覧ください！