Scala Spark Dataset聚合需求：合并同客户账户列表且禁用sql.functions

现有Dataset数据格式 #

+----------+-----------+--------+--------------------+--------------+------------+--------------+
|customerId|   forename| surname|            accounts|numberAccounts|totalBalance|averageBalance|
+----------+-----------+--------+--------------------+--------------+------------+--------------+
|   IND0001|Christopher|   Black|              [null]|        234454|    99949494|       50000.0|
|   IND0002|  Madeleine|    Kerr|[[IND0002,ACC0262...|        234454|    99949494|       50000.0|
|   IND0002|  Madeleine|    Kerr|[[IND0002,ACC0155...|        234454|    99949494|       50000.0|
|   IND0003|      Sarah| Skinner|[[IND0003,ACC0540...|        234454|    99949494|       50000.0|
|   IND0003|      Sarah| Skinner|[[IND0003,ACC0486...|        234454|    99949494|       50000.0|
|   IND0003|      Sarah| Skinner|[[IND0003,ACC0235...|        234454|    99949494|       50000.0|

目标Dataset数据格式 #

+----------+-----------+--------+--------------------+--------------+------------+--------------+
|customerId|   forename| surname|            accounts|numberAccounts|totalBalance|averageBalance|
+----------+-----------+--------+--------------------+--------------+------------+--------------+
|   IND0001|Christopher|   Black|              [null]|        234454|    99949494|       50000.0|
|   IND0002|  Madeleine|    Kerr|[[IND0002,ACC0262,65], [IND0002,ACC0155,98]]|234454|    99949494|       50000.0|
|   IND0003|      Sarah| Skinner|[[IND0003,ACC0540,76],[IND0003,ACC0486,89],[IND0003,ACC0235,23]]| 234454|    99949494|       50000.0|

需求约束 #

• 禁止使用import spark.sql.functions._，仅使用Dataset的强类型API实现
• 将同一客户（相同customerId、forename、surname）的多条记录合并，把对应账户列表聚合到单条记录中

现有代码片段 #

val joinedData = customerDS
  .joinWith(
    accountDS,
    customerDS("customerId") === accountDS("customerId"),
    "left"
  )
val mappedData = joinedData.map{
  case(c,a) => CustomerAccountOutput(c.customerId, c.forename, c.surname, Seq(a), 234454, 99949494, (100000/2))
}

解决方案 #

可以利用Dataset的groupByKey和mapGroups方法实现聚合，无需引入sql functions。以下是完整实现代码：

1. 假设的样例类定义（需与你的实际数据结构匹配） #

case class Customer(customerId: String, forename: String, surname: String)
case class Account(customerId: String, accountId: String, balance: Int) // 根据实际账户字段调整
case class CustomerAccountOutput(
    customerId: String,
    forename: String,
    surname: String,
    accounts: Seq[Account],
    numberAccounts: Int,
    totalBalance: Long,
    averageBalance: Double
)

2. 聚合实现代码 #

val joinedData = customerDS
  .joinWith(
    accountDS,
    customerDS("customerId") === accountDS("customerId"),
    "left"
  )

// 按客户唯一标识分组，聚合账户列表
val resultDS = joinedData
  .groupByKey { case (customer, _) => (customer.customerId, customer.forename, customer.surname) }
  .mapGroups { case ((cid, fname, sname), recordsIter) =>
    // 收集所有非空账户，处理left join后的null情况
    val aggregatedAccounts = recordsIter.flatMap { case (_, account) => Option(account) }.toSeq
    // 匹配原数据中无账户的情况（如IND0001）
    val finalAccounts = if (aggregatedAccounts.isEmpty) Seq(null) else aggregatedAccounts
    // 构建输出对象，numberAccounts等字段保持你原代码的硬编码值，可根据实际需求计算
    CustomerAccountOutput(
      cid,
      fname,
      sname,
      finalAccounts,
      234454,
      99949494L,
      50000.0
    )
  }

代码说明 #

• groupByKey：以客户的customerId、forename、surname作为分组键，确保同一客户的所有关联记录被分到同一组
• mapGroups：遍历每个分组，通过flatMap+Option过滤掉left join产生的null账户，将有效账户收集为列表
• 特殊处理无账户的客户：如果聚合后账户列表为空，设置为Seq(null)，与原数据格式保持一致

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