pandas.DataFrameの行を条件で抽出するquery

Modified: | Tags: Python, pandas

pandas.DataFrameから任意の条件を満たす行を抽出するにはquery()メソッドを使う。比較演算子や文字列メソッドによる条件指定、複数条件の組み合わせなどを簡潔に記述できる。

ブーリアンインデックス(Boolean indexing)による条件指定については以下の記事を参照。

特定の型の列を抽出したり、行名・列名で行・列を抽出したりすることも可能。

本記事のサンプルコードのpandasのバージョンは以下の通り。バージョンによって仕様が異なる可能性があるので注意。以下のDataFrameを例として使う。

import pandas as pd

print(pd.__version__)
# 2.1.4

df = pd.read_csv('data/src/sample_pandas_normal.csv')
print(df)
#       name  age state  point
# 0    Alice   24    NY     64
# 1      Bob   42    CA     92
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70
# 4    Ellen   24    CA     88
# 5    Frank   30    NY     57

比較演算子で条件を指定して行を抽出

pandasでは比較演算子を使って以下のように行を抽出できる。

print(df[df['age'] < 25])
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

query()メソッドを使うと同様の条件を文字列で指定できる。列名に対して条件を記述する。

print(df.query('age < 25'))
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

条件文字列の中で変数を使用するには変数名の前に@をつける。

val = 25
print(df.query('age < @val'))
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

f文字列で変数を埋め込んだ文字列を生成して使ってもよい。

print(f'age < {val}')
# age < 25

print(df.query(f'age < {val}'))
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

Pythonの条件指定のように2つの比較演算子で範囲を指定可能。

print(df.query('30 <= age < 50'))
#     name  age state  point
# 1    Bob   42    CA     92
# 5  Frank   30    NY     57

列同士で比較したり、算術演算子で計算して比較したりすることもできる。

print(df.query('age < point / 3'))
#       name  age state  point
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

一致、不一致は==, !=。条件文字列内の文字列は引用符で囲む必要があるので注意。

シングルクォート'で囲んだ文字列内ではダブルクォート"、ダブルクォート"で囲んだ文字列内ではシングルクォート'を使用できる。バックスラッシュ\でエスケープすれば同じ記号を使用可能。

print(df.query('state == "CA"'))
#       name  age state  point
# 1      Bob   42    CA     92
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

@で変数を使う場合は引用符は不要だが、f文字列はあくまでも文字列なので引用符が必要。

s = 'CA'
print(df.query('state != @s'))
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

# print(df.query(f'state != {s}'))
# UndefinedVariableError: name 'CA' is not defined

print(df.query(f'state != "{s}"'))
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

文字列に対する部分一致の条件には後述の文字列メソッドを使う。

条件文字列内で利用できるシンタックスについての詳細は公式ドキュメントを参照。

in演算子で特定の値を含む行を抽出(isin()と同等)

Seriesisin()メソッドは各要素が引数のリストに含まれているか判定しbool型のSeriesを返す。これを利用して、ある列の要素が特定の値に一致する行を抽出できる。

print(df[df['state'].isin(['NY', 'TX'])])
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

query()メソッドではinを使って同等の処理が可能。

print(df.query('state in ["NY", "TX"]'))
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

条件文字列内での特別な使い方として、リストに対する==も同等の処理となる。

print(df.query('state == ["NY", "TX"]'))
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

リストの変数を用いることもできる。@やf文字列を使う。

l = ['NY', 'TX']
print(df.query('state in @l'))
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

print(df.query(f'state in {l}'))
#     name  age state  point
# 0  Alice   24    NY     64
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

文字列メソッドで条件を指定して行を抽出

文字列に対する完全一致の条件は上述の==inで指定できるが、部分一致の条件は文字列メソッドstr.xxx()を使う。

以下のようなメソッドがある。

  • str.contains(): 特定の文字列を含む
  • str.endswith(): 特定の文字列で終わる
  • str.startswith(): 特定の文字列で始まる
  • str.match(): 正規表現のパターンに一致する

ブーリアンインデックスに比べて大幅に簡潔になる訳ではないが、これらもquery()で使用できる。

print(df[df['name'].str.endswith('e')])
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70

print(df.query('name.str.endswith("e")'))
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70

文字列以外の型dtypeの列はastype()で文字列型strに変換することで文字列メソッドを使える。これもquery()で指定可能。

print(df.query('age.astype("str").str.endswith("8")'))
#       name  age state  point
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70

欠損値NaNやNoneがある場合の注意

欠損値NaNNoneがある列に対して文字列メソッドを条件とするとエラーになる。

df_nan = df.copy()
df_nan.at[0, 'name'] = float('NaN')
print(df_nan)
#       name  age state  point
# 0      NaN   24    NY     64
# 1      Bob   42    CA     92
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70
# 4    Ellen   24    CA     88
# 5    Frank   30    NY     57

# print(df_nan.query('name.str.endswith("e")'))
# ValueError: unknown type object

文字列メソッドの多くは引数naで欠損値を置き換える値を指定できる。引数naTrueとすると欠損値の行も抽出され、Falseとすると欠損値の行は抽出されない。

print(df_nan[df_nan['name'].str.endswith('e', na=True)])
#       name  age state  point
# 0      NaN   24    NY     64
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70

print(df_nan[df_nan['name'].str.endswith('e', na=False)])
#       name  age state  point
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70

query()でも同様に引数naを指定すればよい。

print(df_nan.query('name.str.endswith("e", na=False)'))
#       name  age state  point
# 2  Charlie   18    CA     70
# 3     Dave   68    TX     70

行名indexに対する条件

行名indexに対する条件はindexを使って指定可能。

print(df.query('index % 2 == 0'))
#       name  age state  point
# 0    Alice   24    NY     64
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

indexに名前が付いている場合はその名前でもindexでもどちらでもよい。

紛らわしくて申し訳ないが、以下の例ではindexnameという名前を付けている。

df_name = df.set_index('name')
print(df_name)
#          age state  point
# name                     
# Alice     24    NY     64
# Bob       42    CA     92
# Charlie   18    CA     70
# Dave      68    TX     70
# Ellen     24    CA     88
# Frank     30    NY     57

print(df_name.index.name)
# name

print(df_name.query('name.str.endswith("e")'))
#          age state  point
# name                     
# Alice     24    NY     64
# Charlie   18    CA     70
# Dave      68    TX     70

print(df_name.query('index.str.endswith("e")'))
#          age state  point
# name                     
# Alice     24    NY     64
# Charlie   18    CA     70
# Dave      68    TX     70

列名columnsに対する条件で列を抽出したい場合は以下の記事を参照。

複数条件

ブーリアンインデックスで複数条件を指定する場合は以下のように記述する。

print(df[(df['age'] < 25) & (df['point'] > 65)])
#       name  age state  point
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

query()メソッドだと以下のように書ける。条件ごとの括弧は必要なく、AND(かつ)は&でもandでもどちらでもよい。

print(df.query('age < 25 & point > 65'))
#       name  age state  point
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

print(df.query('age < 25 and point > 65'))
#       name  age state  point
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

OR(または)も|でもorでもどちらでもよい。

print(df.query('age < 20 | point > 80'))
#       name  age state  point
# 1      Bob   42    CA     92
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

print(df.query('age < 20 or point > 80'))
#       name  age state  point
# 1      Bob   42    CA     92
# 2  Charlie   18    CA     70
# 4    Ellen   24    CA     88

NOT(否定)はnot

print(df.query('not age < 25 and not point > 65'))
#     name  age state  point
# 5  Frank   30    NY     57

3つ以上での条件も同様だが、&のほうが|より優先順位が高いなど順番によって結果が異なるので、先に処理したい部分を明示的に括弧で囲んだほうが無難。

print(df.query('age == 24 | point > 80 & state == "CA"'))
#     name  age state  point
# 0  Alice   24    NY     64
# 1    Bob   42    CA     92
# 4  Ellen   24    CA     88

print(df.query('(age == 24 | point > 80) & state == "CA"'))
#     name  age state  point
# 1    Bob   42    CA     92
# 4  Ellen   24    CA     88

スペースやドットを含んだ列名は「`」で囲む

query()メソッドを使う際は列名に注意が必要。例として以下のように列名を変更する。

df_rename = df.set_axis(['0name', 'age.year', 'state name', 3], axis=1)
print(df_rename)
#      0name  age.year state name   3
# 0    Alice        24         NY  64
# 1      Bob        42         CA  92
# 2  Charlie        18         CA  70
# 3     Dave        68         TX  70
# 4    Ellen        24         CA  88
# 5    Frank        30         NY  57

Pythonの変数名として有効でない列名はそのまま使うとエラーになる。例えば、数字で始まる列名、.やスペースが含まれる列名はエラー。

# print(df_rename.query('0name.str.endswith("e")'))
# SyntaxError: invalid syntax

# print(df_rename.query('age.year < 25'))
# UndefinedVariableError: name 'age' is not defined

# print(df_rename.query('state name == "CA"'))
# SyntaxError: invalid syntax

「`」で囲む必要がある。

print(df_rename.query('`0name`.str.endswith("e")'))
#      0name  age.year state name   3
# 0    Alice        24         NY  64
# 2  Charlie        18         CA  70
# 3     Dave        68         TX  70

print(df_rename.query('`age.year` < 25'))
#      0name  age.year state name   3
# 0    Alice        24         NY  64
# 2  Charlie        18         CA  70
# 4    Ellen        24         CA  88

print(df_rename.query('`state name` == "CA"'))
#      0name  age.year state name   3
# 1      Bob        42         CA  92
# 2  Charlie        18         CA  70
# 4    Ellen        24         CA  88

数値の列名は「`」で囲んでもエラー。ブーリアンインデックスを使った条件指定であれば問題ない。

# print(df_rename.query('3 > 75'))
# KeyError: False

# print(df_rename.query('`3` > 75'))
# UndefinedVariableError: name 'BACKTICK_QUOTED_STRING_3' is not defined

print(df_rename[df_rename[3] > 75])
#    0name  age.year state name   3
# 1    Bob        42         CA  92
# 4  Ellen        24         CA  88

引数inplaceで元のオブジェクトを更新

これまでの例ではquery()で行を抽出した新たなDataFrameが返され、元のオブジェクトはそのまま。引数inplaceTrueとすると元のオブジェクト自体が変更される。

df.query('age > 25', inplace=True)
print(df)
#     name  age state  point
# 1    Bob   42    CA     92
# 3   Dave   68    TX     70
# 5  Frank   30    NY     57

パーサーとエンジン(numexpr)

query()は内部でpd.eval()を使っている。

パーサー

pd.eval()ではパーサーとして'pandas''python'を選択可能。query()でも引数parserで指定できる。デフォルトは'pandas'

デフォルトの'pandas'&andでつなげるときに各条件式を括弧で囲む必要がないなど、より直感的に使えるようになっている。

エンジン

pd.eval()では式を評価するエンジンとしてnumexprを使用可能。1万行を超えるような大規模なデータを処理する場合はnumexprを使うと速くなる(らしい)。

numexprはpipでインストールできる。(環境によってはpip3

$ pip install numexpr

query()の引数engine'python''numexpr'かを選択できる。デフォルトはNoneでnumexprが利用できればnumexprが、できなければpythonが使われる。

以前のバージョン(numexprが2.6.5、pandasが0.23.0)ではnumexprで文字列メソッドが使えず、引数engine='python'とする必要があったように、バージョンによっては利用できる機能が異なる場合があるので注意。

関連カテゴリー

関連記事