GCP Dataflow批处理：传入Pandas Dataframe后process_batch未执行的原因及示例

1. ParDo默认处理单个元素：你直接将整个Pandas DataFrame作为输入传给ParDo，Beam会把它当作单个元素处理，而process_batch方法是用于批量元素处理的，只有当输入是分块的批量数据时才会触发。
2. 未启用批量处理模式：process_batch属于Batch DoFn的方法，需要明确配置批量处理的触发条件（比如指定批次大小），否则Beam会默认调用普通的process方法（你未实现该方法，因此无输出）。

下面提供两种适配Pandas DataFrame的批处理实现方式：

方式一：使用Batch DoFn处理分块数据 #

先将DataFrame的行转为PCollection元素，再通过BatchElements分块，最后用实现process_batch的DoFn处理：

import apache_beam as beam
import pandas as pd
from typing import Any

class TestBatch(beam.DoFn):
    def __init__(self, args: Any):
        self.args = args
        print("TestBatch.init")
    
    def setup(self) -> None:
        print("TestBatch.setup")
    
    def process_batch(self, batch: list[dict]) -> pd.DataFrame:
        # 将批量行转为DataFrame
        df_batch = pd.DataFrame(batch)
        print("TestBatch.process_batch")
        print(df_batch)
        
        # 执行列转换示例：新增一列，值为原列的2倍
        df_batch['new_col'] = df_batch['original_col'] * 2
        
        yield df_batch

if __name__ == "__main__":
    # 读取Excel数据
    pd_df = pd.read_excel("your_path.xlsx")
    
    # 初始化Pipeline
    with beam.Pipeline() as p:
        # 将DataFrame转为行的PCollection
        rows = p | "DataFrame to Rows" >> beam.Create(pd_df.to_dict('records'))
        
        # 将元素分批次（比如每100行一个批次）
        batched_rows = rows | "Batch Elements" >> beam.BatchElements(min_batch_size=50, max_batch_size=100)
        
        # 用Batch DoFn处理批次
        result = batched_rows | "Process Batch" >> beam.ParDo(TestBatch(argv=None))
        
        # 输出结果（可选）
        result | "Print Result" >> beam.Map(lambda df: print(df))

方式二：直接使用Beam的Pandas转换API #

如果你的需求是对整个DataFrame做列转换，也可以用beam.Map直接处理DataFrame（适合单批次处理，或结合Partition拆分后处理）：

import apache_beam as beam
import pandas as pd

def transform_df(df: pd.DataFrame) -> pd.DataFrame:
    # 执行列转换逻辑
    df['transformed_col'] = df['original_col'].apply(lambda x: x.upper() if isinstance(x, str) else x)
    return df

if __name__ == "__main__":
    pd_df = pd.read_excel("your_path.xlsx")
    
    with beam.Pipeline() as p:
        # 将DataFrame作为单个元素传入，用Map处理
        result = (
            p | "Create DataFrame PCollection" >> beam.Create([pd_df])
              | "Transform DataFrame" >> beam.Map(transform_df)
              | "Print Result" >> beam.Map(lambda df: print(df))
        )

• 如果使用process_batch，必须确保输入是批量的元素集合，而非单个DataFrame对象。
• 若要处理整个DataFrame，直接用beam.Create([pd_df])将其包装为PCollection的单个元素，再用beam.Map处理更直接。
• 在Dataflow运行时，打印语句的输出会出现在Worker日志中；本地运行时会直接显示在控制台。

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