PySpark如何基于字典正则部分匹配规则为DataFrame新增列

实现方案 #

你需要的模糊匹配无法通过create_map实现，该函数仅支持精确键值匹配，你可以通过链式构造when条件实现需求，全程使用Spark原生函数，性能远高于自定义UDF：

1. 首先修正你给出的映射字典语法错误，Y、Z需要加引号声明为字符串：

d = {
    "abc": "X",
    "some_other": "Y",
    "anything": "Z"
}

2. 构造匹配逻辑并新增C列：

from pyspark.sql import functions as f

# 初始化匹配表达式
match_expr = None
for pattern, mapped_val in d.items():
    # 前缀匹配规则，如需任意位置包含可替换为 f.col("B").rlike(pattern)
    cond = f.col("B").startswith(pattern)
    if match_expr is None:
        match_expr = f.when(cond, mapped_val)
    else:
        match_expr = match_expr.when(cond, mapped_val)

# 匹配不到的场景默认返回Null，可根据需求修改otherwise的返回值
match_expr = match_expr.otherwise(None)

# 新增C列
df = df.withColumn("C", match_expr)

补充说明 #

• 如果匹配规则存在优先级，把优先级高的规则放在字典前面即可，Python 3.7+版本默认保留字典插入顺序，会优先匹配先出现的规则
• 如果需要更复杂的正则匹配，把startswith替换为rlike方法，传入对应的正则表达式即可正常匹配

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