Google Colab の Spark にて Google Cloud Storage と連携する方法

概要

Google Colab の Spark から Google Cloud Storage (GCS)と連携する方法を共有します。Google Cloud 上で GCS から初めて、 Google Colab 上の Spark から書き込みと読み込みを実施する手順です。

環境構築

GCS のバケットを作成

今回は us-cental1 リージョンで作成しています。

gsutil URIの値を控えます。

Service Account の作成

コンソールにて｀IAM & Admin->Service accounts`を選択します。

+ CREATE SERVICE ACCOUNTを選択します。

手順に則り Service Account を作成します。

Service Account の認証情報を取得

作成した Service Account を表示して、KEYSタブにある ADD KEY -> Create new keyを選択します。

JSONをチェックして、CREATEを選択します。

json をダウンロードが開始されるため、ローカル環境に保存します。

GCS のバケットに対する権限を Service Account に付与

作成した GCS のバケットを表示後、PERMISSONSタブにあるGRANT ACCESSを選択します。

New principalsに Service Principal のメールアドレスを入力し、ROLEにストレージに対する適切な権限（例：Storage Admin）を設定した上でSaveを選択します。

Google Colab 上での Spark による操作

Google Drive 上に Service Account の key ファイルを配置

Google Colab 上で新規ノートブック作成後に、 PySpark のバージョンを確認

!pyspark --version

Google Drive を Google Colab のノートブックにマウント

from google.colab import drive
drive.mount('/content/drive')

左側の ファイルタブを選択し、アップロードした Service Account の key ファイルのパスをコピー

パスを変数にセット

gcs_key_path = "/content/drive/MyDrive/Colab Notebooks/bigquery_iceberg/gcs_key.json"

SparkSession の定義

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder
        .appName("gcs_test")
        .config("spark.jars", "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar")
        .config("spark.hadoop.fs.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem")
        .config("google.cloud.auth.service.account.enable", "true")
        .config("spark.hadoop.google.cloud.auth.service.account.enable", "true")
        .config("spark.hadoop.google.cloud.auth.service.account.json.keyfile", gcs_key_path)
        .getOrCreate()
)
spark

GCS のパスを変数にセット

gcs_base_path = "gs://manabianqiita"

Spark から GCS に書き込みを実施

data = [("Alice", 34), ("Bob", 45), ("Cathy", 29)]
columns = ["Name", "Age"]
df = spark.createDataFrame(data, columns)

df.write.mode("overwrite").parquet(f"{gcs_base_path}/output_data3")

GCS 上で parquet ファイルが書き込まれていることを確認できます。

Spark から GCS から読み込みを実施

df = spark.read.parquet(f"{gcs_base_path}/output_data3")
df.show()