如何在Dataflow/Beam中扩展顺序输出的数据源？

问题 #

在Dataflow场景中，若使用的数据源是顺序输出的（比如Go SDK的BigQueryIO实现，通过循环读取查询结果并顺序发射记录），会导致Dataflow无法扩展worker数量，后续消费该数据源的ParDo步骤也始终无法并行，还会触发掉队任务检测。这类问题的核心原因是自定义数据源未实现进度追踪与分片能力。

对应的自定义DoFn Java伪代码如下：

@ProcessElement
public void processElement(@Element byte[] input, OutputReceiver<T> output) throws Exception {
        // ..
        result = bigquery.query(queryConfig);
        for (FieldValueList row : result.iterateAll()) {    
            output.output(mapRowToType(row));
        }
}

管道示例：

var bigqueryRows = pipeline.apply("ReadFromBigQuery", BigQueryIO.read(//...));
var mutations = bigqueryRows.apply("ProcessRows", ParDo.of(/**/ {
    @ProcessElement
    public void processElement(@Element ItemRow row, OutputReceiver<BigtableIO.Write.Mutation> out) {

        String rowKey = row.getId(); 
        BigtableIO.Write.Mutation mutation = BigtableIO.Write.Mutation.create(rowKey);
        // ...
        out.output(mutation);
    }
}));

在无法通过其他方式消费数据的前提下，如何处理这类顺序数据源，让处理过程能并行到多个worker？

解决方案 #

1. 拆分查询为可并行的分片任务 #

将原本的单一大查询拆分为多个独立的分片查询，通过分片参数（比如时间范围、分区键、分片ID）让每个查询只处理一部分数据。把这些分片参数作为初始PCollection的输入，每个参数对应一个ParDo任务去执行查询并发射结果，这样就能让多个worker并行处理不同分片。

比如针对BigQuery，可以按日期分区拆分查询，生成多个查询条件，再将这些条件作为输入分发：

// 生成分片参数，比如按日期范围拆分
PCollection<String> queryShards = pipeline.apply(Create.of(
    "SELECT * FROM table WHERE date >= '2024-01-01' AND date < '2024-01-08'",
    "SELECT * FROM table WHERE date >= '2024-01-08' AND date < '2024-01-15'",
    ...
));

// 每个分片独立执行查询
PCollection<ItemRow> shardedRows = queryShards.apply(ParDo.of(new DoFn<String, ItemRow>() {
    @ProcessElement
    public void process(@Element String query, OutputReceiver<ItemRow> out) {
        var result = bigquery.query(query);
        for (var row : result.iterateAll()) {
            out.output(mapRowToType(row));
        }
    }
}));

2. 实现支持分片与进度追踪的自定义源 #

如果必须用自定义数据源，不要用普通DoFn，而是实现Beam的BoundedSource（批处理）或UnboundedSource（流处理）接口：

• 重写split方法：将整个数据集拆分为多个可独立处理的分片，每个分片可以分配给不同worker。
• 重写getProgress方法：实时报告当前分片的处理进度，让Dataflow能识别任务状态，动态调度worker并检测掉队任务。

核心逻辑是把查询结果按行号或范围拆分，每个分片对应一个查询的LIMIT和OFFSET（或者更高效的分区键过滤），确保每个分片的处理是独立的。

3. 引入中间存储做数据中转 #

如果无法直接拆分查询，先把顺序读取的全量数据写入支持并行读取的分布式存储（比如GCS的多文件输出、BigQuery分区表），再从这个存储并行读取数据：

1. 先用单任务读取全量数据，写入GCS时按一定规则拆分多个文件（比如按行数、大小）。
2. 再用Beam的TextIO/ParquetIO等支持并行读取的Transform读取这些文件，后续ParDo步骤就能自动扩展到多个worker。

4. 优先使用官方并行化读取实现 #

尽量替换成Beam官方提供的、已实现并行能力的数据源Transform。比如Java版的BigQueryIO.read()默认会自动分片处理查询结果，而Go SDK如果当前版本的BigQueryIO不够完善，可以先通过BigQuery的导出功能把数据导出到GCS，再用Go SDK的TextIO或AvroIO并行读取，避免自定义顺序读取的DoFn。

内容的提问来源于stack exchange，提问作者dre-hh