Vertex AI WorkbenchからBigQueryアクセスする

おはこんハロチャオ~!なにもんなんじゃ?じょんすみすです。

当エントリー『クラスメソッド 機械学習チーム アドベントカレンダー 2022』は6日目となります。

みなさん、Vertex AI Workbench(以下Workbench)使ってますか? Vertex AIの機能の一部で、Jupyterの環境を提供してくれています。

このWorkbenchにはGoogle Cloud上で便利に使えるような拡張がされています。 今回は、その中の1つであるBigQueryとの連携周りについて見ていきましょう。

Workbench内からBigQueryにアクセスする

便利な拡張その1です。 Workbench内からBigQueryに接続し、情報を確認したりクエリを実行したりできます。

この画面上から操作できることは以下のようになります。

  • 上部の「Open SQL editor」を選択し、SQLを記述するための画面を開く
  • テーブルを選択肢てそのテーブルに関する情報を取得する
  • 下部の「Query history」からクエリの実行履歴を見る

クエリエディタを開くと以下のように任意のSQLを実行してその結果を閲覧することができます。

特定のテーブルを選択すると、そのテーブルのスキーマに関する情報やデータのプレビューを閲覧できます。

また、この画面から上部の「Query table」を選択するとあらかじめテーブルに対するSELECT分が入力されたクエリエディタを開くことができます。

その隣の「Generate Statistics」を選択すると、マジックコマンドが入力された新規のノートブックが開きます。 これはノートブックからBigQueryの統計情報を取得するものになっており、実行すると以下のような結果が表示されます。

Query historyはクエリの実行履歴になっており、過去に実行したクエリの情報を確認したり同じものをクエリエディタで開いたりといったことが可能になっています。

これらの情報の多くはBigQueryのコンソール上からも確認できるものなのでどちらを使うかはお好みによりますが、Workbench上でだけで完結できるのはありがたいですね。

マジックコマンドを使ってBigQueryにアクセス

続いてはJupyterらしくノートブックファイル内での便利機能を紹介していきます。

Workbenchのノートブックファイル内では %%bigquery というマジックコマンドが利用できます。 これを記述したセルにはSQLを直接記述することで、BigQueryへのアクセスが可能です。

PandasのDataFrameを出力したときのような形式で表示されているのがわかります。 この結果はそのままDataFrameとしてPythonから利用することも可能です。 マジックコマンドを %%bigquery のようにすることでそのまま後続の処理で変数として利用可能です。

また --param にてPythonのdictを渡してクエリ中のパラメータにすることも可能です。 %%bigquery --params $ としてdictのkeyを @key として指定することでvalueの値をパラメータとして利用できます。

WorkbenchのPython Kernelには標準で google-cloud-bigquery もインストールされているため、こちらを利用してBigQueryとPythonの間をシームレスに行き来することも可能です。

おわりに

今回は、Vertex AI WorkbenchにおけるBigQueryとの連携のための便利な仕組みを紹介しました。 Vertex AI自体はWorkbenchの利用は必須ではなく任意の場所から利用できますが、BigQueryをデータソースとした機械学習を行う際のEDAなどには非常に便利なのでぜひお試しください。