pandasのqueryのススメ
こんにちは。株式会社Rosso、AI部です。
前回は個人的な趣味全開のnote(下記UFO調査)を書かせてもらったので、今回はより実務で使えそうな記事を書きました!
どなたかのお役に立てれば嬉しいです。
はじめに
pandasのデータフレームを扱う際、行の絞り込み実行時に、何度もデータフレームの変数名を書くのが嫌という理由で、私はqueryメソッド派です。
df.query("C=='b'")
この記事ではqueryメソッドの魅力をお伝えします。
なお、弊社team_aiデータ分析技術チームでアンケートを取った結果、
圧倒的にdf[df["C"]=="b"]派が優勢でした。
queryの書き方
そもそも、queryの使い方どころか存在すら知られていないかもしれないので、説明したいところですが、すでに分かりやすくまとめられているページがいくつかあったのでそちらを参照してもらえばほとんどOK。私は軽く紹介する程度にしておきます。
pandasのquery書き方オススメページ
DataFrameレシピ: データ抽出条件 FukuharaYohei様
pandas.DataFrame.queryによるデータ抽出10選 dox様
数値絞り込み
日本語の列名でも問題なし
df.query("年齢 ==50")
文字列絞り込み
対象の文字列を、queryと異なるクォーテーションマークで囲めばOK。
内側と外側が異なれば("都道府県 == '東京都'")でも('都道府県 == "東京都"')でもOK。
df.query("都道府県 == '東京都'")
変数の使用
query内のテキストで@を使えば変数を参照させられる。
age = 70
df.query("年齢==@age")
複数条件
queryが本領を発揮するのは複数を組み合わせたとき。
age = [50,70]
area = "北海道"
gender = "男"
df.query("都道府県==@area & 年齢 in @age & 性別==@gender")
queryを使う方が12文字も短くなる。
今回はデータフレームの変数名がシンプルな2文字のdfにもかかわらず。
queryを使う場合、50文字
df.query("都道府県==@area & 年齢 in @age & 性別==@gender")
使わない場合、62文字
df[(df["都道府県"]==area)&(df["年齢"].isin(age))&(df["性別"]==gender)]
と差が出た。
もしデータフレームの変数名がjapanese_worker_dfとかだったらこの差はもっと広がるはず。
よくある誤解
様々なチームでpandasの条件絞り込みの話題になった時に、上がった誤解について触れていきます。
誤解1:変数名に変なのがついてたらqueryは使えない?→使えます!
df.query("`年収 世帯(万円)`<105")
誤解2:変数で値は指定できるけど列名は指定できない?→できます!
先程までのテーブルに月度ごとの労働時間のカラムを追加。
python3.6から導入されたf文字列を使用すれば楽。formatメソッドでもOK。
month = 4
hour = 295
df.query(f"`{month:0=2}月労働時間`==@hour")
誤解3:遅いんじゃない?→自分で検証してみました
以下のように100万行のデータを作成し、「queryあり」と「なし」を順番ランダムで100回ほど実行させて時間を計測し、平均を取ったところ、
queryあり:0.12013(s)
〃なし:0.12893(s)
と、queryありのほうが早かった!!!(よ・・・よかった)
import pandas as pd
import time
def restart_df(row_num = 1000000):
columns = ["都道府県","性別","血液型","年齢"]
columns += [f"{i:0=2}月労働時間" for i in range(1,13)]
columns += ["年収 世帯(万円)"]
values = []
values.append(np.array(["北海道"]*(row_num//2) + ["東京都"]*(row_num-row_num//2)))
values.append(np.array(["男" if i == 0 else "女" for i in np.random.randint(0,2,row_num)]))
values.append(np.array(random.choices(["A","B","O","AB",None,None],k=row_num)))
values.append(np.random.randint(20,10**(2),row_num))
for i in range(len(columns)-5):
values.append(np.random.randint(100,300,row_num))
values.append(np.random.randint(100,1000,row_num))
df = pd.DataFrame(np.array(values).T,columns = columns)
for col in columns[3:]:
df[col] = df[col].astype(int)
return df
age = [50,70]
e = 2
area = "北海道"
gender = "男"
A_time_list = []
B_time_list = []
for i in range(100):
df = restart_df(1000000)
#先に実行した方が遅いらしいので、交互に実施
if i%2==0:
start = time.time()
_ = df[(df["都道府県"]==area)&(df["年齢"].isin(age))&(df["性別"]==gender)]
end = time.time()
B_time_list.append(end-start)
start = time.time()
_ = df.query("都道府県==@area & 年齢 in @age & 性別==@gender")
end = time.time()
A_time_list.append(end-start)
else:
start = time.time()
_ = df.query("都道府県==@area & 年齢 in @age & 性別==@gender")
end = time.time()
A_time_list.append(end-start)
start = time.time()
_ = df[(df["都道府県"]==area)&(df["年齢"].isin(age))&(df["性別"]==gender)]
end = time.time()
B_time_list.append(end-start)
numexprというライブラリがインストールされていれば、pandasのqueryでの絞り込みが早くなるらしい。ということで、numexprの有無でも試してみたが、query有りの方が早いという結果は同じでした。
誤解4(番外編):queryってSQLみたいなの書くんでしょ?→違います
queryの欠点
欠点も挙げないとフェアじゃないと思ったので調べてみました。
欠点:文字列の処理がめんどくさい
XXXX.str.contains("aaa")系の処理ではengine="python"を指定する必要がある。
df.query("都道府県.str.contains('都')",engine="python")
しかし、カラムに欠損が含まれる場合はqueryだとどうしてもエラーとなってしまう
df.query("血液型.str.contains('A')",engine="python")
"""
ValueError: Cannot mask with non-boolean array containing NA / NaN values
"""
これはqueryを使わない方法だとna=False で回避できるので、文字列の含む・含まない系の処理だけはqueryを使うのを我慢しよう。
df[df["血液型"].str.contains('A',na=False)]
まとめ
pandasでのdataframeデータ絞り込みにqueryメソッドはオススメ!!!
文字数少なくて済む!!!
処理速度も遅くはない!!!
もし、queryメソッドを知らなかっただけなら、ぜひ今日から試してみてください。以上!!