PySpark如何将字典类型列拆分为多个独立列？

现有如下PySpark代码：

from pyspark.sql.types import MapType, StringType, FloatType
from pyspark.sql.functions import udf, col

@udf(returnType=MapType(StringType(), FloatType()))
def postprocess(data):
    ret = dict()
    # 此处为生成字典的逻辑
    ...

    return ret

ret = postprocess(col('data'))
print(ret) # Column<'postprocess(data)'>

需要从上述代码生成的MapType字典列中提取键值对，生成独立的列。例如字典内容为{"key1": 0.1, "key2": 0.3}时，期望得到如下结果：

解决方案 #

情况1：已知字典所有键名

如果提前明确Map中的键，可以直接通过getItem()方法或下标访问提取值，生成新列：

from pyspark.sql import functions as F

# 先将Map列添加到DataFrame中
df = df.withColumn("map_column", postprocess(F.col("data")))

# 提取每个键作为独立列
df = df.withColumn("key1", F.col("map_column").getItem("key1")) \
       .withColumn("key2", F.col("map_column")["key2"])  # 两种写法等价

# 可选：删除原始的Map列
df = df.drop("map_column")

df.show()

情况2：字典键名动态不确定

如果Map中的键不固定，需要先获取所有唯一键，再批量生成列：

from pyspark.sql import functions as F

# 添加Map列到DataFrame
df = df.withColumn("map_column", postprocess(F.col("data")))

# 提取所有唯一的键
all_unique_keys = df.select(F.explode(F.map_keys(F.col("map_column")))) \
                    .distinct() \
                    .rdd.flatMap(lambda x: x) \
                    .collect()

# 批量生成新列
for key in all_unique_keys:
    df = df.withColumn(key, F.col("map_column").getItem(key))

# 可选：删除原始Map列
df = df.drop("map_column")

df.show()

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