如何在Apache Beam Python中合并ParDo生成的多个DataFrame/PCollections

首先纠正你代码里的几个明显错误：

• process方法里要引用self.no_of_hits，而不是直接用no_of_hits
• 请求API后要调用.json()获取响应数据，不能直接把Response对象传给pd.json_normalize
• URL拼接逻辑有误，循环里应该用当前迭代的i来构造请求URL，而不是固定的no_of_hits

修正后的GetData类如下：

import requests
import pandas as pd
import apache_beam as beam

class GetData(beam.DoFn):
    def __init__(self, hits):
        self.no_of_hits = hits

    def process(self, url):
        for i in range(self.no_of_hits):
            # 根据你的API格式调整URL拼接方式，这里用page参数示例
            response = requests.get(f"{url}?page={i}")
            response.raise_for_status()  # 抛出请求异常，方便排查错误
            json_data = response.json()
            df = pd.json_normalize(json_data)
            yield df

接下来针对合并DataFrame的需求，提供三种实用方案：

方案一：拆分行再聚合（适合大数据量场景） #

这种方式先把每个DataFrame拆成单行字典，在分布式环境下并行处理，最后再合并成完整DataFrame，避免单节点内存压力。

def df_to_rows(df):
    # 将DataFrame转换成字典格式的行，用yield from输出所有行
    yield from df.to_dict('records')

def combine_rows_to_df(rows):
    # 将所有行字典合并成一个DataFrame
    return pd.DataFrame(rows)

def write_df_to_csv(df, output_path):
    df.to_csv(output_path, index=False)

with beam.Pipeline() as pipeline:
    data = (pipeline
            | "Create URL input" >> beam.Create(["https://your-api-url.com"])
            | "Fetch data from API" >> beam.ParDo(GetData(hits=5)))
    
    # 对每个DataFrame做基础处理，示例：保留指定列
    wrangled = (data
                | "Wrangle each DataFrame" >> beam.Map(lambda df: df[['id', 'name', 'value']]))
    
    # 转换为行集合
    rows = wrangled | "Convert DataFrame to rows" >> beam.FlatMap(df_to_rows)
    
    # 合并所有行成单个DataFrame
    combined_df = rows | "Combine rows to DataFrame" >> beam.CombineGlobally(combine_rows_to_df)
    
    # 写入CSV文件
    _ = combined_df | "Write to CSV" >> beam.Map(lambda df: write_df_to_csv(df, "output.csv"))

方案二：直接合并DataFrame（适合小数据量场景） #

如果数据量不大，可以自定义CombineFn直接合并所有DataFrame，逻辑更直观：

class MergeDataFrames(beam.CombineFn):
    def create_accumulator(self):
        # 初始化累加器，用列表存储DataFrame
        return []
    
    def add_input(self, accumulator, df):
        # 将单个DataFrame加入累加器
        accumulator.append(df)
        return accumulator
    
    def merge_accumulators(self, accumulators):
        # 合并多个累加器中的DataFrame列表
        all_dfs = []
        for acc in accumulators:
            all_dfs.extend(acc)
        # 先局部合并，减少后续全局合并的压力
        return [pd.concat(all_dfs, ignore_index=True)]
    
    def extract_output(self, accumulator):
        # 最终合并所有DataFrame
        return pd.concat(accumulator, ignore_index=True) if accumulator else pd.DataFrame()

with beam.Pipeline() as pipeline:
    data = (pipeline
            | "Create URL input" >> beam.Create(["https://your-api-url.com"])
            | "Fetch data from API" >> beam.ParDo(GetData(hits=5)))
    
    wrangled = (data
                | "Wrangle each DataFrame" >> beam.Map(lambda df: df.drop(columns=['unwanted_col'])))
    
    # 合并所有DataFrame
    combined_df = wrangled | "Merge DataFrames" >> beam.CombineGlobally(MergeDataFrames())
    
    # 写入CSV
    _ = combined_df | "Write to CSV" >> beam.Map(lambda df: df.to_csv("output.csv", index=False))

方案三：使用Beam DataFrame API（推荐） #

Beam提供了延迟DataFrame API，可以像操作普通Pandas DataFrame一样处理分布式数据，自动处理合并逻辑：

from apache_beam.dataframe.convert import to_dataframe

with beam.Pipeline() as pipeline:
    data = (pipeline
            | "Create URL input" >> beam.Create(["https://your-api-url.com"])
            | "Fetch data from API" >> beam.ParDo(GetData(hits=5)))
    
    # 确保所有DataFrame的Schema一致（列名、类型相同）
    wrangled = (data
                | "Wrangle DataFrames" >> beam.Map(lambda df: df.rename(columns={'old_name': 'new_name'})))
    
    # 转换为Beam延迟DataFrame
    beam_df = wrangled | "Convert to Beam DataFrame" >> to_dataframe()
    
    # 直接写入CSV，Beam自动处理分布式合并
    beam_df.to_csv("output.csv", index=False).run()

注意事项 #

• 如果数据量极大，CombineGlobally会将所有数据集中到单个worker，可能导致内存溢出，此时优先选择方案一或方案三
• 确保所有DataFrame的Schema一致，否则合并时会出现列不匹配的错误
• 生产环境中建议添加错误处理（比如API请求超时、重试逻辑）

内容的提问来源于stack exchange，提问作者Jack Froster