如何无循环高效按字典条件过滤PySpark DataFrame？

问题背景 #

现有PySpark DataFrame，需要根据product_pt_family对应的阈值筛选出acceptance_rate不小于阈值的记录，阈值存储在如下字典中：

dict1 = {'Fruits & Vegetables': 85, 'Dairy & Eggs': 90, 'Water':91, 'Bakery':92}

当前采用循环+union的方式实现，但运行耗时过长，原代码如下：

pt_families = list(set(df.select('product_pt_family').toPandas()['product_pt_family']))

schema = df.schema

thresholded_df = spark.createDataFrame([], schema)

for pt_family in pt_families:
        df_1 = df.filter((df.product_pt_family == pt_family ) & (df.acceptance_rate >= dict1[pt_family]))
        thresholded_df = thresholded_df.union(df_1)
     
thresholded_df.show()

解决方案1：用create_map构建阈值映射（推荐） #

利用PySpark的create_map函数把字典转换成列映射，生成对应阈值列后直接过滤，全程只扫描一次原DataFrame：

from pyspark.sql import functions as F

# 将字典键值对转换为map表达式
threshold_map = F.create_map([F.lit(x) for pair in dict1.items() for x in pair])

# 添加临时阈值列，过滤后可按需删除
thresholded_df = df.withColumn("threshold", threshold_map[F.col("product_pt_family")]) \
                   .filter(F.col("acceptance_rate") >= F.col("threshold")) \
                   .drop("threshold")

thresholded_df.show()

核心优势 #

• 仅一次全表扫描，彻底避免循环中重复读取原数据的问题
• 无需创建空DataFrame和多次执行union，减少Spark任务调度和数据合并的开销
• 代码简洁，逻辑直观易懂

解决方案2：字典转临时表关联过滤 #

如果阈值规则需要独立维护或规模较大，可以把字典转换成临时DataFrame，通过关联操作实现过滤：

# 将字典转换为PySpark DataFrame
threshold_df = spark.createDataFrame(dict1.items(), ["product_pt_family", "threshold"])

# 关联原表并过滤
thresholded_df = df.join(threshold_df, on="product_pt_family", how="inner") \
                   .filter(F.col("acceptance_rate") >= F.col("threshold")) \
                   .drop("threshold")

thresholded_df.show()

核心优势 #

• 阈值规则可单独管理，适合需要频繁调整阈值的场景
• 关联操作由Spark优化器自动优化，性能稳定可靠

原方法低效原因 #

循环+union的方式存在两个关键性能瓶颈：

1. 重复扫描数据：每次循环都会对原DataFrame执行一次过滤操作，相当于多次读取相同数据
2. 多次union开销：每次union都会生成新的DataFrame，Spark需要处理多次数据合并，增加任务执行的额外成本

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