pandas の query メソッドで条件抽出を簡潔に書く

pandas で条件抽出するとき、通常は df[df['col'] > 0] のようにブールインデックスを使います。query メソッドを使うと、より読みやすい文字列で同じ処理を書けます。

query の基本

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'name': ['田中', '佐藤', '鈴木', '高橋'],
    'age': [25, 30, 35, 28],
    'score': [85, 72, 90, 68]
})

# 通常の方法
result = df[df['age'] >= 30]

# query を使う
result = df.query('age >= 30')

print(result)
#   name  age  score
# 1   佐藤   30     72
# 2   鈴木   35     90

複雑な条件になるほど、query の可読性が際立ちます。

# 通常の方法(括弧が多くて読みにくい)
result = df[(df['age'] >= 25) & (df['age'] <= 35) & (df['score'] >= 70)]

# query を使う
result = df.query('25 <= age <= 35 and score >= 70')

複数条件の指定

and / or を使って複数条件を組み合わせられます。

# AND 条件
df.query('age >= 30 and score >= 80')

# OR 条件
df.query('age < 25 or score >= 90')

# NOT 条件
df.query('not (age >= 30)')

Python の演算子 & | ~ の代わりに、英単語 and or not が使えるのも query の利点です。

ブールインデックス

& | ~ を使う。各条件を () で囲む必要がある。

query

and or not を使う。数学的な比較演算も直感的に書ける。

変数の参照

query 内で Python の変数を参照するには @ を付けます。

min_age = 30
min_score = 80

# 変数を参照
result = df.query('age >= @min_age and score >= @min_score')

print(result)
#   name  age  score
# 2   鈴木   35     90

リストを使った in 演算も @ で参照できます。

target_names = ['田中', '鈴木']

# リストに含まれるか
result = df.query('name in @target_names')

print(result)
#   name  age  score
# 0   田中   25     85
# 2   鈴木   35     90

文字列の条件

文字列の比較には引用符を使います。

# 文字列の完全一致
df.query('name == "田中"')

# シングルクォートでも可
df.query("name == '田中'")

ただし、文字列メソッド(startswith, contains など)は query 内では使えません。その場合は通常のブールインデックスを使います。

# これはエラーになる
# df.query('name.str.startswith("田")')

# 通常の方法を使う
df[df['name'].str.startswith('')]

スペースを含む列名

列名にスペースや特殊文字が含まれる場合、バッククォートで囲みます。

df = pd.DataFrame({
    'first name': ['田中', '佐藤'],
    'total score': [85, 72]
})

# バッククォートで囲む
result = df.query('`total score` >= 80')

print(result)
#   first name  total score
# 0         田中           85

index の参照

index を条件に使うこともできます。

df = pd.DataFrame({
    'value': [10, 20, 30, 40]
}, index=['a', 'b', 'c', 'd'])

# index を条件に使う
result = df.query('index in ["a", "c"]')

print(result)
#    value
# a     10
# c     30

数値インデックスの場合も同様です。

df = pd.DataFrame({
    'value': [10, 20, 30, 40]
})

# 数値インデックス
result = df.query('index >= 2')

print(result)
#    value
# 2     30
# 3     40

パフォーマンス

query は内部で numexpr ライブラリを使っており、大きなデータセットでは通常のブールインデックスより高速になることがあります。

import pandas as pd
import numpy as np
import time

# 大きなデータ
df = pd.DataFrame({
    'a': np.random.randn(1_000_000),
    'b': np.random.randn(1_000_000)
})

# ブールインデックス
start = time.time()
for _ in range(100):
    result = df[(df['a'] > 0) & (df['b'] > 0)]
print(f"ブールインデックス: {time.time() - start:.2f}")

# query
start = time.time()
for _ in range(100):
    result = df.query('a > 0 and b > 0')
print(f"query: {time.time() - start:.2f}")

ただし、小さなデータセットでは query の方がオーバーヘッドで遅くなることもあります。

データサイズ推奨
小(〜1万行)どちらでも可。可読性で選ぶ
中(1万〜100万行)query が有利なことが多い
大(100万行〜)query + numexpr で高速化

inplace オプション

query は新しい DataFrame を返しますが、inplace=True で元の DataFrame を直接変更できます。

# 通常は新しい DataFrame を返す
result = df.query('a > 0')

# inplace で元を変更(非推奨)
# df.query('a > 0', inplace=True)

ただし、inplace は将来廃止される可能性があり、使用は推奨されません。

query が使えないケース

いくつかの操作は query ではできません。

文字列メソッド

str.contains(), str.startswith() などは使えない

NaN のチェック

isna(), notna() は使えない。代わりに col == col(NaN でない行)を使う

# NaN を除外するトリック
df = pd.DataFrame({'a': [1, np.nan, 3]})

# a == a は NaN のとき False になる
df.query('a == a')  # NaN を除外

query は可読性とパフォーマンスの両面でメリットがあります。複雑な条件抽出を書くときは、まず query で書けないか検討してみてください。