如何基于集合过滤PySpark DataFrame指定行并按规则合并两类数据

原代码错误原因 #

• 直接遍历PySpark DataFrame对象无法获取行数据，你没有调用collect()方法将数据拉取到Driver端，所以循环收集逻辑完全没有执行，sf_DATA_set是空集合，isin过滤自然无效。
• 过滤df2的逻辑遗漏了规则例外：没有判断DATA值为SP时直接保留的条件。

修正后的实现方案 #

方案1：修正原有集合思路 #

适合DATA数据量不大、不会导致Driver端OOM的场景：

import pyspark.sql.functions as F    

# 拆分得到df1和df2
df1 = final_df.filter(F.col("FLAG") == 'S').select("FLAG", "DATA")
df2 = final_df.filter(F.col("FLAG") == 'N').select("FLAG", "DATA")

# 收集df1中非SP的DATA值，必须调用collect()拉取到本地
s_non_sp_list = [row["DATA"] for row in df1.filter(F.col("DATA") != 'SP').select("DATA").collect()]

# 过滤df2：满足SP 或者 不在S的非SP列表中两个条件之一即可保留
df2_filtered = df2.filter(
    (F.col("DATA") == 'SP') | (~F.col("DATA").isin(s_non_sp_list))
)

# 合并结果，df1在前符合你要求的顺序
final_result = df1.unionByName(df2_filtered)
final_result.show()

方案2：分布式JOIN实现（推荐） #

完全基于Spark分布式计算逻辑，不需要把数据拉取到Driver端，适合大数据量场景，性能更稳定：

import pyspark.sql.functions as F

df1 = final_df.filter(F.col("FLAG") == 'S').select("FLAG", "DATA")
df2 = final_df.filter(F.col("FLAG") == 'N').select("FLAG", "DATA")

# 给df1中所有非SP的DATA打存在标记
s_exists_flag = df1.filter(F.col("DATA") != 'SP') \
    .select("DATA") \
    .distinct() \
    .withColumn("is_exist_in_s", F.lit(True))

# df2左联标记表，过滤符合要求的行
df2_filtered = df2.join(s_exists_flag, on="DATA", how="left") \
    .filter(
        (F.col("DATA") == 'SP') | 
        (F.col("is_exist_in_s").isNull())
    ) \
    .select("FLAG", "DATA")

# 合并得到最终结果
final_result = df1.unionByName(df2_filtered)
final_result.show()

两种方案运行后都会得到你预期的输出结果。

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