PandasでSQLクエリ風にデータのサブセットを抽出する2つの方法

はじめに

この記事では、Pandasを使ってSQLスタイルのフィルタリングでデータ分析を行う方法を解説します。企業のデータの多くはデータベースに格納されており、その取得や操作にはSQLが欠かせません。たとえばOracle、IBM、Microsoftといった各社は、独自のSQL実装を持つデータベース製品を提供しています。

データは常にCSVファイルとして保存されているとは限らないため、データサイエンティストはキャリアの中で必ず一度はSQLと向き合うことになります。筆者自身も、会社のデータの大半がOracleに保存されているため、Oracleをよく利用しています。

シナリオ：条件に合う映画を抽出する

映画データセットから、以下の条件をすべて満たす映画を抽出するタスクが与えられたとしましょう。

• 映画の言語が英語（en）またはスペイン語（es）であること
• 人気度（popularity）が500以上1000以下であること
• ステータス（status）が「Released（公開済み）」であること
• 投票数（vote_count）が5000より大きいこと

このシナリオをSQLで書くと、次のようになります。

SELECT title AS movie_title
      ,original_language AS movie_language
      ,popularity AS movie_popularity
      ,status AS movie_status
      ,vote_count AS movie_vote_count
FROM movies_data
WHERE original_language IN ('en', 'es')
  AND status = 'Released'
  AND popularity BETWEEN 500 AND 1000
  AND vote_count > 5000;

SQLでの書き方がイメージできたところで、同じ処理をPandasで段階的に実現してみましょう。ここでは2つの方法を紹介します。

方法1：ブールインデックス（Boolean Indexing）

ステップ1：movies_dataデータセットをDataFrameに読み込む

import pandas as pd
movies = pd.read_csv("https://raw.githubusercontent.com/sasankac/TestDataSet/master/movies_data.csv")

ステップ2：各条件を変数に代入する

languages = ["en", "es"]

condition_on_languages = movies.original_language.isin(languages)
condition_on_status = movies.status == "Released"
condition_on_popularity = movies.popularity.between(500, 1000)
condition_on_votecount = movies.vote_count > 5000

ステップ3：すべての条件（ブール配列）を組み合わせる

final_conditions = (condition_on_languages
                    & condition_on_status
                    & condition_on_popularity
                    & condition_on_votecount)

columns = ["title", "original_language", "status", "popularity", "vote_count"]

# すべてを組み合わせて抽出
movies.loc[final_conditions, columns]

実行結果は以下の通りです。

方法2：.query()メソッド

.query()メソッドは、SQLのWHERE句のような感覚でデータをフィルタリングできる便利な機能です。条件を文字列として渡せる点が特徴ですが、カラム名に空白を含めてはいけないという制約があります。

カラム名に空白が含まれている場合は、Pythonのreplace関数などを使ってアンダースコア（_）に置き換えておきましょう。

なお、筆者の経験上、大規模なDataFrameに対しては、このquery()メソッドの方が前述のブールインデックスよりも高速に動作することが多いです。

ステップ4：クエリ文字列を作成してメソッドを実行する

注意：.query()メソッドは、複数行にまたがるトリプルクォート文字列では正しく動作しません。

final_conditions = (
    "original_language in ['en','es'] "
    "and status == 'Released' "
    "and popularity > 500 "
    "and popularity < 1000 "
    "and vote_count > 5000"
)

final_result = movies.query(final_conditions)
final_result

実行結果：

@記号でPython変数を参照する

さらに実用的なテクニックとして、実際のコーディングでは「in」句でチェックしたい値が多数になるケースがよくあります。その場合、値を直接文字列に埋め込む上記の書き方は扱いづらくなります。

そこで活躍するのが@（アットマーク）記号です。これを使うと、クエリ文字列の中からPythonの変数を参照できます。値をPythonのリストとしてプログラム的に生成し、それを@で参照することも可能です。

movie_languages = ['en', 'es']

final_conditions = (
    "original_language in @movie_languages "
    "and status == 'Released' "
    "and popularity > 500 "
    "and popularity < 1000 "
    "and vote_count > 5000"
)

final_result = movies.query(final_conditions)
final_result

結果は先ほどと同じになります。

まとめ

Pandasでは、①ブールインデックスと②.query()メソッドの2つのアプローチで、SQLのWHERE句のような条件絞り込みが実現できます。

• ブールインデックス：条件を個別の変数として管理できるため、可読性が高くデバッグしやすいのがメリット。
• .query()メソッド：SQLライクな文字列で直感的に書けるうえ、大規模データでは高速に動作する傾向がある。

@記号による変数参照を活用すれば、query()メソッドでも柔軟に条件を組み立てられます。データの規模や可読性の要件に応じて、両者を使い分けてみてください。