pndas.DataFrame.queryを高速化したい

結論

関数を使用せず直接条件を記述したほうが、queryを使用するよりも処理速度が速い

for i in range(0, 100):
    df_iris[df_iris["sepal length (cm)"] >= 5]

背景

queryを使用した処理を作っていたが、処理時間が長かったので短縮できないか検討してみた。

環境

項目	情報
OS	Windows 10
Python	3.9.7
pandas	1.3.4
sklearn	1.1.3

速度比較

事前準備

import time
import pandas as pd
from sklearn.datasets import load_iris

# irisのデータセットを準備
data_iris = load_iris()
df_iris = pd.DataFrame(
    data_iris.data, 
    columns=data_iris.feature_names
)

# データを確認
display(df_iris.shape)
display(df_iris.head())

出力

単一条件で比較

queryを使用する(engine: python)

%%time
# 【単一条件】queryを使用する(engine: python)
for i in range(0, 10000):
    df_iris.query("`sepal length (cm)` >= 5", engine='python')

Wall time: 17.7 s

queryを使用する(engine: numexpr)

%%time
# 【単一条件】queryを使用する(engine: numexpr)
for i in range(0, 10000):
    df_iris.query("`sepal length (cm)` >= 5", engine='numexpr')

Wall time: 23.8 s

関数を使わない

%%time
# 【単一条件】関数を使わない
for i in range(0, 10000):
    df_iris[df_iris["sepal length (cm)"] >= 5]

Wall time: 3.22 s

複数条件で比較

queryを使用する(engine: python)

%%time
# 【複数条件】queryを使用する(engine: python)
for i in range(0, 10000):
    df_iris.query("`sepal length (cm)` >= 5 & `sepal_width (cm)` <= 3", engine='python')

Wall time: 322 ms

queryを使用する(engine: numexpr)

%%time
# 【複数条件】queryを使用する(engine: numexpr)
for i in range(0, 10000):
    df_iris.query("`sepal length (cm)` >= 5 & `sepal_width (cm)` <= 3", engine='numexpr')

Wall time: 32.8 s

関数を使わない

%%time
# 【複数条件】関数を使わない
for i in range(0, 10000):
    df_iris[(df_iris["sepal length (cm)"] >= 5) & (df_iris["sepal length (cm)"] <= 3)]

Wall time: 6.28 s

まとめ

条件	query(engine: python)	query(engine: numexpr)	関数を使わない
単一	17.7 s	23.8 s	3.22 s
複数	32.8 s	39.5 s	6.28 s

queryを使用せず、直接条件を書いて抽出したほうが圧倒的に早い

追加検証候補

• ループ回数を変更した場合の比較
• 条件を複雑にした場合の比較
• DataFrameの行数を増やした場合の比較

参考サイト