优化Databricks环境下PySpark的withColumn when otherwise性能问题

问题根源 #

你遇到的性能问题完全是冗余代码导致的：连续24次调用withColumn嵌套when判断，会让Spark的逻辑执行计划指数级膨胀，Driver节点解析、优化超复杂执行计划时内存被占满引发GC，最终出现无响应、任务卡住的问题。

优化方案 #

方案1：字典映射替换（性能最优，适合固定映射场景） #

把所有日期对应关系做成字典，用PySpark内置的create_map函数一次性完成替换，全程使用Spark原生算子，无额外开销：

from pyspark.sql.functions import create_map, lit, col, coalesce

# 定义法语月份和数字的映射
fr_month_map = {
    "Janvier": "01", "Fevrier": "02", "Mars": "03", "Avril": "04", "Mai": "05", "Juin": "06",
    "Juillet": "07", "Aout": "08", "Septembre": "09", "Octobre": "10", "Novembre": "11", "Decembre": "12"
}
# 生成完整的替换映射字典
replace_map = {}
for year in ["2020", "2021"]:
    for month_name, month_num in fr_month_map.items():
        old_val = f"XXX {month_name} {year}"
        new_val = f"XXX{month_num}{year[-2:]}"
        replace_map[old_val] = new_val

# 生成map表达式一次性替换
map_expr = create_map([lit(x) for pair in replace_map.items() for x in pair])
# coalesce保证没有匹配到的值保持原样
courriers = courriers.withColumn("Vague", coalesce(map_expr[col("Vague")], col("Vague")))

方案2：日期解析转换（扩展性最好，适合后续新增年份/月份的场景） #

直接对字符串做拆分、日期解析，不需要手动维护映射表，后续新增年份无需修改代码：

from pyspark.sql.functions import regexp_replace, to_date, date_format, concat, lit, coalesce

# 若Spark默认locale不是法语，需先执行配置
# spark.conf.set("spark.sql.locale", "fr_FR")

courriers = courriers.withColumn(
    "Vague",
    coalesce(
        concat(
            lit("XXX"),
            date_format(
                to_date(
                    # 去掉前缀XXX后解析法语日期
                    regexp_replace(col("Vague"), "^XXX ", ""),
                    "MMMM yyyy"
                ),
                "MMyy"
            )
        ),
        col("Vague")
    )
)

额外优化建议 #

• 大文件读取后可先执行courriers = courriers.cache()缓存到内存，转换完成后调用unpersist()释放内存，避免重复读取CSV
• 展示DataFrame时务必加上limit(N)仅取前N行，不要直接调用show()或collect()拉取全量数据到Driver端

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