如何优化PySpark逻辑 仅遍历一次DataFrame提取TBD项目名中国家值

原循环写法的性能瓶颈在于重复调用withColumn迭代修改DataFrame，相当于在执行计划中叠加75层列计算逻辑，即使Catalyst优化器做部分合并，仍会产生大量冗余计算。以下两个方案均可实现仅单次遍历DataFrame完成计算，大幅降低耗时：

方案1：正则一次性匹配（最优性能） #

通过构造包含所有国家名的正则表达式，仅需一次字符串扫描即可完成匹配提取，时间复杂度最低。

实现代码 #

import re
import pyspark.sql.functions as F

# 按国家名长度降序排序，避免短名称被优先匹配导致错误（如"United States"优先于"States"匹配）
sorted_countries = sorted(country_list, key=lambda x: len(x), reverse=True)
# 转义国家名中的正则特殊字符，避免匹配异常
escaped_countries = [re.escape(c) for c in sorted_countries]
# 构造正则规则：仅匹配包含TBD的项目名，捕获对应国家名
regex_pattern = rf".*TBD.*\b({'|'.join(escaped_countries)})\b.*"

df = df.withColumn(
    "ExtractColumn",
    F.when(
        F.regexp_extract(F.col("ProjectName"), regex_pattern, 1) != "",
        F.regexp_extract(F.col("ProjectName"), regex_pattern, 1)
    ).otherwise(None)
)

方案2：coalesce 合并判断逻辑（逻辑完全兼容原写法） #

如果不想使用正则，可以将所有判断逻辑合并为单次列计算，和原实现逻辑完全一致，性能远优于循环写法。

实现代码 #

import pyspark.sql.functions as F

# 批量生成所有国家的匹配判断条件
match_conditions = [
    F.when(
        F.col("ProjectName").contains("TBD") & F.col("ProjectName").contains(country),
        F.lit(country)
    ) for country in country_list
]

# coalesce会按顺序返回第一个非空的匹配结果，无匹配则返回null
df = df.withColumn("ExtractColumn", F.coalesce(*match_conditions))

性能说明 #

• 正则方案每行仅执行1次匹配操作，是性能最优的选择，适合超大规模数据集。
• coalesce方案每行最多执行75次字符串包含判断，性能略低于正则，但逻辑和原写法完全对齐，不需要处理正则转义、排序问题，适合对准确性要求极高的场景。

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