如何在PySpark中实现自定义聚合函数(UDAF)以复用现有逻辑？

你的需求是对分组后的DataFrame执行自定义复杂处理，PySpark中没有直接对应Scala UDAF的原生Python类，但可以通过以下两种方式实现类似效果：

方案一：使用GroupedData.apply()（完全贴合你的函数逻辑） #

这个方法直接支持对每个分组的完整DataFrame进行处理，完美匹配你process_data接收DataFrame的设计：

1. 定义返回结果的Schema
   因为Spark需要明确知道分组处理后的返回结构，先定义对应schema：

   from pyspark.sql.types import StructType, StructField, StringType, BinaryType
   
   result_schema = StructType([
       StructField("Foo_ID", StringType(), nullable=False),
       StructField("processed_result", BinaryType(), nullable=True)
   ])
   

2. 封装分组处理逻辑
   把你的process_data包装成适配apply的函数，需要返回可迭代的结果（元组或Row）：

   def process_group(group_df):
       # 获取当前分组的Foo_ID值
       foo_id = group_df.select("Foo_ID").first()[0]
       # 调用你的自定义处理函数
       processed_bytes = process_data(group_df)
       return [(foo_id, processed_bytes)]
   

3. 执行分组聚合
   直接调用groupBy后的apply方法：

   result_df = source_df.groupBy("Foo_ID").apply(process_group, schema=result_schema)
   

方案二：使用Pandas UDAF（适合基于单列/多列的聚合） #

如果你的process_data可以适配成基于Pandas对象的处理（不需要完整DataFrame结构），可以用性能更优的Pandas聚合UDF：

1. 改造处理函数为Pandas UDAF
   用pandas_udf装饰器标记为聚合类型，适配Pandas Series输入：

   import pandas as pd
   from pyspark.sql.functions import pandas_udf
   
   @pandas_udf(BinaryType(), functionType=pandas_udf.Type.AGGREGATE)
   def pandas_process_data(col: pd.Series) -> bytes:
       # 将Series转为DataFrame，适配原process_data的入参要求
       temp_df = col.to_frame("target_column")
       return process_data(temp_df)
   

2. 执行聚合操作
   直接在agg中调用该函数：

   result_df = source_df.groupBy("Foo_ID").agg(
       pandas_process_data("your_target_column").alias("processed_result")
   )
   

注意事项 #

• 方案一的apply支持任意复杂的DataFrame处理，但要确保process_data函数可序列化（能被分发到Spark Executor执行）。
• 方案二的Pandas UDAF基于Apache Arrow传输数据，性能比普通Python UDF更高，但仅支持基于列的聚合输入。
• 如果对性能要求极高，可考虑用Scala编写原生UDAF，再在Python代码中调用。

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