Databricks にて BigQuery から JDBC により Spark データフレームを作成する方法

## 概要

Databricks にて BigQuery から JDBC により Spark データフレームを作成する方法を共有します。本記事で紹介する方法を検証目的としており、多くのケースにおいて`format`オプションを`bigquery`で指定して Spark コネクターを利用する方が最適です。

:::note info
本記事では Serverless クラスターで実行することを想定しておりません。
:::

本記事は下記シリーズの一部です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/1bb4805e-9f67-49e9-8e72-36036bf4310c.png)

出所：[Databricks で BigQuery のデータを取得する方法の徹底ガイド - Qiita](https://qiita.com/manabian/items/7097a9c31de23752d58a)

## 事前準備

### BigQuery に対する認証情報を取得

下記の記事を参考に BigQuery に対する認証情報（キー）を取得してください。

- [BigQuery × Databricks ：サンドボックス環境と Databricks Community edition で始める無償で高速データ分析 #Python - Qiita](https://qiita.com/manabian/items/a1e9d2619c4afa221251)

### BigQuery の JDBC の完全修飾クラス名の取得

`/databricks/jars`配下にある BigQuery の JDBC ファイルのパスを取得します。

```bash
%sh
ls -1 /databricks/jars/*GoogleBigQueryJDBC*.jar
```

```text
/databricks/jars/----ws_3_5--third_party--bigquery-jdbc--bigquery-driver-shaded---846918551--GoogleBigQueryJDBC42.jar
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b94a85f3-face-4eef-985f-b4a795bd7b05.png)

BigQuery の JDBC の完全修飾クラス名（`bigquery.shaded.com.simba.googlebigquery.jdbc42.Driver`）を取得します。

```bash
%sh
jar tf /databricks/jars/----ws_3_5--third_party--bigquery-jdbc--bigquery-driver-shaded---846918551--GoogleBigQueryJDBC42.jar | grep jdbc42/Driver.class
```

```python
import zipfile
import re

jar_path = "/databricks/jars/----ws_3_5--third_party--bigquery-jdbc--bigquery-driver-shaded---846918551--GoogleBigQueryJDBC42.jar"

with zipfile.ZipFile(jar_path) as zf:
    # namelist() で JAR 内のファイル一覧を取得
    for entry in zf.namelist():
        # jdbc42/Driver.class だけを対象にする
        if entry.endswith("jdbc42/Driver.class"):
            # / → . に置換し、末尾の .class を削除
            class_name = re.sub(r"\.class$", "", entry.replace("/", "."))
            print(class_name)
```

```text
bigquery.shaded.com.simba.googlebigquery.jdbc42.Driver
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/bf1221fb-7786-46ce-8326-706292910b52.png)

:::note info

`jdbc42`という指定値は変更される可能性があり、上記コードの値も対応させる必要があります。

:::

## 手順

### 1. `OAuthPvtKeyPath`プロパティを指定する手順

#### 1-1. 認証情報を変数にセット

```python
json_str = """{json_key}"""
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/f2538dfc-65a7-4617-9e4f-31d868294aeb.png)

#### 1-2. Spark データフレームを作成

```python
# Google Cloud の プロジェクト ID をセット
project_id = "axial-triode-XXXXX"

# サービスアカウントのメールアドレスをセット
service_account_mail = "dbxfederation01@axial-triode-XXXXX.iam.gserviceaccount.com"

# 取得した driver の値をセット
driver = "bigquery.shaded.com.simba.googlebigquery.jdbc42.Driver"
```

```python
# データ取得元のテーブル名を設定
table_name = "bigquery-public-data.google_analytics_sample.ga_sessions_20170801"

jdbc_url = (
    "jdbc:bigquery://https://www.googleapis.com/bigquery/v2;"
    f"ProjectId={project_id};"
    "OAuthType=0;"
    f"OAuthServiceAcctEmail={service_account_mail};"
    f"OAuthPvtKey={json_str};"
    "LogLevel=0;"
    "Timeout=600;"
)

bq_df = (
    spark.read.format("jdbc")
    .option("url", jdbc_url)
    .option("driver", driver)
    .option("dbtable", table_name)
    .load()
)

bq_df.limit(10).display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/61ddda77-dfd3-4166-bb16-6abd89b83ab6.png)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/49f05c7d-b774-45e4-9611-c10919bdf11c.png)

:::note warn

Serverless compute for notebooks で実行した場合には下記のエラーが発生します。

> The input query contains unsupported data source(s). Only csv, json, avro, delta, kafka, parquet, orc, text, unity_catalog, binaryFile, xml, simplescan, iceberg, mysql, postgresql, sqlserver, redshift, snowflake, sqldw, databricks, bigquery, oracle, salesforce, salesforce_data_cloud, teradata, workday_raas, mongodb data sources are supported on serverless compute, and only csv, json, avro, delta, kafka, parquet, orc, text, unity_catalog, binaryFile, xml, simplescan, iceberg data sources are allowed to run DML on serverless compute.

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/488ae5ac-0ed2-47bd-a85b-1364124fb017.png)

:::

### 2. `OAuthPvtKeyPath`プロパティを指定する手順

#### 2-1. Databricks Runtime のバージョンを確認

```python
print(spark.conf.get("spark.databricks.clusterUsageTags.sparkVersion"))
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/03a4436a-697f-4432-902c-448f9157a783.png)

#### 2-2. Volume 配下に認証情報（キー）を配置して、パスを取得

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/1b745392-5467-4721-823c-88979077d0ae.png)

#### 2-3. 認証情報のキーを FileStore のディレクトリに配置

```python
# Volume に配置したキーのファイルパスを指定
src_key_path = "/Volumes/bq_work/default/bq-keys/axial-triode-XXXXX-1ba03e864e42.json"
local_key_path  = "/FileStore/bq/axial-triode-XXXXX-1ba03e864e42.json"

dbutils.fs.cp(src_key_path, local_key_path, True)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/bd00fd1d-9351-4326-810e-74ae12e1734b.png)

#### 2-4. Spark データフレームを作成

```python
# Google Cloud の プロジェクト ID をセット
project_id = "axial-triode-XXXXX"

# サービスアカウントのメールアドレスをセット
service_account_mail = "dbxfederation01@axial-triode-XXXXX.iam.gserviceaccount.com"

# 取得した driver の値をセット
driver = "bigquery.shaded.com.simba.googlebigquery.jdbc42.Driver"
```

```python
# データ取得元のテーブル名を設定
table_name = "bigquery-public-data.google_analytics_sample.ga_sessions_20170801"

spark_key_path = "/dbfs" + local_key_path
jdbc_url = (
    "jdbc:bigquery://https://www.googleapis.com/bigquery/v2;"
    f"ProjectId={project_id};"
    "OAuthType=0;"
    f"OAuthServiceAcctEmail={service_account_mail};"
    f"OAuthPvtKeyPath={spark_key_path};"
    "LogLevel=0;"
    "Timeout=600;"
)

bq_df = (
    spark.read.format("jdbc")
    .option("url", jdbc_url)
    .option("driver", driver)
    .option("dbtable", table_name)
    .load()
)

bq_df.limit(10).display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/c10bdfbc-15ca-4684-9e7c-3b44781e52b1.png)