如何在Dataproc集群的PySpark任务中写入Cloud Bigtable及解决依赖问题

我来帮你解决这个问题——在Dataproc的PySpark任务里用Google Cloud Bigtable客户端写入数据，同时搞定依赖包的部署问题。你遇到的wheel转egg失效的情况很常见，很多Python包的结构并不适配egg格式，我们换更可靠的方式来部署依赖，再给你一个可运行的代码示例。

一、解决依赖包部署问题的几种可行方案 #

方案1：集群初始化时预安装依赖包 #

这是最省心的方式，创建集群时通过初始化脚本让所有节点自动安装google-cloud-bigtable：

1. 创建一个初始化脚本（比如install_bigtable_deps.sh），内容如下：

#!/bin/bash
# 确保pip是针对集群默认Python版本的
sudo pip3 install google-cloud-bigtable

2. 将脚本上传到你的GCS存储桶，比如gs://your-bucket-name/scripts/install_bigtable_deps.sh
3. 创建Dataproc集群时指定初始化动作：

gcloud dataproc clusters create your-cluster-name \
  --region your-region \
  --initialization-actions gs://your-bucket-name/scripts/install_bigtable_deps.sh \
  --service-account your-service-account@your-project.iam.gserviceaccount.com

注意：服务账号需要拥有Bigtable的读写权限，建议绑定roles/bigtable.user角色

方案2：提交任务时动态加载依赖包 #

如果不想重新创建集群，可以在提交PySpark任务时上传wheel包并加载：

1. 把你下载的google-cloud-bigtable wheel包上传到GCS，比如gs://your-bucket-name/deps/google_cloud_bigtable-2.23.0-py2.py3-none-any.whl
2. 提交任务时用--archives参数挂载wheel包，同时在脚本中添加路径：

gcloud dataproc jobs submit pyspark your-write-script.py \
  --cluster your-cluster-name \
  --region your-region \
  --archives gs://your-bucket-name/deps/google_cloud_bigtable-2.23.0-py2.py3-none-any.whl#bigtable_deps

3. 在你的PySpark脚本开头添加路径加载：

import sys
# 挂载的wheel包会被解压到指定的目录下，把路径加入sys.path
sys.path.insert(0, './bigtable_deps')

二、PySpark写入Bigtable的代码示例 #

注意：Bigtable客户端不能被序列化，所以必须在每个分区的mapPartitions算子内初始化客户端，避免Spark序列化报错。

from pyspark.sql import SparkSession
from google.cloud import bigtable
from google.cloud.bigtable.row import DirectRow

def batch_write_to_bigtable(partition_data):
    # 每个分区初始化一次Bigtable客户端（客户端不可序列化，必须在分区内创建）
    project_id = "your-gcp-project-id"
    instance_id = "your-bigtable-instance-id"
    table_id = "your-bigtable-table-id"

    client = bigtable.Client(project=project_id, admin=True)
    instance = client.instance(instance_id)
    table = instance.table(table_id)

    # 批量写入分区内的数据
    batch = table.mutate_rows()
    for row_item in partition_data:
        # 假设每个row_item是(row_key, column_family, column_name, value)的元组
        row_key, cf, col, val = row_item
        row = DirectRow(row_key.encode("utf-8"))
        # 设置单元格数据，第三个参数可以指定时间戳（可选）
        row.set_cell(cf, col, val.encode("utf-8"))
        batch.add(row)
    
    # 提交批量写入并处理异常
    try:
        responses = batch.commit()
        for resp in responses:
            if resp.code != 0:
                print(f"行写入失败，错误码: {resp.code}, 行键: {resp.row_key.decode('utf-8')}")
    except Exception as e:
        print(f"批量写入出错: {str(e)}")
    
    return []

if __name__ == "__main__":
    # 初始化SparkSession
    spark = SparkSession.builder \
        .appName("PySpark-Bigtable-Writer") \
        .getOrCreate()

    # 替换成你的业务数据，可以从HDFS/GCS/其他数据源读取
    sample_data = [
        ("user_1001", "profile", "name", "Alice"),
        ("user_1002", "profile", "name", "Bob"),
        ("user_1003", "activity", "last_login", "2024-05-20"),
        ("user_1004", "activity", "last_login", "2024-05-19")
    ]

    # 并行化数据并调用批量写入函数
    rdd = spark.sparkContext.parallelize(sample_data)
    # 用mapPartitions处理每个分区，避免重复创建客户端
    rdd.mapPartitions(batch_write_to_bigtable).count()

    spark.stop()

三、关键注意事项 #

• 客户端初始化位置：一定要在mapPartitions内创建Bigtable客户端，不能在Driver端创建后广播，因为客户端对象无法被Spark序列化到Executor节点。
• 批量写入优化：使用mutate_rows进行批量写入，比单条写入效率高很多，适合Spark的并行处理场景。
• 权限配置：确保Dataproc集群的服务账号拥有Bigtable实例的bigtable.tables.mutateRows权限，否则会出现权限拒绝错误。

内容的提问来源于stack exchange，提问作者MANISH ZOPE