如何使用Apache Beam将单个文本文件内容拆分到不同PCollection

实现思路 #

• 首先通过带状态的DoFn逐行标记所属区间：在DoFn内部维护当前所处的区间标识，遇到BEGIN标记时更新当前区间，遇到END标记时重置区间，中间的行都绑定当前的区间标识输出。
• 然后使用Beam内置的Partition转换，按照绑定的区间标识将PCollection拆分为两个独立的PCollection，分别对应burger和fish区间的内容。

完整可运行代码 #

import apache_beam as beam
from apache_beam import Create, ParDo, Partition

class TagWithGroup(beam.DoFn):
    def __init__(self):
        # 维护当前所处的分组
        self.current_group = None

    def process(self, element):
        element = element.strip()
        if element.startswith("BEGIN="):
            # 遇到起始标记，更新当前分组，需要保留标记行可取消下一行注释
            # yield (element.split("=")[1], element)
            self.current_group = element.split("=")[1]
            return
        if element.startswith("END="):
            # 遇到结束标记，重置当前分组，需要保留标记行可取消下一行注释
            # yield (element.split("=")[1], element)
            self.current_group = None
            return
        # 当前处于有效区间内，输出分组+内容
        if self.current_group is not None:
            yield (self.current_group, element)

def partition_fn(element, num_partitions):
    # 0号分区对应burger，1号对应fish
    if element[0] == "burger":
        return 0
    elif element[0] == "fish":
        return 1

with beam.Pipeline() as p:
    all_lines = (
        p 
        # | "Read From Text" >> beam.io.ReadFromText("gs://abc.txt")
        | "Create dummy text" >> Create([
              'BEGIN=burger',
              'blue',
              'lettuce',
              'mayonise',
              'END=burger',
              'BEGIN=fish',
              'green',
              'strawberry',
              'ketchup',
              'END=fish',
          ])
        | "Tag group for each line" >> ParDo(TagWithGroup())
    )

    # 拆分为2个PCollection
    burger_collection, fish_collection = all_lines | "Split into two collections" >> Partition(partition_fn, 2)

    # 测试输出，后续可替换为各自的处理逻辑
    burger_collection | "Print burger content" >> beam.Map(lambda x: print(f"burger组内容：{x[1]}"))
    fish_collection | "Print fish content" >> beam.Map(lambda x: print(f"fish组内容：{x[1]}"))

注意事项 #

• 如果输入是多个独立文本文件，确保每个区间的完整内容都在同一个文件中，避免单个区间被拆分到多个文件导致状态丢失。
• 如果需要保留BEGIN、END标记行，取消TagWithGroup中对应yield行的注释即可。

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