如何使用 PCollection 在光束管道中构建数据存储查询？

Question

我正在构建一个光束管道，以按照一些规则将 csv 文件处理/转换为 xml 文件。到目前为止，我的方法是在管道开始时按行拆分输入 csv 文件，并将每一行输入到管道中。在管道中，将每一行转换为一个 xml 标记，在管道结束时，我将全局上的所有内容合并到最终的 xml 文件中。现在的问题是我需要存储在 Google Datastore 中的一些额外信息来为 csv 文件中的每一行构建 xml 标记，我不知道该怎么做，因为从数据存储区检索数据的查询是运行时参数(https://beam.apache.org/releases/pydoc/2.16.0/apache_beam.io.gcp.datastore.v1new.datastoreio.html#apache_beam.io.gcp.datastore.v1new.datastoreio.ReadFromDatastore) 并且查询取决于 PCollection。我需要构建一个这样的查询：

Select from xxx where id in PCollection

有没有办法做到这一点？比如 combineGlobally 来构建查询，然后以某种方式将查询传递给 ReadFromDatastore 函数？或者有什么方法可以做我需要的吗？

我现在有这样的事情：

with beam.Pipeline(options=pipeline_options) as p:
    items = (
        p |
        'ReadCsvFile' >> beam.io.Read(CsvFileSource(input_name))
        'PrepareToJoin' >> beam.ParDo(PrepareToJoin())
    )

    datastore_items = (p |
        'DatastoreDataP' >> ReadFromDatastore(project_id, query)
    )

    new_items = (
        {'data': items, 'datastore': datastore_items} |
        'JoinWithDatastore' >> beam.CoGroupByKey() |
        'PostJoinProcess' >> beam.ParDo(PostJoinProcess())
    )

    xml_file = (new_items |
        'ItemToXmlTag' >> beam.ParDo(ItemToXmlTag()) |
        'MakeXMLFile' >> beam.CombineGlobally(XMLCombineFn()) |
        WriteToText(output_name)
    )

如您所见，查询是管道中的一个参数。

提前感谢您的帮助！

Answer 1

您可以使用 ParDo 来实现此目的。下面是如何编写 ParDo 的伪代码，其中您处理 PCollection 的每个元素，然后进行 DataStore 查找

class EnrichEntities(beam.DoFn):
"""Updates Datastore entity"""
def process(self, element):
    key = client.key('Task', 'sample_task')
    task = client.get(element.key)

        return [element]
    else:
        return [element]

然后在您的管道中，您可以使用 ParDo，如下所示

with beam.Pipeline() as pipeline:
  results = (
      pipeline
      | 'Create inputs' >> beam.Create()
      | 'DoFn methods' >> beam.ParDo(EnrichEntities())
      | beam.Map(print)

为了提高查找的性能，您还可以使用BagState，其中您可以一次查找 N 条记录。这是一个实现这个https://beam.apache.org/blog/2017/08/28/timely-processing.html的链接