Apache Beam调用Google自然语言API遇序列化错误求助

问题 #

在DoFn类中调用Google Natural Language API时，触发序列化错误：

TypeError: cannot pickle 'google.protobuf.pyext._message.RepeatedCompositeContainer' object [while running 'ParDo(RequestAPI)']

使用的代码如下：

DoFn实现代码 #

class RequestAPI(beam.DoFn):
    def setup(self):
        self._client = language_v1.LanguageServiceClient()

    def process(self, element):
        classify = False
        if element['lang'] =='en':
            classify = True
        document = {"document": {
            "type_": "PLAIN_TEXT", "content": element['text']},
            "features": {
                "extract_entities": True,
                "extract_document_sentiment": True,
                "extract_entity_sentiment": True,
                "classify_text": classify,
            },
            "encoding_type": "UTF8"
        }

        response = self._client.annotate_text(request=document)
        new_element = element
        new_element['analysis'] = {'sentences': response.sentences,
                                   'score': response.document_sentiment.score,
                                   'magnitude': response.document_sentiment.magnitude,
                                   'tokens': response.tokens,
                                   'lang_api': response.language,
                                   'entities': response.entities
                                   }
        print(new_element)
        return new_element

调用代码 #

data_enrich: PCollection = (
            tweets_data | beam.ParDo(RequestAPI())
            | 'write data' >> beam.io.WriteToText('data.txt')
    )

环境：Python 3.7、Apache Beam 2.43，此前使用Map transform也出现相同错误。

解决方案 #

错误核心原因是：Google Natural Language API返回的response.sentences、response.tokens、response.entities等对象属于Protobuf的RepeatedCompositeContainer类型，这类对象无法被Python的pickle机制序列化，而Apache Beam在处理PCollection元素时必须序列化数据以支持分布式传输或持久化，因此直接将这些对象存入字典会触发错误。

解决方法是将Protobuf对象转换为Python原生可序列化类型（如dict、list），具体修改如下：

1. 导入Protobuf转字典的工具：

from google.protobuf.json_format import MessageToDict

2. 修改process方法中处理响应的逻辑：

def process(self, element):
    # 原有document构建逻辑保持不变
    classify = False
    if element['lang'] =='en':
        classify = True
    document = {"document": {
        "type_": "PLAIN_TEXT", "content": element['text']},
        "features": {
            "extract_entities": True,
            "extract_document_sentiment": True,
            "extract_entity_sentiment": True,
            "classify_text": classify,
        },
        "encoding_type": "UTF8"
    }

    response = self._client.annotate_text(request=document)
    new_element = element.copy()  # 创建原元素副本，避免修改原数据引发意外
    # 将Protobuf对象转换为可序列化的Python类型
    analysis_data = {
        'sentences': [MessageToDict(sent) for sent in response.sentences],
        'score': response.document_sentiment.score,
        'magnitude': response.document_sentiment.magnitude,
        'tokens': [MessageToDict(token) for token in response.tokens],
        'lang_api': response.language,
        'entities': [MessageToDict(entity) for entity in response.entities]
    }
    new_element['analysis'] = analysis_data
    print(new_element)
    return [new_element]  # DoFn的process方法必须返回可迭代对象，这里用列表包裹单个元素

关键注意点 #

• MessageToDict负责将单个Protobuf Message对象转为Python字典，对于集合类型的Protobuf对象（如response.sentences），需要用列表推导式逐个转换。
• 原代码中new_element = element是引用赋值，会直接修改原元素，建议用element.copy()创建副本，避免影响后续数据处理。
• DoFn的process方法要求返回可迭代对象（如列表、生成器），原代码直接返回单个字典可能导致Beam处理异常，因此改为返回[new_element]。

内容的提问来源于stack exchange，提问作者jose manuel rojas granado