PySpark如何基于配置字典高效实现DataFrame条件映射？

针对200万行的PySpark DataFrame，循环过滤后union的方式会触发多次Shuffle和Stage提交，性能自然极差。推荐以下两种矢量化、低开销的实现方式：

方案一：使用when()+otherwise直接映射（适合规则数量较少的场景） #

这种方式无需额外Shuffle，直接在原DataFrame上进行列计算，性能最优。

假设你的配置字典示例如下：

mapping_rules = {
    1: {"condition": "sub_category_id IN (10, 11)", "new_sub": 100},
    2: {"condition": "sub_category_id = 20", "new_sub": 200},
    3: {"condition": "sub_category_id >= 30", "new_sub": 300}
}

实现代码：

from pyspark.sql import functions as F

# 初始化为原sub_category_id，后续覆盖匹配规则的行
transformed_df = inventory.withColumn("sub_category_id", F.col("sub_category_id"))

# 遍历规则，逐个应用when条件
for cat_id, rule in mapping_rules.items():
    transformed_df = transformed_df.withColumn(
        "sub_category_id",
        F.when(
            (F.col("category_id") == cat_id) & F.expr(rule["condition"]),
            F.lit(rule["new_sub"])
        ).otherwise(F.col("sub_category_id"))
    )

方案二：映射DataFrame关联（适合规则数量多、需灵活维护的场景） #

将配置规则转换成PySpark DataFrame，通过join关联原表实现映射，同样是批量矢量化操作，避免循环开销。

步骤1：将配置字典转换为结构化的映射DataFrame

# 展开配置规则为列表
rule_list = []
for cat_id, rule in mapping_rules.items():
    rule_list.append({
        "category_id": cat_id,
        "condition_expr": rule["condition"],
        "new_sub_category_id": rule["new_sub"]
    })

# 创建映射DataFrame
mapping_df = spark.createDataFrame(rule_list)

步骤2：关联原表并应用映射

# 关联原表与映射表，匹配category_id
joined_df = inventory.join(mapping_df, on="category_id", how="left")

# 计算新的sub_category_id：匹配到规则则用映射值，否则保留原字段
final_df = joined_df.withColumn(
    "sub_category_id",
    F.when(
        F.expr("condition_expr"),
        F.col("new_sub_category_id")
    ).otherwise(F.col("sub_category_id"))
).drop("condition_expr", "new_sub_category_id")

关键性能说明 #

• 两种方案均为矢量化操作，仅需1-2次Stage处理，避免了循环过滤+union带来的多次Shuffle和重复计算
• 若规则中包含复杂条件，F.expr()可直接复用SQL风格的条件表达式，无需额外转换
• 200万行属于PySpark的常规处理量级，以上两种方案均可在数秒内完成计算

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