PySpark：无需UDF，从数组列匹配字典键获取对应值数组

可以通过Spark内置的Map与数组操作函数实现需求，无需编写UDF，具体步骤如下：

1. 将广播变量的键值对转换为Spark Map类型
   把广播变量中的字典转换成Spark原生的Map列，方便后续用内置函数快速匹配键值。

2. 遍历数组匹配键并提取对应值
   用transform遍历cities_array的每个元素，通过map_get获取匹配的国家代码，再用filter过滤掉不匹配产生的null值，最终得到有效国家代码组成的数组。

完整代码实现 #

from pyspark.sql import functions as F

# 修正广播变量键的大小写，与处理后的address列匹配
city_map = F.create_map(*[F.lit(k.lower()), F.lit(v) for k, v in cities_name_to_code.value.items()])

# 处理数组，提取匹配的国家代码
df_result = df_with_codes.withColumn(
    'country_codes',
    F.filter(
        F.transform(F.col('cities_array'), lambda x: F.map_get(city_map, x)),
        lambda x: x.isNotNull()
    )
).drop('cities_array')  # 可选：移除中间生成的cities_array列

# 查看最终结果
df_result.show(truncate=False)

结果说明 #

运行后输出如下（匹配到的国家代码会组成数组，无匹配则为空数组）：

+-------------------------------+-------------+
|address                        |country_codes|
+-------------------------------+-------------+
|definitely somewhere           |[]           |
|Las Vegas                      |[US]         |
|其他                           |[]           |
|null                           |[]           |
|                               |[]           |
|Pucela Madrid Langreo, España  |[ES]         |
|Trenches, With Egbon Adugbo    |[]           |
+-------------------------------+-------------+

注意：原代码中对address做了小写转换，因此创建Map时需将广播变量的键统一转为小写，避免大小写不匹配导致的匹配失败。

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