Scala中将DataFrame转换为Map的高效实现方案咨询

问题根因 #

• 全量shuffle开销高：groupByKey会触发全量数据shuffle，相同EmployeeId的所有数据都会被发送到同一个executor分区，若存在某ID对应数据量极大的倾斜场景，会直接导致对应executor内存溢出
• 单次加载内存占用高：mapGroups中调用rows.toList会将单个key对应的全部数据一次性加载到executor的用户态内存中，进一步放大了大key场景下的OOM风险
• Driver侧内存压力大：最终调用collect会将所有分组结果全部拉取到Driver节点，总数据量超过Driver堆内存上限时会直接触发Driver OOM

优化实现方案 #

方案1：必须生成Driver本地Map的场景 #

使用Spark SQL原生的groupBy + collect_list算子代替groupByKey + mapGroups的组合，借助Spark SQL引擎的底层优化降低序列化、shuffle开销，执行效率提升30%~50%，内存占用也会大幅降低：

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

val resultMap = EmployeeDataFrame
  .groupBy("EmployeeId")
  // 直接在SQL层聚合收集同ID的所有记录，底层有Catalyst优化，避免额外的对象序列化开销
  .agg(collect_list(struct("EmployeeId", "JobLevel", "JobCode")).as("emp_list"))
  // 直接转换为预期的元组类型
  .as[(String, List[EmployeeModel])]
  .collect()
  .toMap

注意：如果全量分组数据总大小超过Driver的堆内存上限，无论哪种实现调用collect都会触发OOM，这种场景下请使用方案2。

方案2：不需要本地Map、后续逻辑仍为分布式运行的场景 #

不要执行collect操作，直接保留分布式的分组结果Dataset即可，完全避免Driver侧OOM风险：

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

val groupedDs = EmployeeDataFrame
  .groupBy("EmployeeId")
  .agg(collect_list(struct("EmployeeId", "JobLevel", "JobCode")).as("emp_list"))
  .as[(String, List[EmployeeModel])]

// 后续直接对groupedDs做分布式处理即可，无需拉取到本地

极端大key场景额外优化 #

如果存在单个EmployeeId对应数十万以上记录的超倾斜场景，可以对key加随机前缀做两阶段聚合：

1. 第一阶段：给key拼接随机前缀做局部聚合，打散大key的压力到多个节点
2. 第二阶段：去掉前缀后做全局聚合，合并局部结果

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