PySpark：能否通过cogroup+applyInPandas处理2个以上DataFrame？

由于PySpark 3.2.0的cogroup仅能返回支持两两处理的PandasCogroupedOps对象，无法链式添加更多DataFrame，可通过以下两种方式实现多DataFrame的分组处理：

方法一：合并多DataFrame后分组处理 #

先给每个DataFrame添加来源标识列，合并为单个DataFrame后再分组，在applyInPandas的处理函数中拆分出各个原始DataFrame进行逻辑处理：

from pyspark.sql.functions import lit

# 给每个DataFrame添加来源标识列
df1_with_src = df1.withColumn("source", lit("df1"))
df2_with_src = df2.withColumn("source", lit("df2"))
df3_with_src = df3.withColumn("source", lit("df3"))
df4_with_src = df4.withColumn("source", lit("df4"))

# 合并所有DataFrame（确保各DataFrame列名一致）
combined_df = df1_with_src.unionByName(df2_with_src) \
                          .unionByName(df3_with_src) \
                          .unionByName(df4_with_src)

def process_multiple_pdfs(pdf):
    # 按标识列拆分出各个原始DataFrame
    pdf1 = pdf[pdf["source"] == "df1"].drop("source", axis=1)
    pdf2 = pdf[pdf["source"] == "df2"].drop("source", axis=1)
    pdf3 = pdf[pdf["source"] == "df3"].drop("source", axis=1)
    pdf4 = pdf[pdf["source"] == "df4"].drop("source", axis=1)
    
    # 执行你的业务逻辑，比如关联、聚合等操作
    # result_pdf = <处理逻辑>
    return result_pdf

# 分组并调用applyInPandas
final_result = combined_df.groupby("id").applyInPandas(
    process_multiple_pdfs, 
    schema="time int, id int, v1 double, v2 string"
)

方法二：嵌套cogroup逐步合并处理 #

通过多次两两cogroup逐步合并结果，每次处理两个DataFrame，最终完成多DataFrame的逻辑处理：

# 第一步：处理df1和df2
def process_df1_df2(pdf1, pdf2):
    # 合并前两个DataFrame的结果，返回包含id及所需列的DataFrame
    combined = pd.merge(pdf1, pdf2, on="id", how="outer")
    return combined

temp_result1 = df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas(
    process_df1_df2, 
    schema="id int, <列名及类型，需匹配合并后的结构>"
)

# 第二步：用第一步结果处理df3
def process_temp_df3(pdf_temp, pdf3):
    combined = pd.merge(pdf_temp, pdf3, on="id", how="outer")
    return combined

temp_result2 = temp_result1.groupby("id").cogroup(df3.groupby("id")).applyInPandas(
    process_temp_df3, 
    schema="id int, <更新后的列结构>"
)

# 第三步：用第二步结果处理df4
def process_final(pdf_temp, pdf4):
    # 执行最终业务逻辑
    # result_pdf = <处理逻辑>
    return result_pdf

final_result = temp_result2.groupby("id").cogroup(df4.groupby("id")).applyInPandas(
    process_final, 
    schema="time int, id int, v1 double, v2 string"
)

注意事项 #

• 方法一适合各DataFrame结构相近的场景，需确保合并时列名一致；
• 方法二需注意每次applyInPandas的schema必须匹配当前处理后返回的DataFrame结构；
• 两种方法中，处理函数内的pandas逻辑需保证分组后的每个id对应的子DataFrame处理正确。

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