如何将Apache Beam中两个PCollection的计数合并为FinalResult对象

要将两个统计得到的PCollection<Long>合并为FinalResult对象，可通过键关联+分组合并的方式实现，具体步骤及代码如下：

实现步骤 #

1. 为两个计数集合添加统一的固定键，确保后续能通过键关联；
2. 使用CoGroupByKey将带键的两个集合合并，得到包含双计数的分组结果；
3. 通过ParDo将分组结果转换为FinalResult对象。

修改后的完整代码 #

首先导入必要的依赖类：

import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.CoGbkResult;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.CoGroupByKey;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ProcessElement;
import org.apache.beam.sdk.values.KeyedPCollectionTuple;

修改build方法：

public Pipeline build() {
    // 定义TupleTag用于区分两个计数集合
    final TupleTag<Long> queryDataCountTag = new TupleTag<>();
    final TupleTag<Long> incidentesCountTag = new TupleTag<>();

    PCollection<QueryData> queryData = pipeline.apply("Read from BigQuery",
            BigQueryIO.read(reader).fromQuery(reader.getQueryString()).usingStandardSql().withoutValidation());

    PCollection<Incident> incidentes = pipeline.apply("READ CSV", TextIO.read().from("gs://bucket/incidents.csv"))
            .apply("Convert to bean", ParDo.of(new IncidentCsvToBeanFunction()));

    PCollection<Long> queryDataCount = queryData.apply("Count QueryData", Count.globally());
    PCollection<Long> incidentesCount = incidentes.apply("Count Incidentes", Count.globally());

    // 为两个计数添加统一固定键，用于后续分组合并
    PCollection<KV<String, Long>> keyedQueryCount = queryDataCount.apply(WithKeys.of("count_key"));
    PCollection<KV<String, Long>> keyedIncidentCount = incidentesCount.apply(WithKeys.of("count_key"));

    // 合并两个带键的计数集合
    PCollection<KV<String, CoGbkResult>> groupedCounts = KeyedPCollectionTuple
            .of(queryDataCountTag, keyedQueryCount)
            .and(incidentesCountTag, keyedIncidentCount)
            .apply(CoGroupByKey.create());

    // 将分组结果转换为FinalResult对象
    PCollection<FinalResult> finalResults = groupedCounts.apply("Create FinalResult", ParDo.of(new DoFn<KV<String, CoGbkResult>, FinalResult>() {
        @ProcessElement
        public void processElement(ProcessContext c) {
            CoGbkResult result = c.element().getValue();
            // 提取两个计数并转换为Integer类型
            Long queryCount = result.getOnly(queryDataCountTag);
            Long incidentCount = result.getOnly(incidentesCountTag);

            FinalResult finalResult = new FinalResult();
            finalResult.setNumberOfQueryData(queryCount.intValue());
            finalResult.setNumberOfIncidents(incidentCount.intValue());

            c.output(finalResult);
        }
    }));

    // 此处可添加将finalResults写入目标BigQuery表的逻辑
    // finalResults.apply("Write to BigQuery", BigQueryIO.writeTableRows()...);

    return pipeline;
}

关键说明 #

• TupleTag：用于在合并后的分组结果中区分两个来源的计数数据；
• WithKeys：给每个计数元素添加统一键，保证CoGroupByKey能将两个计数归到同一组；
• CoGbkResult.getOnly()：由于每个键下仅对应一个计数元素，可安全使用该方法直接获取值；
• 类型转换：将统计得到的Long类型计数转换为Integer，匹配FinalResult的字段定义。

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