け日記

最近はPythonでいろいろやってます

小ネタ: Pandasでqueryを使って行を選択する

PandasのDataFrameから行を抽出する簡便な方法として、queryメソッドが提供されています。
SQLで言えば選択 (WHERE句) にあたる処理を、文字列で記述できます。

pandas.pydata.org

queryメソッドを使った選択の例

今回の投稿で使うデータを準備します。

import pandas as pd

df = pd.DataFrame({
    'name': ['荒岩 一味', '荒岩 虹子', '荒岩 まこと', '荒岩 みゆき', '田中 一', '梅田 よしお'],
    'age': [44, 44, 21, 13, 36, 31],
    'sex': ['M', 'F', 'M', 'F', 'M', 'M'],
})

以下のように第1引数に条件文を記述することで、行を選択できます。完全一致で選択する場合は、列名 == 値とします。

df.query("name == '荒岩 一味'")
#    age   name sex
# 0   44  荒岩 一味   M

条件文内で@変数名とすると、Pythonの変数の値を参照できます。

search_name = '田中 一'

df.query("name == @search_name")
#    age  name sex
# 4   36  田中 一   M

複数条件で選択することもできます。&でAND条件、|でOR条件となります。ちなみにBETWEEN句などはありません。

df.query("sex == 'M' & age >= 35")
#    age   name sex
# 0   44  荒岩 一味   M
# 4   36   田中 一   M

df.query("age >= 40 | age < 20")
#    age    name sex
# 0   44   荒岩 一味   M
# 1   44   荒岩 虹子   F
# 3   13  荒岩 みゆき   F

IN句も以下のようにすることで記述できます。

df.query("name in ('荒岩 まこと', '荒岩 みゆき')")
#    age    name sex
# 2   21  荒岩 まこと   M
# 3   13  荒岩 みゆき   F

LIKE検索は少し特殊です。LIKE句はありませんが、以下のようにengine='python'とすることで、条件文中でSeriesのstrメソッドを利用できます。あとはstartswithメソッドを使うことで、name LIKE '荒岩%'といった選択もできます。startswith以外にも、endswith, containsなども利用できます。

df.query("name.str.startswith('荒岩')", engine='python')
#    age    name sex
# 0   44   荒岩 一味   M
# 1   44   荒岩 虹子   F
# 2   21  荒岩 まこと   M
# 3   13  荒岩 みゆき   F

queryメソッドの注意点

上のように簡単に行選択を記述できるqueryメソッドですが、Series (bool型) による選択と比較すると、若干遅いようです。速度を気にする場合は、利用を避けたほうが無難そうです。

stackoverflow.com

cocodrips.hateblo.jp