python - 如何为 Apache Beam 中的复合变换提供参数？

Question

我正在使用 Apache Beam 的 Python SDK。

我有几个转换步骤并希望使它们可重用，这使我可以编写一个自定义的复合转换，如下所示:

class MyCompositeTransform(beam.PTransform):
def expand(self, pcoll, arg1, kwarg1=u'default'):
    result = (pcoll
              | 'Step 1' >> beam.Map(lambda f: SomeFn(f, arg1))
              | 'Last step' >> beam.Map(lambda f: SomeOtherFn(f, kwarg1))
              )
    return result

我想要的是提供一些额外的参数 arg1和 kwarg1内部的其他转换需要这些。但我不知道这是否是一种有效的方式，也不知道如何在管道中使用它。

有人可以指出我的方向吗？

Answer 1

您可以通过 PTransform 提供参数构造函数。参数也可以采用侧输入的形式(即来自另一个变换的数据输出)。这是一个同时使用“正常”参数和侧输入的示例。

from typing import Dict, Any, Iterable
import apache_beam as beam


class MyCompositeTransform(beam.PTransform):

    def __init__(self, my_arg, my_side_input):
        super().__init__()
        self.my_arg= my_arg
        self.my_side_input= my_side_input

    @staticmethod
    def transform(
        element: Dict[str, Any], my_arg: int, my_side_input: Iterable[int]
    ) -> Dict[str, Any]:
        pass

    def expand(self, pcoll):
        return pcoll | "MyCompositeTransform" >> beam.Map(
            MyCompositeTransform.transform,
            self.my_arg,
            beam.pvalue.AsIter(self.my_side_input),
        )

使用 beam.pvalue 定义侧输入如何传递给变换，例如它是单个值吗， Iterable或具体化为 List ?

来自 Beam 的其他示例:(参见 PTransform) https://beam.apache.org/releases/pydoc/2.20.0/_modules/apache_beam/transforms/stats.html