Apache Beam管道中的组元素

Question

我有一个管道可以解析AVRO文件中的记录。

我需要将传入的记录分成500个项目的块，以便调用一个同时接受多个输入的API。

有没有办法用PythonSDK做到这一点？

Answer 1

我假设您是指Batch用例。您有几个选项：

如果您的PCollection足够大，并且您对包的大小有一定的灵活性，则可以在按随机/循环顺序为元素分配键后使用GroupByKey变换。e、 g.：

my_collection = p | ReadRecordsFromAvro()

element_bundles = (my_collection 
                     # Choose a number of keys that works for you (I chose 50 here)
                   | 'AddKeys' >> beam.Map(lambda x: (randint(0, 50), x))
                   | 'MakeBundles' >> beam.GroupByKey()
                   | 'DropKeys' >> beam.Map(lambda (k, bundle): bundle)
                   | beam.ParDo(ProcessBundlesDoFn()))

其中ProcessBundlesDoFn是这样的：

class ProcessBundlesDoFn(beam.DoFn):
  def process(self, bundle):
    while bundle.has_next():
      # Fetch in batches of 500 until you're done
      result = fetch_n_elements(bundle, 500)
      yield result

如果您需要所有捆绑包都包含500个元素，那么您可能需要：

1. 计算PCollection中的元素数

希望这有帮助。