🐼
DataFrameを快適にクエリするライブラリを作った
はじめに
pandas.DataFrameを使ったデータ分析をするときに、queryを使って中身を眺めることが良くあります。queryは強力ですが、抽出したい条件をクエリ文字列で記述する必要があります。シンプルな条件であれば苦もないですが、少し複雑な条件になると、f-stringを駆使したり、str.containsなどの関数が必要になり、書き辛さを感じていました。特に単純文字列なのでIDEの補完が効かない不便さもあります。
そこで、クエリを手軽に生成できるライブラリpandaqを作りました。
Getting Started
インストール
pip install pandaq
準備
Titanic号のデータセットを例にします。
import pandas as pd
df = pd.read_csv('https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv')
| PassengerId | Survived | Pclass | Name | Sex | Age | ... | Fare | Cabin | |
|---|---|---|---|---|---|---|---|---|---|
| 0 | 1 | 0 | 3 | Braund, Mr. Owen Harris | male | 22 | ... | 7.25 | nan |
| 1 | 2 | 1 | 1 | Cumings, Mrs. John Bradley (Florence Briggs Thayer) | female | 38 | ... | 71.2833 | C85 |
| ... | ... | ... | ... | ... | ... | ... | ... | ... | ... |
使い方
A. query-stringを生成する方法
from pandaq import Q
qstr = Q().q(PassengerId=1) # -> "PassengerId==1"
df.query(qstr)
B. qメソッドを追加する方法
import pandaq.patch
df.q(PassengerId=1)
例題
以下のようなdiff形式で紹介します。
- # pandas標準のクエリ
+ # pandaqの記法
生存者
DataFrameのcolumn名を、キーワード引数で与えることができます。
- df.query("Survived==1")
+ df.q(Survived=1)
PassengerId Survived Pclass ... Fare Cabin Embarked
1 2 1 1 ... 71.2833 C85 C
2 3 1 3 ... 7.9250 NaN S
3 4 1 1 ... 53.1000 C123 S
...
CabinがD
文字列の完全一致検索です。
- df.query('Cabin=="D"')
+ df.q(Cabin="D")
PassengerId Survived Pclass ... Fare Cabin Embarked
292 293 0 2 ... 12.8750 D C
327 328 1 2 ... 13.0000 D S
473 474 1 2 ... 13.7917 D C
CabinがD35もしくはD36
リストでinを表現します。
- df.query('Cabin in ["D35", "D36"]')
+ df.q(Cabin=["D35", "D36"])
PassengerId Survived Pclass ... Fare Cabin Embarked
215 216 1 1 ... 113.2750 D36 C
248 249 1 1 ... 52.5542 D35 S
393 394 1 1 ... 113.2750 D36 C
871 872 1 1 ... 52.5542 D35 S
Brienさん
?のprefixがあると部分一致検索になります。
- df.query('Name.str.contains("Brien", regex=False, na=False)')
+ df.q(Name="?Brien")
PassengerId Survived Pclass Name Sex Age SibSp Parch Ticket Fare Cabin Embarked
186 187 1 3 O'Brien, Mrs. Thomas (Johanna "Hannah" Godfrey) female NaN 1 0 370365 15.5000 NaN Q
364 365 0 3 O'Brien, Mr. Thomas male NaN 1 0 370365 15.5000 NaN Q
552 553 0 3 O'Brien, Mr. Timothy male NaN 0 0 330979 7.8292 NaN Q
maleじゃないBrienさん
!のprefixは非完全一致です。このように引数は組み合わせることもできます。
- df.query('(Name.str.contains("Brien", regex=False, na=False) & Sex!="male")')
+ df.q(Name="?Brien", Sex="!male")
PassengerId Survived Pclass Name Sex Age SibSp Parch Ticket Fare Cabin Embarked
186 187 1 3 O'Brien, Mrs. Thomas (Johanna "Hannah" Godfrey) female NaN 1 0 370365 15.5000 NaN Q
JackかRose
/のprefixは正規表現です。
- df.query('Name.str.contains("Jack|Rose", regex=True, na=False)')
+ df.q(Name="/Jack|Rose")
PassengerId Survived Pclass Name Sex Age SibSp Parch Ticket Fare Cabin Embarked
766 767 0 1 Brewe, Dr. Arthur Jackson male NaN 0 0 112379 39.60 NaN C
855 856 1 3 Aks, Mrs. Sam (Leah Rosen) female 18.0 0 1 392091 9.35 NaN S
71才以上の人
タプルで条件を表現します。(条件のオペレータ、値)と書きます。
- df.query("Age >= 71")
+ df.q(Age=(">=", 71))
PassengerId Survived Pclass Name Sex Age SibSp Parch Ticket Fare Cabin Embarked
96 97 0 1 Goldschmidt, Mr. George B male 71.0 0 0 PC 17754 34.6542 A5 C
493 494 0 1 Artagaveytia, Mr. Ramon male 71.0 0 0 PC 17609 49.5042 NaN C
630 631 1 1 Barkworth, Mr. Algernon Henry Wilson male 80.0 0 0 27042 30.0000 A23 S
851 852 0 3 Svensson, Mr. Johan male 74.0 0 0 347060 7.7750 NaN S
71才以上72才未満の人
タプルで条件は組み合わせできます。(条件のオペレータ1、値1、 条件のオペレータ2、値2、・・・)と書きます。
- df.query("71 <= Age < 72")
+ df.q(Age=(">=", 71, "<", 72))
PassengerId Survived Pclass Name Sex Age SibSp Parch Ticket Fare Cabin Embarked
96 97 0 1 Goldschmidt, Mr. George B male 71.0 0 0 PC 17754 34.6542 A5 C
493 494 0 1 Artagaveytia, Mr. Ramon male 71.0 0 0 PC 17609 49.5042 NaN C
色々な書き方
どれも同じクエリですが、色々な書き方ができます。
df.q(Survived=1, Pclass=3, Fare=(">", 50))
# メソッドチェーン的
df.q(Survived=1).q(Pclass=3).q(Fare=(">", 50))
# 辞書渡し
df.q({"Survived": 1,
"Pclass": 3,
"Fare": (">", 50)})
おわりに
我ながら中途半端で奇妙なツールを作ってしまった自負はあるのですが、私用では便利に使っています。また、ryeやCodecovを導入してみたり、初めてPyPIに登録したりで、いい勉強になりました。GitHub eholic/pandaqにも、ドキュメントを記載してるのでご覧ください。
Discussion