如何在PySpark中将未知键的字典类型列拆分为多个独立列

前置说明 #

本方案适用于DataFrame中features列为MapType类型、所有行字典键完全一致的场景，无需提前手动录入所有键名即可完成拆分。

操作步骤 #

• 导入依赖函数

from pyspark.sql import functions as F

• 自动提取所有唯一的特征键

feature_keys = df.select(F.explode(F.map_keys("features"))) \
                 .agg(F.collect_set("key").alias("keys")) \
                 .first()["keys"]
# 若需要特征列按固定顺序排列，可补充排序逻辑：
# feature_keys = sorted(feature_keys)

• 动态构造目标DataFrame

result_df = df.select(
    "id1",
    "id2",
    *[F.col("features").getItem(k).alias(k) for k in feature_keys]
)

特殊情况处理 #

如果你的features列存储的是JSON字符串而非原生MapType，需要先做类型转换：

# 可根据实际值的类型调整schema中的值类型，比如值为整数则改为map<string, int>
df = df.withColumn("features", F.from_json(F.col("features"), "map<string, double>"))

完整测试示例 #

可直接运行以下代码验证效果：

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

# 初始化Spark会话
spark = SparkSession.builder.appName("split_map_column").getOrCreate()

# 构造示例数据
test_data = [
    (1341205, "a232523", {"attr1.feature1": 0.25, "attr1.feature2": 0.0, "attr2.feature1": -0.43}),
    (553654, "a325933", {"attr1.feature1": 0.3, "attr1.feature2": 0.70, "attr2.feature1": 0.11}),
    (573786, "a9923823", {"attr1.feature1": -0.1, "attr1.feature2": 0.20, "attr2.feature1": 0.12})
]
df = spark.createDataFrame(test_data, schema=["id1", "id2", "features"])

# 执行拆分逻辑
feature_keys = df.select(F.explode(F.map_keys("features"))).agg(F.collect_set("key")).first()[0]
feature_keys = sorted(feature_keys)
result_df = df.select("id1", "id2", *[F.col("features")[k].alias(k) for k in feature_keys])

# 输出结果
result_df.show()

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