PySpark：如何判断DataFrame列含Map键并生成对应值新列

场景说明 #

假设我们有一个DataFrame，其中某列（如target_col）存储着单个字符串或字符串数组，需要判断该列的每条记录是否存在于预定义的Map键中，若存在则将对应Map值填充到新列（如new_col）中。

1. 预定义Map的准备 #

可以通过两种方式创建Spark可识别的Map对象：

• 方式一：使用PySpark内置的create_map函数构建
• 方式二：将Python字典转换为Spark的MapType类型

示例代码：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, create_map, lit, array_contains, first, explode
from pyspark.sql.types import MapType, StringType

spark = SparkSession.builder.appName("MapMatch").getOrCreate()

# 方式一：用create_map构建
predefined_map = create_map(
    lit("key1"), lit("value1"),
    lit("key2"), lit("value2"),
    lit("key3"), lit("value3")
)

# 方式二：Python字典转Spark MapType
py_dict = {"key1": "value1", "key2": "value2", "key3": "value3"}
predefined_map = lit(py_dict).cast(MapType(StringType(), StringType()))

2. 目标列为单个字符串的情况 #

如果target_col的每条记录是单个字符串（即可能的Map键），直接通过Map的get方法取值，不存在时可设为默认值（如null或自定义文本）：

示例DataFrame：

data = [("key1",), ("key4",), ("key2",), ("key5",)]
df = spark.createDataFrame(data, ["target_col"])

填充新列：

# 直接用Map的get方法，不存在返回null
df_with_new_col = df.withColumn("new_col", predefined_map.getItem(col("target_col")))

# 或者用when+otherwise指定自定义默认值（比如"no_match"）
df_with_new_col = df.withColumn(
    "new_col",
    when(col("target_col").isin(list(py_dict.keys())), predefined_map.getItem(col("target_col")))
    .otherwise("no_match")
)

df_with_new_col.show()

输出结果：

+----------+--------+
|target_col| new_col|
+----------+--------+
|      key1|  value1|
|      key4|no_match|
|      key2|  value2|
|      key5|no_match|
+----------+--------+

3. 目标列为字符串数组的情况 #

如果target_col是字符串数组，需要先筛选出数组中存在于Map键的元素，再取对应值（以下示例以取第一个匹配值为例）：

示例DataFrame：

data = [(["key1", "key4"],), (["key5", "key2"],), (["key6"],)]
df = spark.createDataFrame(data, ["target_col"])

实现逻辑：

# 展开数组元素→匹配Map值→过滤有效匹配→聚合取第一个匹配值→补全无匹配的记录
df_with_new_col = df.withColumn("exploded", explode(col("target_col"))) \
    .withColumn("matched_value", predefined_map.getItem(col("exploded"))) \
    .filter(col("matched_value").isNotNull()) \
    .groupBy("target_col") \
    .agg(first("matched_value").alias("new_col")) \
    .join(df, on="target_col", how="right") \
    .fillna("no_match", subset=["new_col"])

df_with_new_col.show()

输出结果：

+----------+--------+
|target_col| new_col|
+----------+--------+
|[key1,key4]|  value1|
|[key5,key2]|  value2|
|    [key6]|no_match|
+----------+--------+

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