PySpark：如何根据列中数值与字典映射生成status列？

问题场景 #

现有PySpark DataFrame包含value列和dict列，每行的字典映射规则不同，需要根据value的值从对应行的dict中取出映射文本，生成status列，预期效果如下：

问题原因 #

最初尝试的代码df.withColumn("status", F.col("dict").getItem(F.col("value")))无法运行，核心原因是**value列类型为double，而字典的键是字符串类型**，类型不匹配导致无法匹配到对应的键值对。硬编码字符串键能运行，但无法满足动态映射的需求。

解决方案 #

将value列的类型转换为与字典键一致的字符串类型即可，步骤如下：

1. 先把double类型的value转为int，避免小数干扰
2. 再将int转为string，确保与字典键类型完全匹配
3. 使用转换后的列作为getItem的参数

完整示例代码 #

from pyspark.sql import SparkSession
import pyspark.sql.functions as F

# 创建测试DataFrame
spark = SparkSession.builder.appName("DictMapExample").getOrCreate()
data = [
    (1.0, {"1": "Text A", "2": "Text B"}),
    (2.0, {"1": "Text A", "2": "Text B"}),
    (0.0, {"0": "Other A", "1": "Other B"})
]
df = spark.createDataFrame(data, ["value", "dict"])

# 转换value类型并生成status列
df_result = df.withColumn(
    "status",
    F.col("dict").getItem(F.col("value").cast("int").cast("string"))
)

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

运行后输出结果符合预期：

+-----+--------------------------------+-------+
|value|dict                            |status |
+-----+--------------------------------+-------+
|1.0  |{1 -> Text A, 2 -> Text B}      |Text A |
|2.0  |{1 -> Text A, 2 -> Text B}      |Text B |
|0.0  |{0 -> Other A, 1 -> Other B}    |Other A|
+-----+--------------------------------+-------+

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