Python: pandas と Google BigQuery を連携させる

ぶっちゃけ pandas は大規模なデータセットを扱うのが苦手だ。だいたい一桁 GB なら我慢と工夫で何とかなるけど、二桁 GB を超えると現実的な処理時間で捌けなくなってくる。そこで、今回は pandas を Google BigQuery と連携させることで重たい処理をオフロードする方法を試してみる。

使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.13.5
BuildVersion:   17F77
$ python -V         
Python 3.6.5
$ pip list --format=columns | grep -i pandas
pandas                   0.23.3
pandas-gbq               0.5.0

Google BigQuery を使う下準備

ここから BigQuery を使うための下準備が結構長いので、既に使ったことがある人は読み飛ばしてもらって構わない。

まずは Homebrew Cask を使って GCP を操作するためのコマンドラインツールをインストールしておく。

$ brew cask install google-cloud-sdk

インストールしたら自分の GCP アカウントで認証する。

$ gcloud init

認証ができたら pandas との連携を試すためのプロジェクトを新しく作っておこう。

$ gcloud projects create pandas-bq

プロジェクトを作ったら支払いの設定を有効にする。

console.cloud.google.com

続いてプロジェクトにデータセットとテーブルを作る。

$ bq mk pandas-bq:example
Dataset 'pandas-bq:example' successfully created.
$ bq mk pandas-bq:example.iris
Table 'pandas-bq:example.iris' successfully created.

名前から分かる通り、みんな大好き Iris データセットを読み込む。

UCI のリポジトリから Iris データセットをダウンロードする。

$ wget https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data

上記からダウンロードしたデータセットには末尾に空行が入っているので取り除く。

$ brew install gnu-sed
$ alias sed="gsed"
$ sed -i -e "/^$/d" iris.data

空行が入っていると bq コマンドが無理やり内容を解釈しようとするので。

続いてテーブルのスキーマを用意する。

$ cat << 'EOF' > schema.json
[
  {"name": "sepal_length", "type": "float", "mode": "required"},
  {"name": "sepal_width", "type": "float", "mode": "required"},
  {"name": "petal_length", "type": "float", "mode": "required"},
  {"name": "petal_width", "type": "float", "mode": "required"},
  {"name": "label", "type": "string", "mode": "required"}
]
EOF

あとはデータをアップロードする。

$ bq load --replace --project_id pandas-bq example.iris iris.data schema.json

これで、次の通りデータが読み込まれた。

$ bq show pandas-bq:example.iris
Table pandas-bq:example.iris

   Last modified                 Schema                 Total Rows   Total Bytes   Expiration   Time Partitioning   Labels  
 ----------------- ----------------------------------- ------------ ------------- ------------ ------------------- -------- 
  05 Aug 20:50:57   |- sepal_length: float (required)   150          7100                                                   
                    |- sepal_width: float (required)                                                                        
                    |- petal_length: float (required)                                                                       
                    |- petal_width: float (required)                                                                        
                    |- label: string (required)                                                                             
                                                                
$ bq query --project_id pandas-bq "SELECT * FROM example.iris LIMIT 10"
Waiting on bqjob_r457fdda779482a15_0000016509ed844f_1 ... (0s) Current status: DONE   
+--------------+-------------+--------------+-------------+-----------------+
| sepal_length | sepal_width | petal_length | petal_width |      label      |
+--------------+-------------+--------------+-------------+-----------------+
|          5.1 |         2.5 |          3.0 |         1.1 | Iris-versicolor |
|          5.0 |         2.0 |          3.5 |         1.0 | Iris-versicolor |
|          5.7 |         2.6 |          3.5 |         1.0 | Iris-versicolor |
|          6.0 |         2.2 |          4.0 |         1.0 | Iris-versicolor |
|          5.8 |         2.6 |          4.0 |         1.2 | Iris-versicolor |
|          5.5 |         2.3 |          4.0 |         1.3 | Iris-versicolor |
|          6.1 |         2.8 |          4.0 |         1.3 | Iris-versicolor |
|          5.5 |         2.5 |          4.0 |         1.3 | Iris-versicolor |
|          6.4 |         3.2 |          4.5 |         1.5 | Iris-versicolor |
|          5.6 |         3.0 |          4.5 |         1.5 | Iris-versicolor |
+--------------+-------------+--------------+-------------+-----------------+

これで Google BigQuery 側の準備は整った。

pandas で BigQuery を使えるようにする下準備

続いて pandas と Google BigQuery を連携させるための準備を進める。

まずは必要なパッケージとして pandas と pandas-gbq をインストールする。

$ pip install pandas pandas-gbq

Python の REPL を起動する。

$ python

pandas をインポートする。

>>> import pandas as pd

これで下準備が整った。

pandas と BigQuery を連携させる

さっき作ったプロジェクトの ID と BigQuery で実行したいクエリを用意する。

>>> project_id = 'pandas-bq'
>>> query = 'SELECT * FROM example.iris LIMIT 10'

クエリは先ほど実行したものと同じ。

あとは pandas.read_gbq() に実行したいクエリとプロジェクトの ID を渡すだけ。すると認証の URL が表示するので、それをブラウザで開く。認証が成功すると認証コードが表示されるので、それを REPL に貼り付けよう。

>>> pd.read_gbq(query, project_id, dialect='standard')
Please visit this URL to authorize this application: https://accounts.google.com/o/oauth2/auth?...
Enter the authorization code: ...
sepal_length  sepal_width       ...         petal_width            label
0           5.1          2.5       ...                 1.1  Iris-versicolor
1           5.0          2.0       ...                 1.0  Iris-versicolor
2           5.7          2.6       ...                 1.0  Iris-versicolor
3           6.0          2.2       ...                 1.0  Iris-versicolor
4           5.8          2.6       ...                 1.2  Iris-versicolor
5           5.5          2.3       ...                 1.3  Iris-versicolor
6           6.1          2.8       ...                 1.3  Iris-versicolor
7           5.5          2.5       ...                 1.3  Iris-versicolor
8           6.4          3.2       ...                 1.5  Iris-versicolor
9           5.6          3.0       ...                 1.5  Iris-versicolor

[10 rows x 5 columns]

これでクエリが実行される。

二回目の実行からは結果がキャッシュされるので、認証について聞かれることはない。

>>> query = 'SELECT COUNT(1) AS count FROM example.iris'
>>> pd.read_gbq(query, project_id, dialect='standard')
   count
0    150

pandas.read_gbq() の結果は DataFrame として得られる。

>>> query = '''
... SELECT
...   AVG(sepal_length) AS sepal_length_avg,
...   AVG(sepal_width) AS sepal_width_avg,
...   AVG(petal_length) AS petal_length_avg,
...   AVG(petal_width) AS petal_width_avg
... FROM example.iris
... '''
>>> new_df = pd.read_gbq(query, project_id, dialect='standard')
>>> type(new_df)
<class 'pandas.core.frame.DataFrame'>
>>> new_df
   sepal_length_avg       ...         petal_width_avg
0          5.843333       ...                1.198667

[1 rows x 4 columns]

DataFrame を BigQuery に書き込む

先ほどの例は BigQuery のテーブルにクエリを発行して結果を読み出すものだった。今度はローカルの DataFrame の内容を BigQuery に書き出してみる。

サンプルの DataFrame を用意する。

>>> data = [
...   ('alice', 150),
...   ('bob', 160),
...   ('carol', 170),
... ]
>>> df = pd.DataFrame(data, columns=['name', 'height'])

あとは DataFrame#to_gbq() メソッドを実行する。このときデータセット名、テーブル名、プロジェクト ID を指定する。

>>> df.to_gbq('example.users', project_id)

コマンドラインツールで確認すると、ちゃんとテーブルができてデータが入っていることが分かる。

$ bq show pandas-bq:example.users
Table pandas-bq:example.users

   Last modified          Schema         Total Rows   Total Bytes   Expiration   Time Partitioning   Labels  
 ----------------- -------------------- ------------ ------------- ------------ ------------------- -------- 
  05 Aug 21:05:36   |- name: string      3            43                                                     
                    |- height: integer                                                                       

$ bq query --project_id pandas-bq "SELECT * FROM example.users"
+-------+--------+
| name  | height |
+-------+--------+
| alice |    150 |
| bob   |    160 |
| carol |    170 |
+-------+--------+