Spark DataFrame中Map类型列自动展开为多列的实现方法

我之前处理这类需求的时候，最烦的就是手动列所有键——要是Map里的键经常变，那代码根本没法维护！下面给你分享两种实用的自动展开方法，分别用Scala和Python实现，核心思路都是先提取所有唯一的键，再遍历键自动生成新列：

Scala 实现 #

步骤1：获取所有唯一的键 #

首先需要从properties列的Map中提取出所有存在的键，去重后得到完整的键列表：

import org.apache.spark.sql.functions.{explode, map_keys}

// 提取所有键并去重，收集到Driver端
val allKeys = df.select(explode(map_keys($"properties"))).distinct().as[String].collect()

这里map_keys会取出每个Map的所有键返回数组，explode把数组拆成单行，distinct去重后，collect()把结果转为本地字符串数组。

步骤2：自动生成所有新列 #

用foldLeft遍历键列表，逐个给DataFrame添加对应列：

val dfWithExpandedCols = allKeys.foldLeft(df) { (accDF, key) =>
  // 从Map中提取对应键的值作为新列，不存在的键会返回null
  accDF.withColumn(key, $"properties".getItem(key))
}

如果需要给不存在的键设置默认值，可以用coalesce处理：

import org.apache.spark.sql.functions.lit

val dfWithExpandedCols = allKeys.foldLeft(df) { (accDF, key) =>
  accDF.withColumn(key, coalesce($"properties".getItem(key), lit("默认值")))
}

Python 实现 #

步骤1：获取所有唯一的键 #

和Scala思路一致，提取并去重所有键：

from pyspark.sql.functions import explode, map_keys

# 提取键、去重后转为Python列表
all_keys = df.select(explode(map_keys(df.properties))).distinct().rdd.flatMap(lambda x: x).collect()

步骤2：自动生成所有新列 #

可以用简单的循环遍历，或者用reduce函数实现更简洁的函数式风格：

方式1：普通循环

df_with_expanded_cols = df
for key in all_keys:
    # 提取对应键的值作为新列
    df_with_expanded_cols = df_with_expanded_cols.withColumn(key, df.properties.getItem(key))

方式2：reduce函数式实现

from functools import reduce
from pyspark.sql.functions import col

df_with_expanded_cols = reduce(
    lambda acc_df, key: acc_df.withColumn(key, col("properties").getItem(key)),
    all_keys,
    df
)

同样，要设置默认值的话，结合coalesce即可：

from pyspark.sql.functions import lit, coalesce

df_with_expanded_cols = reduce(
    lambda acc_df, key: acc_df.withColumn(key, coalesce(col("properties").getItem(key), lit("默认值"))),
    all_keys,
    df
)

注意事项 #

• 如果你的数据集非常大，collect()会把键列表加载到Driver端内存，如果键的数量特别多，可能会导致内存压力，这种情况下可以考虑先抽样获取键，或者确认键的数量在可控范围内。
• 不同行的Map中不存在的键，对应新列的值会是null，根据业务需求可以选择保留null或者设置默认值。

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