Pyspark如何将DataFrame作为数组类型列关联到另一个DataFrame

实现代码 #

from pyspark.sql import Row
from pyspark.sql.functions import collect_list, struct

# 构造示例数据
df1 = spark.createDataFrame([
    Row(a = 1, b = 'C', c = 26, d = 'abc'),
    Row(a = 1, b = 'C', c = 27, d = 'def'),
    Row(a = 1, b = 'D', c = 51, d = 'ghi'),
    Row(a = 2, b = 'C', c = 40, d = 'abc'),
    Row(a = 2, b = 'D', c = 45, d = 'abc'),
    Row(a = 2, b = 'D', c = 38, d = 'def')
])

df2 = spark.createDataFrame([
    Row(a = 1, b = 'C', e = 2, f = 'cba'),
    Row(a = 1, b = 'D', e = 3, f = 'ihg'),
    Row(a = 2, b = 'C', e = 7, f = 'cba'),
    Row(a = 2, b = 'D', e = 9, f = 'cba')
])

# 第一步：对df1按a、b分组，将同组的c、d打包为结构体后收集为数组
df1_agg = df1.groupBy("a", "b") \
    .agg(collect_list(struct("c", "d")).alias("df1"))

# 第二步：将聚合后的df1和df2按a、b关联
result_df = df2.join(df1_agg, on=["a", "b"], how="left")

# 验证输出，转为JSON查看格式
result_df.toJSON().foreach(print)

逻辑说明 #

• 使用struct("c", "d")将每行的c、d字段打包为结构体，会自动保留原有字段名，不会丢失元信息
• 使用collect_list将同一个a、b分组下的所有结构体收集为数组，对应你需要的df1数组列
• 关联时选择left连接可以保留df2的所有行，即使某行a、b在df1中没有匹配记录，对应df1列会生成空数组，如果你需要只保留两边都匹配的行可以改为inner连接
• 输出的JSON格式完全符合你给出的示例要求

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