Apache Beam Dataflow运行报错：找不到beam_nuggets等第三方模块

Dataflow Worker节点默认没有预装beam_nuggets这类第三方库，而DirectRunner是在本地环境运行（已安装该库），因此会出现模块找不到的错误。以下是几种解决方法：

方法1：使用requirements文件指定依赖 #

这是最简单的解决方案，通过指定依赖文件让Dataflow自动在Worker节点安装所需包：

1. 创建requirements.txt文件，包含所需依赖：

beam-nuggets>=0.16.0
psycopg2-binary>=2.9.0  # PostgreSQL驱动，beam_nuggets依赖此包

2. 提交Dataflow作业时添加--requirements_file参数：

python your_script.py \
  --runner=DataflowRunner \
  --project=你的GCP项目ID \
  --region=你的区域 \
  --requirements_file=requirements.txt \
  # 其他必要参数（如host、port、dbname等）

如果是通过代码配置PipelineOptions，可直接设置：

from apache_beam.options.pipeline_options import PipelineOptions, DataflowOptions

options = PipelineOptions()
dataflow_options = options.view_as(DataflowOptions)
dataflow_options.requirements_file = "requirements.txt"

方法2：使用自定义Worker镜像 #

如果依赖安装存在兼容性问题，可构建包含所有依赖的自定义Docker镜像：

1. 创建Dockerfile：

# 基于对应版本的Beam官方Python镜像
FROM apache/beam_python:2.46.0
RUN pip install beam-nuggets>=0.16.0 psycopg2-binary>=2.9.0

2. 构建并推送到GCR（Google Container Registry）：

docker build -t gcr.io/你的GCP项目ID/beam-worker:latest .
docker push gcr.io/你的GCP项目ID/beam-worker:latest

3. 提交作业时指定自定义镜像：

python your_script.py \
  --runner=DataflowRunner \
  --project=你的GCP项目ID \
  --region=你的区域 \
  --worker_harness_container_image=gcr.io/你的GCP项目ID/beam-worker:latest \
  # 其他必要参数

方法3：替换为Beam官方JDBC IO（避免第三方依赖） #

如果不想依赖beam_nuggets，可使用Apache Beam官方的JDBC读取能力（需使用Beam 2.40+版本）：

import apache_beam as beam
from apache_beam.io.jdbc import ReadFromJdbc

# ... 其他代码保持不变 ...

def execute_pipeline(options: IngestionOptions) -> None:
    # ... 其他变量初始化 ...

    with beam.Pipeline(options=options) as p:
        records = p | "Reading from PostgreSQL" >> ReadFromJdbc(
            query=query,
            driver_class_name="org.postgresql.Driver",
            jdbc_url=f"jdbc:postgresql://{options.host}:{options.port}/{options.dbname}",
            username=options.user,
            password=options.password,
        )
        # ... 后续转换和写入逻辑保持不变 ...

额外注意 #

原代码中存在一个小错误：file_name_suffix=".arvo"应为file_name_suffix=".avro"，需修正避免生成错误后缀的文件。

内容的提问来源于stack exchange，提问作者Roman Dryndik