DataFrameレシピ: データ抽出条件


PandasのDataFrameはたまにしか使っていませんでした。いつもググりながら使っていましたが、本格的に使うに当たり、整理をしてみました。query関数でできることや注意点、ブールインデックスとの比較などを整理しています。
似た内容として記事「DataFrameレシピ: 行列指定して出力」locilocプロパティなどを使った行列指定方法も書いています。

まとめ(早見表)

私が参照したい部分を一番上に持ってきました。全般的にquery関数の方がシンプルに記述できますが、文字列はブールインデックスの方がやや読みやすいです。

文字列・数値共通

条件 query ブールインデックス
等号 df.query('C == "b"') df[df['C'] == 'b']
不等号 df.query('C != "b"') df[df['C'] != 'b']
否定 df.query('not C == "b"') df[~(df['C'] == 'b')]
AND df.query('A > 0 and C == "b"')
または
df.query('A > 0 & C == "b"')
df[(df['A'] > 0)&(df['C'] == 'b')]
OR df.query('A > 0 or C == "b"')
または
df.query('A > 0 | C == "b"')
df[(df['A'] > 0)|(df['C'] == 'b')]
論理演算優先 df.query('(A > 0 or A == 0)
and C == "b"')
df[((df['A'] > 0)|(df['A'] == 0))
&(df['C'] == 'b')]
配列 df.query('A in [10, 20]') df[df['A'].isin([10, 20])]
変数使用 df.query('A == @var') df[df['A']==var]
null判定 df.query('New == New') df[df['New'].isnull()]
null以外 df.query('New != New') df[df['New'].notnull()]

null判定に関してはdf.query('New.isnull()', engine='python')でもいいのですが、長いので同じ列名使った方がシンプル。

数値系

条件 query ブールインデックス
比較 df.query('A > 0') df[df['A']>0]
等号あり比較 df.query('A >= 0') df[df['A']>=0]
範囲 df.query('-10 <= A <= 0') df[(df['A'] >= -10)&(df['A'] <= 0)]
算術使用 df.query('A < B * -10') df[df['A'] < df['B'] * -10]

範囲のブールインデックスはもっとシンプルにできる方法あるかもしれません(あまり調べていない)。

文字列

strメソッドと同じなんでしょうね。

条件 query ブールインデックス
前方一致 df.query('D.str.startswith("a")', engine='python') df[df['D'].str.startswith('a')]
中間一致 df.query('D.str.contains("c")', engine='python') df[df['D'].str.contains('c')]
後方一致 df.query('D.str.endswith("b")', engine='python') df[df['D'].str.endswith('b')]
正規表現 df.query('D.str.match(".b.")', engine='python') df[df['D'].str.match('.b.')]

文字列のengine

engineはpythonのみ使用可能。ライブラリmumexprがインストールされていない場合は当然、pythonが選択されます。しかし、インストール時にはmumexprが選択されますが、文字列クエリに使用できないためエラーが起きます。インストール有無によってエラーが起こるのは堅牢性が低いので、engine='python'を加えるのがベター

NoneとNaN

列にNoneや欠損値NaNがあると上記命令はエラーとなるので注意。パラメータ'na'を使い欠損値があった場合に抽出するかのオプションが可能。

ブールインデックスでNoneや欠損値処理
# Noneや欠損値は抽出対象
df[df['D'].str.endswith('d', na=True)]
df.query('D.str.endswith("d", na=True)', engine='python')

# Noneや欠損値は抽出対象外
df[df['D'].str.endswith('d', na=False)]
df.query('D.str.endswith("d", na=False)', engine='python')

containsと正規表現

regex=Falseにしないと正規表現になってしまうので注意。

> df[df['D'].str.contains('a.', regex=False)]
> df.query('D.str.contains("a.", regex=False)', engine='python')
    A  B  C     D
0 -20  1  a  aa.a
> df[df['D'].str.contains('a.', regex=True)]
> df.query('D.str.contains("a.", regex=True)', engine='python')

    A  B  C     D
0 -20  1  a  aa.a
4  20  5  a   abc

正規表現は、記事「ゼロから覚えるPython正規表現の基本とTips」参照です。

環境

2021年2月にGoogle Colaboratory使っています。そのため、Pythonやそのパッケージはそのままのバージョンで使っています。

種類 バージョン
Python 3.6.9

上記環境で、以下のPython追加パッケージバージョンでした。

種類 バージョン
jupyter 1.0.0
numexpr 2.7.2
numpy 1.19.5
pandas 1.1.5

前提データ

以下のDataFrameで上記の命令を実行しました。

import pandas as pd
import numpy as np

df = pd.DataFrame({'A': [-20, -10, 0, 10, 20],
                   'B': [1, 2, 3, 4, 5],
                   'C': ['a', 'b', 'b', 'b', 'a'],
                   'D': ['aa.a', 'bb.b', None, 'bcd', 'abc']})
df
    A  B  C     D
0 -20  1  a  aa.a
1 -10  2  b  bb.b
2   0  3  b  None
3  10  4  b   bcd
4  20  5  a   abc

query関数とブールインデックス

データ抽出するのに2つの方法があり、query関数とブールインデックスです。
query関数はSQLのような書き方でシンプルに条件を書けます(文字列は少し長くなりがち)。一方、ブールインデックスは汎用的に使え使用可能範囲が広いです。

query関数

処理エンジン

query()の引数engineで'python'か'numexpr'かを選択できます。指定しない場合、numexprがインストールされていればnumexprが、無ければpythonが使われる。numexprは処理が速いらしい。
ただ、文字列系処理はengine='python'でないと使えずにエラーが起きます。
numexprライブラリはnumpyのための高速evaluatorです。

返り値

パラメータinplaceがFalseの場合(デフォルトはFalse)にはDataFrameのオブジェクトを返します。ブールインデックスがTrue/Falseを返すのと対照的です。

df_copied = df.query('(A > 0 or A == 0) and C == "b"')
# df.query('(A > 0 or A == 0) and C == "b"', inplace=True)にするとdf自体が変わる

注意点

以下がquery関数の注意点です。

  • Seriesには使えない
  • 列名が数字の場合はNGらしい(未確認)
  • 列名にスペースを含む場合はバッククォートで囲みましょう

参考リンク