如何在Scala中将嵌套映射数组合并为单个映射？

方法一：使用Spark内置函数（推荐） #

无需自定义UDF，通过调整分组逻辑直接生成目标格式的嵌套映射：

import org.apache.spark.sql.functions.{col, map_from_entries, struct, collect_list}

// 假设原始DataFrame名为df，包含ID、LOCATION、GROUP、COUNT列
val resultDF = df
  // 第一步：按ID和LOCATION分组，将同一地点的所有GROUP-COUNT对合并为一个映射
  .groupBy(col("ID"), col("LOCATION"))
  .agg(
    map_from_entries(collect_list(struct(col("GROUP"), col("COUNT"))))
      .alias("group_count_map")
  )
  // 第二步：按ID分组，将同一ID的所有LOCATION-group_count_map对合并为最终嵌套映射
  .groupBy(col("ID"))
  .agg(
    map_from_entries(collect_list(struct(col("LOCATION"), col("group_count_map"))))
      .alias("map")
  )

方法二：使用UDF合并现有嵌套映射数组 #

如果无法回到原始数据，只能基于已生成的map列（数组类型）处理，可以用UDF实现合并：

import org.apache.spark.sql.functions.udf
import scala.collection.mutable

// 定义UDF：合并数组中的嵌套映射，同一地点的相同GROUP会累加计数
val mergeNestedMaps = udf((mapArray: Array[Map[String, Map[String, Int]]]) => {
  val merged = mutable.HashMap[String, mutable.HashMap[String, Int]]()
  
  mapArray.foreach { singleMap =>
    singleMap.foreach { case (location, groupMap) =>
      val groupCounts = merged.getOrElseUpdate(location, mutable.HashMap())
      groupMap.foreach { case (group, count) =>
        // 若需覆盖而不是累加，替换为 groupCounts(group) = count
        groupCounts(group) = groupCounts.getOrElse(group, 0) + count
      }
    }
  }
  
  // 转换为不可变Map返回
  merged.mapValues(_.toMap).toMap
})

// 应用UDF到现有DataFrame
val resultDF = existingDF.withColumn("map", mergeNestedMaps(col("map")))

说明 #

• 方法一利用Spark原生聚合函数，性能更优且避免自定义代码维护成本，是首选方案。
• 方法二中的UDF默认对同一地点的相同GROUP进行计数累加，若需求是覆盖而非累加，修改注释处的代码即可。

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