Apache Beam PTransform输出传递与复杂数据流构建实践

AiMusic.so

aimusic.so是一款全新的在线免费AI音乐生成器

47

查看详情

import apache_beam as beam import requests # 实际API调用可能用到  # 1. 模拟从数据库读取数据的PTransform class ReadFromDatabase(beam.PTransform):     def expand(self, pcoll):         # 在实际应用中，这里会使用 beam.io.ReadFromJdbc 或其他数据库连接器         # 模拟读取两行数据，每行是一个字典         print("Executing ReadFromDatabase...")         return pcoll | 'ReadDatabaseEntries' >> beam.Create([             {'id': 1, 'name': 'Alice', 'email': 'alice@example.com'},             {'id': 2, 'name': 'Bob', 'email': 'bob@example.com'}         ])  # 2. 调用第一个REST API的PTransform class CallFirstAPI(beam.PTransform):     class ProcessElement(beam.DoFn):         def process(self, element):             # 模拟调用第一个API，并获取一个包含数组的响应             # 实际中会使用 requests.get(f"http://api.example.com/first/{element['id']}")             print(f"CallFirstAPI - Processing element: {element['id']}")             api_response = {                 'status': 'success',                 'data': {                     'id': element['id'],                     'details': f"details_for_{element['name']}",                     'items': [f"itemA_{element['id']}", f"itemB_{element['id']}"] # 模拟数组                 }             }             # 将原始数据与API响应合并，并传递给下一步             yield {**element, 'first_api_data': api_response['data']}      def expand(self, pcoll):         return pcoll | 'CallFirstAPIProcess' >> beam.ParDo(self.ProcessElement())  # 3. 调用第二个REST API的PTransform (针对数组中的每个元素) class CallSecondAPI(beam.PTransform):     class ProcessElement(beam.DoFn):         def process(self, element):             first_api_data = element['first_api_data']             items = first_api_data.get('items', [])              # 对第一个API响应中的每个item调用第二个API             for item in items:                 # 模拟调用第二个API                 # 实际中会使用 requests.get(f"http://api.example.com/second/{item}")                 print(f"CallSecondAPI - Processing item: {item} for element: {element['id']}")                 second_api_response = {                     'item_name': item,                     'additional_info': f"info_for_{item}"                 }                 # 将原始数据、第一个API数据和当前第二个API响应合并                 # 注意：这里可能会产生多个输出元素，每个对应一个item                 yield {                     **element,                     'current_item_data': second_api_response                 }      def expand(self, pcoll):         # 使用ParDo处理每个元素，并可能产生多个输出         return pcoll | 'CallSecondAPIProcess' >> beam.ParDo(self.ProcessElement())  # 4. 模拟更新数据库的PTransform (仅作示意) class UpdateDatabase(beam.PTransform):     class ProcessElement(beam.DoFn):         def process(self, element):             # 在实际应用中，这里会使用 beam.io.WriteToJdbc 或其他数据库写入器             # 可能需要根据element中的id和API数据构建SQL UPDATE语句             print(f"UpdateDatabase - Updating record: {element['id']} with data: {element}")             # 实际场景中，此DoFn可能不会yield任何元素，或者yield一个更新成功标记             yield element # 仅仅为了在管道末尾查看数据流      def expand(self, pcoll):         return pcoll | 'UpdateDatabaseEntries' >> beam.ParDo(self.ProcessElement())  # 构建并运行Beam管道 with beam.Pipeline() as pipeline:     # 步骤1: 从数据库读取数据     read_from_db_pcoll = pipeline | 'Start' >> ReadFromDatabase()      # 步骤2: 调用第一个API，其输入是 read_from_db_pcoll 的输出     call_first_api_pcoll = read_from_db_pcoll | 'CallFirstAPI' >> CallFirstAPI()      # 步骤3: 调用第二个API，其输入是 call_first_api_pcoll 的输出     # 注意：CallSecondAPI可能会将一个输入元素扩展为多个输出元素     call_second_api_pcoll = call_first_api_pcoll | 'CallSecondAPI' >> CallSecondAPI()      # 步骤4: 更新数据库，其输入是 call_second_api_pcoll 的输出     # 在实际场景中，可能需要进一步聚合或处理 call_second_api_pcoll 的输出     # 例如，如果需要将多个item的结果聚合回原始记录，可能需要使用GroupByKey     call_second_api_pcoll | 'UpdateDatabase' >> UpdateDatabase()      # 如果需要查看最终结果，可以写入文件或打印     # call_second_api_pcoll | 'WriteToConsole' >> beam.Map(print)