PySpark中如何以可扩展方式链式串联多个WHEN条件

问题场景 #

我有一个变量名为pats的字典，其中存储了大量when函数的入参：正则匹配条件与对应返回值。
原始硬编码实现逐行编写每个when判断分支，可维护性差，需要找到可扩展的方式链式串联所有when条件，无需为字典中的每个键单独编写一行判断逻辑。

原始实现代码如下：

from pyspark.sql import functions as F
df = spark.createDataFrame([("ė",), ("2",), ("",), ("@",)], ["col1"])

pats = {
  r"^\d$"          : "digit",
  r"^\p{L}$"       : "letter",
  r"^[\p{P}\p{S}]$": "spec_char",
  r"^$"            : "empty"
}

whens = (
    F.when(F.col("col1").rlike(list(pats)[0]), pats[list(pats)[0]])
     .when(F.col("col1").rlike(list(pats)[1]), pats[list(pats)[1]])
     .when(F.col("col1").rlike(list(pats)[2]), pats[list(pats)[2]])
     .when(F.col("col1").rlike(list(pats)[3]), pats[list(pats)[3]])
     .otherwise(F.col("col1"))
)
df = df.withColumn("col2", whens)

df.show()
# +----+---------+
# |col1|     col2|
# +----+---------+
# |   ė|   letter|
# |   2|    digit|
# |    |    empty|
# |   @|spec_char|
# +----+---------+

实现方式 #

直接遍历字典中的规则项，动态追加when分支即可，不需要逐行硬编码。

from pyspark.sql import functions as F

df = spark.createDataFrame([("ė",), ("2",), ("",), ("@",)], ["col1"])

pats = {
  r"^\d$"          : "digit",
  r"^\p{L}$"       : "letter",
  r"^[\p{P}\p{S}]$": "spec_char",
  r"^$"            : "empty"
}

# 转换规则为列表，固定匹配顺序
pat_rules = list(pats.items())
# 初始化第一个判断分支
whens = F.when(F.col("col1").rlike(pat_rules[0][0]), pat_rules[0][1])
# 循环追加剩余分支
for regex, match_val in pat_rules[1:]:
    whens = whens.when(F.col("col1").rlike(regex), match_val)
# 添加兜底逻辑
whens = whens.otherwise(F.col("col1"))

df = df.withColumn("col2", whens)
df.show()

运行后输出和原始硬编码写法完全一致：

+----+---------+
|col1|     col2|
+----+---------+
|   ė|   letter|
|   2|    digit|
|    |    empty|
|   @|spec_char|
+----+---------+

方案说明 #

• 逻辑和原始硬编码完全等价，匹配顺序和字典键的插入顺序保持一致
• 后续新增、修改、删除匹配规则只需要调整pats字典内容，不需要改动判断逻辑代码
• 如果使用Python 3.7以下版本（字典不保证插入顺序），直接将pats改为(正则规则, 返回值)格式的元组列表存储即可，遍历逻辑无需任何修改

如果偏好函数式写法，也可以用reduce省略显式循环，逻辑完全等价：

from functools import reduce

# 省略df和pats初始化代码
pat_rules = list(pats.items())
whens = reduce(
    lambda acc, rule: acc.when(F.col("col1").rlike(rule[0]), rule[1]),
    pat_rules[1:],
    F.when(F.col("col1").rlike(pat_rules[0][0]), pat_rules[0][1])
).otherwise(F.col("col1"))

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