PySpark代码优化：行对比与公共列值查询性能提升

需要对比PySpark DataFrame中Column1的所有不重复两两组合，判断每组组合在Column2中是否存在公共值，输出对应结果。

示例输入DataFrame：

Column1  Column2
abc      111
def      666
def      111
tyu      777
abc      777
def      222
tyu      333
ewq      888

期望输出：

abc,def,CommonRow  <-- because of 111
abc,ewq,NoCommonRow
abc,tyu,CommonRow  <-- because of 777
def,ewq,NoCommonRow
def,tyu,NoCommonRow
ewq,tyu,NoCommonRow

原始代码采用本地双重循环，对每组Column1组合都执行一次filter和join操作，每一次循环都会触发Spark作业。面对百万级数据时，会产生大量重复计算和作业调度开销，导致运行效率极低。

利用Spark的分布式聚合和集合操作，将所有计算逻辑转为分布式执行，仅需少数几次作业即可完成：

1. 预处理：构建Column1到Column2的唯一值映射
   先通过groupBy和collect_set，把每个Column1对应的所有唯一Column2值聚合为一个集合，减少后续计算量：

   from pyspark.sql import functions as F
   
   # 聚合每个Column1对应的Column2唯一值集合
   grouped_df = df.groupBy("Column1").agg(F.collect_set("Column2").alias("col2_set"))
   

2. 生成不重复的Column1两两组合
   通过自连接并设置a.Column1 < b.Column1的条件，生成所有无序且不重复的Column1组合，避免重复对比：

   # 自连接生成所有不重复的两两组合
   pair_df = grouped_df.alias("a").join(
       grouped_df.alias("b"),
       F.col("a.Column1") < F.col("b.Column1"),
       how="inner"
   )
   

3. 判断公共值并生成结果
   使用Spark内置的array_intersect函数计算两个集合的交集，根据交集是否为空判断结果：

   # 判断是否存在公共值，生成结果列
   result_df = pair_df.withColumn(
       "result",
       F.when(F.size(F.array_intersect(F.col("a.col2_set"), F.col("b.col2_set"))) > 0, "CommonRow")
       .otherwise("NoCommonRow")
   )
   
   # 格式化输出为期望的字符串格式
   final_df = result_df.select(
       F.concat_ws(",", F.col("a.Column1"), F.col("b.Column1"), F.col("result")).alias("output")
   )
   

4. 输出结果
   可以直接展示DataFrame，或者逐行打印：

   # 直接展示结果
   final_df.show(truncate=False)
   
   # 或者逐行打印成示例格式
   for row in final_df.collect():
       print(row["output"])
   

• 所有操作均为Spark分布式执行，仅触发3-4次作业，避免了本地循环的多次调度开销
• 聚合和集合操作的计算效率远高于多次filter+join，适合处理百万级甚至更大规模的数据
• 代码逻辑简洁，易于维护和扩展

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