如何用多列作为嵌套字典映射创建PySpark DataFrame新列？

问题场景 #

给定多层嵌套字典：

dict_prob = {"a":{"x1":"y1","x2":"y2"},"b":{"m1":"n1","m2":"n2"}}

以及包含index、col1、col2的PySpark DataFrame：

需要生成新列col3，值为通过col1和col2逐层查找嵌套字典得到的结果，最终输出：

以下提供两种适配多层嵌套（支持4-5层）的解决方案：

方案一：自定义UDF（灵活适配任意层数） #

这种方法通过循环逐层查找嵌套字典，适配任意层数的嵌套结构，是最通用的方案。

步骤1：创建示例数据 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

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

dict_prob = {"a":{"x1":"y1","x2":"y2"},"b":{"m1":"n1","m2":"n2"}}
data = [(0, "a", "x1"), (1, "a", "x2"), (2, "b", "m2")]
df = spark.createDataFrame(data, ["index", "col1", "col2"])

步骤2：编写查找函数并注册UDF #

def lookup_nested_dict(d, *keys):
    current = d
    for key in keys:
        # 逐层查找，键不存在则返回None（可根据需求修改为默认值）
        if current and key in current:
            current = current[key]
        else:
            return None
    return current

# 注册UDF，传入col1和col2作为查找键
lookup_udf = udf(lambda k1, k2: lookup_nested_dict(dict_prob, k1, k2), StringType())

步骤3：生成新列 #

df_with_col3 = df.withColumn("col3", lookup_udf(df.col1, df.col2))
df_with_col3.show()

扩展到4-5层嵌套 #

如果需要用4列作为键查找5层嵌套字典，只需修改UDF的参数：

# 假设新增col3、col4作为额外的查找键
lookup_udf_4level = udf(lambda k1,k2,k3,k4: lookup_nested_dict(dict_prob, k1,k2,k3,k4), StringType())
df_with_col3 = df.withColumn("col3", lookup_udf_4level(df.col1, df.col2, df.col3, df.col4))

性能优化：广播字典 #

如果字典体积较大，建议广播字典避免每个任务重复加载：

from pyspark.sql.functions import broadcast

broadcast_dict = spark.sparkContext.broadcast(dict_prob)

def lookup_nested_dict_broadcast(*keys):
    current = broadcast_dict.value
    for key in keys:
        if current and key in current:
            current = current[key]
        else:
            return None
    return current

lookup_udf_broadcast = udf(lookup_nested_dict_broadcast, StringType())
df_with_col3 = df.withColumn("col3", lookup_udf_broadcast(df.col1, df.col2))

方案二：使用PySpark内置函数（固定层数场景） #

如果嵌套层数固定，可使用PySpark内置的create_map和getItem函数实现，性能优于UDF。

步骤1：创建嵌套Map类型 #

from pyspark.sql.functions import lit, create_map, col

# 逐层创建嵌套Map，对应原字典的结构
map_level1 = create_map(
    lit("a"), create_map(lit("x1"), lit("y1"), lit("x2"), lit("y2")),
    lit("b"), create_map(lit("m1"), lit("n1"), lit("m2"), lit("n2"))
)

步骤2：链式调用getItem生成新列 #

df_with_col3 = df.withColumn("col3", map_level1.getItem(col("col1")).getItem(col("col2")))
df_with_col3.show()

扩展到4-5层嵌套 #

对于5层嵌套，只需继续链式调用getItem：

# 假设map_level4是4层嵌套的Map，col1-col4为查找键
df_with_col3 = df.withColumn(
    "col3", 
    map_level4.getItem(col("col1"))
              .getItem(col("col2"))
              .getItem(col("col3"))
              .getItem(col("col4"))
)

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