如何在Apache Beam中让PTransform接收并返回DoOutputsTuple

问题背景 #

我有一些返回多输出PCollection的DoFn，一条是正常数据路径，另一条捕获错误并重定向，示例代码如下：

output, error = (
        pcol
        | "Fix Timestamps"
        >> ParDo(ConvertTimestamp(), timestamp_field)
        .with_outputs(self.fail_tag, main=self.success_tag)
)

我希望将多个ParDo/DoFn封装成一个可直接调用的PTransform，尝试了如下实现：

from apache_beam import PTransform, ParDo, DoFn, pvalue
from datetime import datetime
from typing import Iterable, Dict, Any, Tuple

# 假设Failure和TimestampError是已定义的数据类和异常类
FAIL_TAG = 'fail'
SUCCESS_TAG = 'success'

class ConvertToBQFriendlyTypes(PTransform):
    def __init__(
        self,
        timestamp_fields: Tuple[str, ...],
        fail_tag: str = FAIL_TAG,
        success_tag: str = SUCCESS_TAG,
    ):
        super().__init__()
        self.fail_tag = fail_tag
        self.success_tag = success_tag
        self.timestamp_fields = timestamp_fields

    class _ConvertSingleTimestamp(DoFn):
        def __init__(self, fail_tag: str = FAIL_TAG):
            super().__init__()
            self.fail_tag = fail_tag

        def process(
            self,
            element: Dict[str, Any],
            field_name: str,
        ) -> Iterable[Dict[str, Any]] | Iterable[pvalue.TaggedOutput]:
            timestamp_raw = element[field_name]
            if hasattr(timestamp_raw, "to_utc_datetime"):
                timestamp_utc = timestamp_raw.to_utc_datetime(has_tz=True)  # type: ignore
            else:
                timestamp_utc = timestamp_raw
            if hasattr(timestamp_utc, "timestamp"):
                timestamp_utc = datetime.fromtimestamp(round(timestamp_utc.timestamp()))
            if hasattr(timestamp_utc, "strftime"):
                result = timestamp_utc.strftime("%Y-%m-%d %H:%M:%S.%f")  # type: ignore
            elif isinstance(timestamp_utc, str) or timestamp_utc is None:
                result = timestamp_utc
            else:
                result = Failure( # `Failure`是包含pipeline_step/element/exception的简单数据类
                    pipeline_step="ConvertToBQFriendlyTypes",
                    element=element,
                    exception=TimestampError(
                        f'Field "{field_name}" has no means to convert time to a '
                        "string, which is needed for writing to BigQuery."
                    ),
                )
            if isinstance(result, Failure):
                yield pvalue.TaggedOutput(self.fail_tag, result)
            else:
                element[field_name] = result
                yield element

    def expand(
        self, pcoll: PCollection[Dict[str, Any]] | pvalue.PValue
    ) -> PCollection[Dict[str, Any]] | pvalue.PValue:
        for timestamp_field in self.timestamp_fields:
            pcoll = pcoll | f'Convert "{timestamp_field}"' >> ParDo(
                self._ConvertSingleTimestamp(self.fail_tag), timestamp_field
            ).with_outputs(self.fail_tag, main=self.success_tag)
        return pcoll

但该实现失败，原因是：

• PTransform的expand方法默认只接收PCollection，无法处理DoOutputsTuple（多输出ParDo的返回对象）
• PTransform默认只能返回PCollection，不能返回DoOutputsTuple

尝试拆分输出管理时，出现TypeError: cannot unpack non-iterable PCollection object或TypeError: 'PCollection' object is not subscriptable错误；直接操作解包后的DoOutputsTuple时，出现TypeError: '_InvalidUnpickledPCollection' object is not subscriptable错误。

曾考虑使用pasgarde包，但不想依赖低采用率的第三方库，且它仅支持Map、Flatmap和Filter，限制过多。

解决方案 #

方案一：修改PTransform以支持DoOutputsTuple的输入输出 #

自定义PTransform，在expand方法中处理输入是否为DoOutputsTuple，并最终返回统一的DoOutputsTuple。核心思路是：

1. 从输入中分离出成功流和失败流（如果输入是DoOutputsTuple）
2. 对成功流依次应用每个字段的转换，收集每一步产生的失败流
3. 合并所有失败流，返回包含最终成功流和合并后失败流的DoOutputsTuple

代码实现如下：

from apache_beam import PTransform, ParDo, DoFn, pvalue
from apache_beam.pvalue import DoOutputsTuple
from datetime import datetime
from typing import Iterable, Dict, Any, Tuple

FAIL_TAG = 'fail'
SUCCESS_TAG = 'success'

# 定义Failure和TimestampError
class Failure:
    def __init__(self, pipeline_step, element, exception):
        self.pipeline_step = pipeline_step
        self.element = element
        self.exception = exception

class TimestampError(Exception):
    pass

class ConvertToBQFriendlyTypes(PTransform):
    def __init__(
        self,
        timestamp_fields: Tuple[str, ...],
        fail_tag: str = FAIL_TAG,
        success_tag: str = SUCCESS_TAG,
    ):
        super().__init__()
        self.fail_tag = fail_tag
        self.success_tag = success_tag
        self.timestamp_fields = timestamp_fields

    class _ConvertSingleTimestamp(DoFn):
        def __init__(self, fail_tag: str = FAIL_TAG, field_name: str = ""):
            super().__init__()
            self.fail_tag = fail_tag
            self.field_name = field_name

        def process(
            self,
            element: Dict[str, Any],
        ) -> Iterable[Dict[str, Any]] | Iterable[pvalue.TaggedOutput]:
            timestamp_raw = element[self.field_name]
            if hasattr(timestamp_raw, "to_utc_datetime"):
                timestamp_utc = timestamp_raw.to_utc_datetime(has_tz=True)  # type: ignore
            else:
                timestamp_utc = timestamp_raw
            if hasattr(timestamp_utc, "timestamp"):
                timestamp_utc = datetime.fromtimestamp(round(timestamp_utc.timestamp()))
            if hasattr(timestamp_utc, "strftime"):
                result = timestamp_utc.strftime("%Y-%m-%d %H:%M:%S.%f")  # type: ignore
            elif isinstance(timestamp_utc, str) or timestamp_utc is None:
                result = timestamp_utc
            else:
                result = Failure(
                    pipeline_step=f"ConvertToBQFriendlyTypes - {self.field_name}",
                    element=element,
                    exception=TimestampError(
                        f'Field "{self.field_name}" has no means to convert time to a '
                        "string, which is needed for writing to BigQuery."
                    ),
                )
            if isinstance(result, Failure):
                yield pvalue.TaggedOutput(self.fail_tag, result)
            else:
                element[self.field_name] = result
                yield element

    def expand(self, input_pcoll):
        # 分离初始的成功流和失败流
        if isinstance(input_pcoll, DoOutputsTuple):
            success_stream = input_pcoll[self.success_tag]
            error_streams = [input_pcoll[self.fail_tag]]
        else:
            success_stream = input_pcoll
            error_streams = []

        # 对每个时间字段依次处理
        for field in self.timestamp_fields:
            transformed = (
                success_stream
                | f'Convert "{field}"' >> ParDo(
                    self._ConvertSingleTimestamp(self.fail_tag, field)
                ).with_outputs(self.fail_tag, main=self.success_tag)
            )
            success_stream = transformed[self.success_tag]
            error_streams.append(transformed[self.fail_tag])

        # 合并所有失败流
        merged_errors = (
            error_streams
            | "Flatten Error Streams" >> pvalue.Flatten()
        )

        # 返回包含成功和失败流的DoOutputsTuple
        return pvalue.DoOutputsTuple(
            main=success_stream,
            tagged={self.fail_tag: merged_errors}
        )

使用方式：

# 输入是PCollection时
result = pcol | ConvertToBQFriendlyTypes(timestamp_fields=('ts1', 'ts2'))
output = result.success
error = result.fail

# 输入是DoOutputsTuple时（比如前一个多输出Transform的结果）
prev_result = pcol | PreviousMultiOutputTransform()
new_result = prev_result | ConvertToBQFriendlyTypes(timestamp_fields=('ts1', 'ts2'))
final_output = new_result.success
final_error = new_result.fail

方案二：使用Composite Transform统一管理流 #

如果不想处理DoOutputsTuple的复杂逻辑，可以将成功流和失败流作为两个独立的输入输出，通过显式方式传递：

1. 创建Transform，接收成功流和可选的初始失败流
2. 处理成功流，收集新的失败流
3. 返回最终的成功流和合并后的失败流

代码示例：

class ConvertToBQFriendlyTypes(PTransform):
    def __init__(
        self,
        timestamp_fields: Tuple[str, ...],
        fail_tag: str = FAIL_TAG,
        success_tag: str = SUCCESS_TAG,
    ):
        super().__init__()
        self.fail_tag = fail_tag
        self.success_tag = success_tag
        self.timestamp_fields = timestamp_fields

    # _ConvertSingleTimestamp DoFn同方案一，略

    def expand(self, inputs):
        # 输入为元组：(success_stream, error_stream)
        success_stream, initial_errors = inputs

        error_streams = [initial_errors] if initial_errors else []

        for field in self.timestamp_fields:
            transformed = (
                success_stream
                | f'Convert "{field}"' >> ParDo(
                    self._ConvertSingleTimestamp(self.fail_tag, field)
                ).with_outputs(self.fail_tag, main=self.success_tag)
            )
            success_stream = transformed[self.success_tag]
            error_streams.append(transformed[self.fail_tag])

        merged_errors = error_streams | "Flatten Errors" >> pvalue.Flatten() if error_streams else None

        return success_stream, merged_errors

使用方式：

# 初始无失败流
output, error = (pcol, None) | ConvertToBQFriendlyTypes(('ts1', 'ts2'))

# 初始有失败流
prev_output, prev_error = pcol | PreviousMultiOutputTransform()
final_output, final_error = (prev_output, prev_error) | ConvertToBQFriendlyTypes(('ts1', 'ts2'))

关键要点 #

• Beam的DoOutputsTuple本质是包含多个PCollection的容器，需要显式处理其内部的流
• 合并失败流时使用pvalue.Flatten()将多个错误流合并为一个
• 自定义PTransform的expand方法可以接收任意类型的输入（不限于PCollection），只要在方法内部正确处理即可

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