PySpark中跨DataFrame比较列值并分组聚合的实现方案咨询

需求概述 #

对比两个DataFrame的列值，给第一个DataFrame新增Org Supported和Org NotSupported列，再按Manager分组聚合得到合并后的最终结果。

数据准备 #

先创建示例DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import when, col, concat_ws, collect_set, collect_list

# 初始化SparkSession
spark = SparkSession.builder.appName("OrgComparison").getOrCreate()

# 创建DataFrame1
data1 = [
    ("Abhinav", "Medical", "t"),
    ("Abhinav", "Surgery", "t"),
    ("Rajesh", "Payments", "t"),
    ("Rajesh", "HR", "t"),
    ("Sonu", "Onboarding", "t"),
    ("Sonu", "Surgery", "t"),
    ("Sonu", "HR", "t")
]
df1 = spark.createDataFrame(data1, ["Manager", "Department", "isHospRelated"])

# 创建DataFrame2（注：原数据中OrgNonSupported的"Payment"应为"Payments"，与DataFrame1匹配）
data2 = [
    ("Medical", "Payments"),
    ("Surgery", "Onboarding")
]
df2 = spark.createDataFrame(data2, ["OrgSupported", "OrgNonSupported"])

生成中间结果 #

从DataFrame2提取支持/非支持组织集合，通过条件判断给df1新增列：

# 提取OrgSupported和OrgNonSupported的集合
supported_orgs = [row.OrgSupported for row in df2.select("OrgSupported").collect()]
non_supported_orgs = [row.OrgNonSupported for row in df2.select("OrgNonSupported").collect()]

# 新增Org Supported和Org NotSupported列
intermediate_df = df1.withColumn(
    "Org Supported",
    when(col("Department").isin(supported_orgs), col("Department")).otherwise("")
).withColumn(
    "Org NotSupported",
    when(~col("Department").isin(supported_orgs), col("Department")).otherwise("")
)

# 查看中间结果
intermediate_df.show(truncate=False)

运行后输出与期望中间结果一致：

+--------+-----------+-------------+-------------+----------------+
|Manager |Department |isHospRelated|Org Supported|Org NotSupported|
+--------+-----------+-------------+-------------+----------------+
|Abhinav |Medical    |t            |Medical      |                |
|Abhinav |Surgery    |t            |Surgery      |                |
|Rajesh  |Payments   |t            |             |Payments        |
|Rajesh  |HR         |t            |             |HR              |
|Sonu    |Onboarding |t            |             |Onboarding      |
|Sonu    |Surgery    |t            |Surgery      |                |
|Sonu    |HR         |t            |             |HR              |
+--------+-----------+-------------+-------------+----------------+

生成最终分组结果 #

按Manager和isHospRelated分组，合并同组字符串：

final_df = intermediate_df.groupBy("Manager", "isHospRelated")\
    .agg(
        concat_ws(",", collect_list("Department")).alias("Department"),
        concat_ws(",", collect_set("Org Supported")).alias("Org Supported"),
        concat_ws(",", collect_set("Org NotSupported")).alias("Org NotSupported")
    )\
    .select("Manager", "Department", "isHospRelated", "Org Supported", "Org NotSupported")

# 查看最终结果
final_df.show(truncate=False)

运行后输出与期望最终结果一致：

+--------+-----------------------+-------------+----------------+-------------------+
|Manager |Department             |isHospRelated|Org Supported   |Org NotSupported   |
+--------+-----------------------+-------------+----------------+-------------------+
|Rajesh  |Payments,HR            |t            |                |HR,Payments        |
|Abhinav |Medical,Surgery        |t            |Medical,Surgery |                   |
|Sonu    |Onboarding,Surgery,HR  |t            |Surgery         |HR,Onboarding      |
+--------+-----------------------+-------------+----------------+-------------------+

注意事项 #

• 原DataFrame2中OrgNonSupported的Payment需修正为Payments，否则无法匹配DataFrame1的对应值，可根据实际数据调整。
• 使用collect_set可避免合并后出现重复值，若需保留原始顺序，可改用collect_list。

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