如何处理Apache Beam Pipeline中的大型内存数据以在Google Dataflow Runner上运行

Question

我有一个简单的以下代码。内存中变量word_to_id的大小约为50MB。这会导致在将管道提交到Dataflow Runner时出错。

413请求实体太大

  word_to_id = {tok: idx for idx, tok in enumerate(vocab)}

  def extract_word_ids(tokens):
    return [word_to_id[w] for w in tokens if word_to_id.get(w, None)]

  with beam.pipeline.Pipeline(
    options=get_pipeline_option()) as p:
    lines = p | 'Read' >> beam.io.ReadFromText(path)

    word_ids = (
        lines
        | 'TokenizeLines' >> beam.Map(words)
        | 'IntergerizeTokens' >> beam.Map(extract_word_ids)
    )

请为此提供替代解决方案。

Answer 1

您可以将GCS存储桶用作文本和变量的来源，并将变量用作side input。您可以将此侧面输入用作列表，字典或单例。

这里有一个wordcount删除存储在GCS存储桶中的停用词的示例

with beam.Pipeline() as p:

    path = "gs://dataflow-samples/shakespeare/kinglear.txt"
    stopwords_path = "<BUCKET/stopwords>"

    output_path = "<BUCKET>"

    def split_words(text, stopwords):
            words = re.split('\W+', text)
            try:
                words.remove('')
            except:
                pass
            return [x for x in words if x.lower() not in stopwords]

    stopwords_p = (p | "Read Stop Words" >> ReadFromText(stopwords_path)
                     | FlatMap(lambda x: x.split(", "))) 

    text = p | "Read Text" >> ReadFromText(path)

    (text | "Split Words" >> FlatMap(split_words, stopwords=beam.pvalue.AsList(stopwords_p))
          | "Count" >> Count.PerElement()
          | "Write" >> WriteToText(file_path_prefix=output_path, file_name_suffix=".txt"))