如何在PySpark DataFrame中通过嵌套字典实现值查找？

原始数据与需求 #

现有如下PySpark DataFrame：

from pyspark.sql import Row
from pyspark.sql import SparkSession

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

data = [
    {"foo": "foo1", "buzz": "buzz1"},
    {"foo": "foo2", "buzz": "buzz1"},
    {"foo": "foo1", "buzz": "buzz2"},
    {"foo": "foo2", "buzz": "buzz2"},
]
df = spark.createDataFrame(Row(**x) for x in data)
df.show()

输出：

+-----+----+
| buzz| foo|
+-----+----+
|buzz1|foo1|
|buzz1|foo2|
|buzz2|foo1|
|buzz2|foo2|
+-----+----+

以及用于匹配的嵌套字典：

mapping = {
    "buzz1": {"foo1": "oneone", "foo2": "onetwo"},
    "buzz2": {"foo1": "twoone", "foo2": "twotwo"},
}

需要通过buzz和foo两列的值在嵌套字典中查找对应结果，生成包含combo列的目标DataFrame：

+-----+----+------+
| buzz| foo| combo|
+-----+----+------+
|buzz1|foo1|oneone|
|buzz1|foo2|onetwo|
|buzz2|foo1|twoone|
|buzz2|foo2|twotwo|
+-----+----+------+

实现方法 #

方法一：使用UDF（用户自定义函数）

直接定义UDF接收buzz和foo的值，从嵌套字典中提取对应结果。注意将mapping作为参数传入UDF，避免序列化问题：

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

# 定义UDF，从mapping中查找对应值
lookup_combo = udf(lambda b, f: mapping.get(b, {}).get(f), StringType())

# 添加combo列
result_df = df.withColumn("combo", lookup_combo(df["buzz"], df["foo"]))
result_df.show()

这种方式简单直观，但属于逐行处理，数据量较大时性能不如Spark原生操作。

方法二：转换字典为DataFrame后关联（推荐）

将嵌套字典展开为扁平结构的lookup表，再通过Spark原生的join操作关联原DataFrame，更符合分布式计算的最佳实践：

# 展开嵌套字典，生成lookup数据
lookup_data = []
for buzz_key, foo_map in mapping.items():
    for foo_key, combo_val in foo_map.items():
        lookup_data.append({"buzz": buzz_key, "foo": foo_key, "combo": combo_val})

# 创建lookup DataFrame
lookup_df = spark.createDataFrame(lookup_data)

# 关联原DataFrame与lookup表
result_df = df.join(lookup_df, on=["buzz", "foo"], how="left")
result_df.show()

该方法利用Spark的分布式join优化，在大数据量场景下性能更优，且无需自定义函数，维护性更好。

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