在PySpark DataFrame中递归展平所有Map类型列

我有一个包含多个Map类型列的PySpark DataFrame，需要递归展平所有Map类型列。当前已知personal和financial是Map类型列，后续可能还会有更多这类列。

输入DataFrame：

输出DataFrame：

解决方案 #

自动识别并展平所有Map列 #

以下代码会自动检测DataFrame中的所有Map类型列，提取其中的所有键作为独立列，同时保留原有非Map列的数据：

from pyspark.sql import SparkSession
from pyspark.sql.types import MapType
from pyspark.sql.functions import col

# 初始化SparkSession
spark = SparkSession.builder.appName("FlattenMapColumns").getOrCreate()

# 构造示例输入DataFrame（可替换为你的实际数据）
data = [
    (1, "A", "M", {"age":20, "city":"Dallas", "State":"Texas"}, {"salary":10000, "bonus":2000, "tax":1500}),
    (2, "B", "F", {"city":"Houston", "State":"Texas", "Zipcode":77001}, {"salary":12000, "tax":1800}),
    (3, "C", "M", {"age":22, "city":"San Jose", "Zipcode":940088}, {"salary":2000, "bonus":500})
]
schema = ["id", "name", "Gender", "personal", "financial"]
df = spark.createDataFrame(data, schema=schema)

def flatten_map_columns(df):
    # 收集所有非Map类型的列
    non_map_cols = [col(c) for c, dtype in df.dtypes if not isinstance(df.schema[c].dataType, MapType)]
    # 筛选出所有Map类型的列
    map_cols = [c for c, dtype in df.dtypes if isinstance(df.schema[c].dataType, MapType)]
    
    # 处理每个Map列
    for map_col in map_cols:
        # 获取该Map列包含的所有唯一键
        all_keys = df.select(f"{map_col}.keys").distinct().rdd.flatMap(lambda x: x[0]).collect()
        # 为每个键生成新列，值为Map中对应键的内容，不存在则为null
        for key in all_keys:
            non_map_cols.append(col(map_col).getItem(key).alias(key.lower()))
    
    # 返回展平后的DataFrame
    return df.select(non_map_cols)

# 执行展平操作
flattened_df = flatten_map_columns(df)
# 查看结果
flattened_df.show()

关键说明 #

• 自动适配新增Map列：无需手动指定Map列名，代码会自动识别所有Map类型列，后续新增此类列也能直接处理。
• 完整提取所有键：通过keys方法获取Map列的所有唯一键，确保不会遗漏任何可能的字段。
• 格式统一：将提取出的列名转为小写，与输出示例格式保持一致；不存在的键会自动填充null，保证数据完整性。

内容的提问来源于stack exchange，提问作者Ramineni Ravi Teja