如何高效将Scala DataFrame转换为以Code为键的HashMap结构

实现思路 #

利用Spark原生分布式算子完成分组聚合，避免Driver端单节点处理全量数据，性能最优：

1. 先将非分组字段封装为结构体，保证分组后同组数据的字段关联关系
2. 按Code字段分组，用collect_list收集同组所有结构体数据
3. （仅小数据集适用）将聚合后的数据拉取到Driver端，转换为目标HashMap结构

实现代码 #

import org.apache.spark.sql.functions._
import scala.collection.mutable.HashMap

// 假设你的原始DataFrame变量名为df
val resultDf = df.groupBy("Code")
  .agg(collect_list(struct("Name", "Age", "ID")).as("records"))

// 拉取到Driver端转HashMap，仅适合结果集不大的场景
val targetMap: HashMap[String, List[Map[String, Any]]] = HashMap(
  resultDf.collect().map(row => {
    val code = row.getAs[String]("Code")
    val records = row.getAs[Seq[org.apache.spark.sql.Row]]("records").map(recordRow => 
      Map(
        "Name" -> recordRow.getAs[String]("Name"),
        "Age" -> recordRow.getAs[Int]("Age"),
        "ID" -> recordRow.getAs[Int]("ID")
      )
    ).toList
    code -> records
  }): _*
)

注意事项 #

• 若数据量极大，不要直接调用collect拉取全量数据到Driver，可直接将resultDf写入外部存储（如HDFS、数据库等）
• 如需类型安全，可提前定义样例类封装用户字段，避免Row的弱类型问题：

  case class User(Name: String, Age: Int, ID: Int)
  // 转换Row为样例类的逻辑替换为：
  val records = row.getAs[Seq[Row]]("records").map(r => User(r.getAs[String]("Name"), r.getAs[Int]("Age"), r.getAs[Int]("ID"))).toList
  

内容的提问来源于stack exchange，提问作者Asad Bin Imtiaz