如何在Apache Beam Python中并行化多个ReadFromBigQuery操作

我需要在Apache Beam Python管道中并行化多个ReadFromBigQuery操作，所有待读取的BigQuery表结构一致，最终合并为一个数据集。尝试过以下方案但遇到问题：

• 使用for循环读取（测试用4张表，生产环境会有数百张）：部署到Dataflow后无法扩缩容，仅激活一个Worker；
• 使用ReadAllFromBigQuery：忽略实验性警告后，读取大量表时出现数据重复问题；
• 尝试用ParDo实现：因对Beam了解不足，无法在DoFn类中使用读取函数，不知如何操作。

之前尝试的代码 #

For循环读取代码 #

tables_list = [table1, table2, table3, table4]
query = 'SELECT id, value FROM `my-project.my-dataset.{table_name}`'

with beam.Pipeline(options=pipeline_options) as p:
    input_tables = [(
        p
        | 'Read {table}'.format(table = t) >> ReadFromBigQuery(
            query = query.format(table_name = t),
            use_standard_sql = True
        )
    ) for t in tables_list]
    final_table = (
        input_tables
        | 'Flatten in one table' >> beam.Flatten()
    )

ReadAllFromBigQuery代码 #

tables_list = [table1, table2, table3, table4]
query = 'SELECT id, value FROM `my-project.my-dataset.{table_name}`'

def to_bq_request(t, query):
    from apache_beam.io import ReadFromBigQueryRequest
    return ReadFromBigQueryRequest(query = query.format(table_name = t))

with beam.Pipeline(options=pipeline_options) as p:
    pre_table = (
        p
        | 'Create Table' >> beam.Create(tables_list)
    )
    final_table = (
        pre_table
        | 'Map before ReadAll' >> beam.Map(to_bq_request, query = query)
        | 'Read All Tables ' >> ReadAllFromBigQuery(temp_dataset = 'tds')
    )

可行解决方案 #

1. 修复For循环扩缩容问题 #

for循环方式无法扩缩容是因为每个ReadFromBigQuery都作为独立的根变换，Dataflow默认会将这些小变换打包到同一个Worker执行。可通过调整资源配置和并行度参数解决：

• 在PipelineOptions中设置--number_of_worker_harness_threads（每个Worker的线程数），比如设为8；
• 为每个Read变换启用flatten_results=False，让BigQuery返回分片文件，提升读取并行度：

# 修改ReadFromBigQuery部分
ReadFromBigQuery(
    query=query.format(table_name=t),
    use_standard_sql=True,
    flatten_results=False  # 开启分片读取，提升并行能力
)

同时通过--max_num_workers配置足够的Worker数量。

2. 解决ReadAllFromBigQuery数据重复问题 #

数据重复通常源于临时表未正确清理或请求幂等性不足，可通过以下方式修复：

• 为每个请求添加唯一的job_id_prefix，确保读取任务的唯一性；
• 显式指定temp_location，确保路径权限正确且临时文件自动清理；
• 修改请求生成函数：

def to_bq_request(t, query):
    from apache_beam.io import ReadFromBigQueryRequest
    return ReadFromBigQueryRequest(
        query=query.format(table_name=t),
        job_id_prefix=f"read-table-{t}"  # 每个表对应唯一的Job前缀
    )

同时确保使用较新版本的Beam，旧版本的ReadAllFromBigQuery存在重复数据bug。

3. 使用ParDo实现并行读取（正确方式） #

DoFn内不能直接使用ReadFromBigQuery（它是PTransform而非可调用函数），但可以用BigQuery Python客户端库直接读取，注意控制并发避免限流：

from google.cloud import bigquery
import apache_beam as beam

class ReadBigQueryTable(beam.DoFn):
    def setup(self):
        # Worker初始化时创建客户端，避免重复实例化
        self.client = bigquery.Client()
    
    def process(self, table_name, query_template):
        query = query_template.format(table_name=table_name)
        query_job = self.client.query(query)
        # 迭代返回查询结果
        for row in query_job.result():
            yield {"id": row.id, "value": row.value}

# 管道执行代码
tables_list = [table1, table2, table3, table4]
query_template = 'SELECT id, value FROM `my-project.my-dataset.{table_name}`'

with beam.Pipeline(options=pipeline_options) as p:
    final_table = (
        p
        | 'Create Table List' >> beam.Create(tables_list)
        | 'Read Each Table' >> beam.ParDo(ReadBigQueryTable(), query_template=query_template)
    )

注意：需确保Dataflow服务账号拥有BigQuery访问权限，可通过--max_num_workers调整并行读取的Worker数量。

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