将具有许多元素的“PCollection”转换为单个元素答案

【问题标题】：Transforming `PCollection` with many elements into a single element将具有许多元素的“PCollection”转换为单个元素
【发布时间】：2022-08-16 06:35:21
【问题描述】：

我正在尝试将具有许多元素的PCollection 转换为具有一个元素的PCollection。基本上，我想从：

[1,2,3,4,5,6]

至：

[[1,2,3,4,5,6]]

这样我就可以在DoFn 中处理整个PCollection。

我试过CombineGlobally(lamdba x: x)，但一次只有一部分元素组合成一个数组，结果如下：

[1,2,3,4,5,6] -> [[1,2],[3,4],[5,6]]

或者类似的东西。

这是我试图运行的脚本的相关部分：

import apache_beam as beam
raw_input = range(1024)
def run_test():
    with TestPipeline() as test_pl:
        input = test_pl | \"Create\" >> beam.Create(raw_input)
        def combine(x):
            print(x)
            return x
        
        (
            input
            | \"Global aggregation\" >> beam.CombineGlobally(combine)
        )
        pl.run()

run_test()

标签： apache-beam

【解决方案1】：

我想出了一个非常轻松的方法来做到这一点，我在docs 中错过了它：

组合元素的更通用方式，也是最灵活的方式是使用从 CombineFn 继承的类。

CombineFn.create_accumulator()：这将创建一个空累加器。为了例如，求和的空累加器为 0，而空的累加器为乘积（乘法）的累加器将为 1。

CombineFn.add_input()：每个元素调用一次。需要一个蓄能器和一个输入元素，将它们组合起来并返回更新的累加器。

CombineFn.merge_accumulators()：多个累加器可以是并行处理，因此此功能有助于将它们合并为单蓄电池。

CombineFn.extract_output()：它允许做额外的计算在提取结果之前。

我想提供一个简单地将其参数传递给“香草”CombineGlobally 的 lambda 函数不会达到我最初的预期。该功能必须由我指定（尽管我仍然认为它没有内置在 API 中很奇怪）。

你可以找到更多关于子类化CombineFnhere，我觉得这很有帮助：
CombineFn 指定如何在全部或部分 PCollection 可以合并为一个值——本质上是提供与 Python “reduce”的参数相同的信息内置（输入参数除外，它是组合FnProcessContext）。合并过程如下：
输入值被分成一个或多个批次。

对于每个批次，都会调用 create_accumulator 方法来创建一个新的初始“累加器”值，该值表示零值。

对于批次中的每个输入值，调用 add_input 方法以将更多值与该批次的累加器组合。

调用 merge_accumulators 方法将来自不同批次的累加器组合成单个组合输出累加器值，一旦所有的累加器都在它们的所有输入值批次添加到他们。重复调用此操作，直到只剩下一个累加器值。
在最终累加器上调用 extract_output 操作以获取输出值。注意：如果此 CombineFn 与变换一起使用具有默认值的应用程序将在扩展时间以获取默认值。
因此，通过继承 CombineFn，我编写了这个简单的实现，Aggregated，这正是我想要的：
import apache_beam as beam
raw_input = range(1024)


class Aggregated(beam.CombineFn):
    def create_accumulator(self):
        return []
    
    def add_input(self, accumulator, element):
        accumulator.append(element)
        return accumulator
    
    def merge_accumulators(self, accumulators):
        merged = []
        for a in accumulators:
            for item in a:
                merged.append(item)
        return merged
    
    def extract_output(self, accumulator):
        return accumulator


def run_test():
    with TestPipeline() as test_pl:
        input = test_pl | "Create" >> beam.Create(raw_input)
        (
            input
            | "Global aggregation" >> beam.CombineGlobally(Aggregated())
            | "print" >> beam.Map(print)
        )
        pl.run()

run_test()

【讨论】：

【解决方案2】：

您还可以使用side inputs 完成您想要的操作，例如

with beam.Pipeline() as p:
    pcoll = ...
    (p
     # Create a PCollection with a single element.
     | beam.Create([None])
     # This will process the singleton exactly once,
     # with the entirity of pcoll passed in as a second argument as a list.
     | beam.Map(
        lambda _, pcoll_as_side: ...consume pcoll_as_side here...,
        pcoll_as_side=beam.pvalue.AsList(pcoll))

【讨论】：