Spark 2.2分组聚合求和及收集关联数据的实现求助

嗨，作为Spark老玩家，我来帮你搞定这个Spark 2.2的分组聚合需求，一步步实现你想要的结果～

需求回顾 #

我们要完成这些目标：

• 按a_id和b_sum字段对记录分组
• 对每个分组内的m_cd对应的td_cnt求和，把这些键值对收集成数组
• 新增td_cnt_sum字段，存储分组内所有td_cnt的总和
• 过滤掉输入中m_cd为["null"]、td_cnt为[null]的无效记录

实现步骤（Python版本） #

1. 初始化SparkSession并加载数据

先创建SparkSession，这里用模拟数据演示，实际可以直接用spark.read.json("path/to/your/file.json")读取本地或HDFS上的JSON文件：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum, collect_list, struct

# 初始化SparkSession
spark = SparkSession.builder.appName("McdTdCntAgg").getOrCreate()

# 模拟输入的JSON数据
data = [
    {"a_id":6336,"b_sum":10.0,"m_cd":["abc00053"],"td_cnt":[10.0]},
    {"a_id":6336,"b_sum":10.0,"m_cd":["abc00053"],"td_cnt":[5.0]},
    {"a_id":6336,"b_sum":10.0,"m_cd":["abc00054"],"td_cnt":[20.0]},
    {"a_id":6336,"b_sum":10.0,"m_cd":["abc00056"],"td_cnt":[30.0]},
    {"a_id":6339,"b_sum":10.0,"m_cd":["abc00051"],"td_cnt":[12.0]},
    {"a_id":6339,"b_sum":10.0,"m_cd":["abc00057"],"td_cnt":[7.0]},
    {"a_id":6339,"b_sum":10.0,"m_cd":["abc00055"],"td_cnt":[10.0]},
    {"a_id":6339,"b_sum":10.0,"m_cd":["abc00058"],"td_cnt":[20.0]},
    {"a_id":6339,"b_sum":10.0,"m_cd":["null"],"td_cnt":[None]}
]

df = spark.createDataFrame(data)

2. 清理数据：提取数组元素并过滤无效记录

输入里的m_cd和td_cnt都是单元素数组，我们先把数组中的实际值提取出来，同时过滤掉无效的null记录：

df_clean = df.withColumn("m_cd_val", col("m_cd").getItem(0)) \
             .withColumn("td_cnt_val", col("td_cnt").getItem(0)) \
             .filter((col("m_cd_val") != "null") & (col("td_cnt_val").isNotNull()))

3. 两次分组聚合得到最终结果

第一次分组先按a_id、b_sum、m_cd_val聚合，计算每个m_cd对应的td_cnt总和；第二次再按a_id、b_sum分组，收集键值对数组并计算分组的总求和：

# 第一步：计算每个m_cd在分组内的td_cnt总和
grouped_mcd = df_clean.groupBy("a_id", "b_sum", "m_cd_val") \
                      .agg(sum("td_cnt_val").alias("total_td_cnt"))

# 第二步：收集键值对数组并计算分组总和
final_df = grouped_mcd.groupBy("a_id", "b_sum") \
                      .agg(
                          collect_list(struct("m_cd_val", "total_td_cnt")).alias("td_cnt"),
                          sum("total_td_cnt").alias("td_cnt_sum")
                      )

# 查看结果
final_df.show(truncate=False)

执行后得到的结果结构和你预期的一致，td_cnt是包含键值对结构体的数组，td_cnt_sum是分组的总和。

实现步骤（Scala版本） #

如果习惯用Scala开发，代码逻辑完全一致：

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

val spark = SparkSession.builder.appName("McdTdCntAgg").getOrCreate()

// 模拟输入数据
val data = Seq(
  (6336, 10.0, Array("abc00053"), Array(10.0)),
  (6336, 10.0, Array("abc00053"), Array(5.0)),
  (6336, 10.0, Array("abc00054"), Array(20.0)),
  (6336, 10.0, Array("abc00056"), Array(30.0)),
  (6339, 10.0, Array("abc00051"), Array(12.0)),
  (6339, 10.0, Array("abc00057"), Array(7.0)),
  (6339, 10.0, Array("abc00055"), Array(10.0)),
  (6339, 10.0, Array("abc00058"), Array(20.0)),
  (6339, 10.0, Array("null"), Array(null.asInstanceOf[Double]))
).toDF("a_id", "b_sum", "m_cd", "td_cnt")

// 清理数据
val dfClean = data.withColumn("m_cd_val", col("m_cd")(0))
                  .withColumn("td_cnt_val", col("td_cnt")(0))
                  .filter(col("m_cd_val") =!= "null" && col("td_cnt_val").isNotNull)

// 分组聚合
val groupedMcd = dfClean.groupBy("a_id", "b_sum", "m_cd_val")
                        .agg(sum("td_cnt_val").alias("total_td_cnt"))

val finalDf = groupedMcd.groupBy("a_id", "b_sum")
                        .agg(
                          collect_list(struct("m_cd_val", "total_td_cnt")).alias("td_cnt"),
                          sum("total_td_cnt").alias("td_cnt_sum")
                        )

finalDf.show(false)

可选：将td_cnt转为字符串格式 #

如果需要把td_cnt数组转成你示例中的["abc00053":15.0,...]字符串格式，可以用UDF处理：

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 定义UDF格式化键值对
format_kv = udf(lambda kv: f'"{kv["m_cd_val"]}":{kv["total_td_cnt"]}', StringType())

# 应用UDF并重新收集数组
final_df_str = final_df.withColumn("td_cnt", collect_list(format_kv(col("td_cnt"))))
final_df_str.show(truncate=False)

这样得到的td_cnt就是字符串数组，更贴近你给出的示例格式。

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