Spark DataFrame多日期列分组计数聚合的高效优化方案问询

问题背景 #

现有Spark DataFrame结构如下，每行包含唯一GUID、所属Group、Occurence及多个日期列（实际场景有15个），日期值分散在不同列中。需求是按Group统计每个日期在各日期列的出现次数，原方案对每个日期列单独分组聚合再关联，数据量大时性能开销极高，需要更高效的实现方式。

示例DataFrame定义：

val df= Seq(
("GUID11",  "Grp1",     1,      "01-01-2023",       None),
("GUID12",  "Grp2",     1,      "01-01-2023",       Some("01-01-2023")),
("GUID21",  "Grp1",     2,      "01-02-2023",       Some("01-01-2023")),
("GUID22",  "Grp1",     2,      "01-02-2023",       Some("01-02-2023")),
("GUID23",  "Grp1",     2,      "01-02-2023",       None),
("GUID31",  "Grp2",     3,      "01-03-2023",       Some("01-03-2023")),
("GUID32",  "Grp1",     3,      "01-03-2023",       Some("01-04-2023")),
("GUID41",  "Grp1",     4,      "01-04-2023",       Some("01-04-2023")),
("GUID42",  "Grp1",     4,      "01-04-2023",       Some("01-04-2023")),
("GUID44",  "Grp3",     4,      "01-04-2023",       None),
("GUID45",  "Grp1",     4,      "01-04-2023",       Some("01-04-2023")),
("GUID51",  "Grp1",     5,      "01-05-2023",       Some("01-05-2023")),
("GUID61",  "Grp2",     6,      "01-06-2023",       Some("01-06-2023")),
("GUID62",  "Grp1",     6,      "01-06-2023",       None),
("GUID63",  "Grp3",     6,      "01-06-2023",       Some("01-06-2023")),
("GUID64",  "Grp3",     6,      "01-06-2023",       Some("01-06-2023")),
("GUID65",  "Grp1",     6,      "01-06-2023",       Some("01-06-2023"))
).toDF("GUID", "Group", "Occurence","Date1","Date2")

预期输出：

+-----+----------+------+------+
|Group|      Date|Count1|Count2|
+-----+----------+------+------+
| Grp1|2023-01-01|     1|     1|
| Grp2|2023-01-01|     1|     1|
| Grp1|2023-01-02|     3|     1|
| Grp1|2023-01-03|     1|  null|
| Grp2|2023-01-03|     1|     1|
| Grp1|2023-01-04|     3|     4|
| Grp3|2023-01-04|     1|  null|
| Grp1|2023-01-05|     1|     1|
| Grp1|2023-01-06|     2|     1|
| Grp2|2023-01-06|     1|     1|
| Grp3|2023-01-06|     2|     2|
+-----+----------+------+------+

优化方案 #

核心思路是将宽表转长表后聚合，再转回宽表，仅需一次分组聚合操作，避免多次join的性能损耗。

实现代码 #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.StringType

// 1. 转换所有日期列的格式为date类型
val dateCols = List("Date1", "Date2") // 实际场景替换为15个日期列的列表
val dfWithTypedDates = dateCols.foldLeft(df) { (accDF, colName) =>
  accDF.withColumn(colName, to_date(col(colName).cast(StringType), "MM-dd-yyyy"))
}

// 2. 将宽表转为长表：保留Group，将每个日期列转为（日期列名，日期值）的行
val longDF = dfWithTypedDates
  .select(
    col("Group"),
    explode(
      array(
        dateCols.map(colName => struct(lit(colName).as("ColName"), col(colName).as("Date"))): _*
      )
    ).as("DateInfo")
  )
  .select(
    col("Group"),
    col("DateInfo.ColName"),
    col("DateInfo.Date")
  )
  .filter(col("Date").isNotNull) // 过滤空日期

// 3. 按Group、Date、ColName分组统计次数
val aggregatedDF = longDF
  .groupBy("Group", "Date", "ColName")
  .agg(count("*").as("Count"))

// 4. 将长表转回宽表，ColName作为列名，Count作为值
val resultDF = aggregatedDF
  .groupBy("Group", "Date")
  .pivot("ColName", dateCols) // 指定pivot的列名列表，保证顺序一致
  .agg(first("Count")) // 每个分组只有一个Count值，用first即可
  .withColumnRenamed("Date1", "Count1")
  .withColumnRenamed("Date2", "Count2")
  .orderBy("Date", "Group")

// 查看结果
resultDF.show()

方案优势 #

• 原方案针对15个日期列需要执行15次分组聚合+14次关联操作，数据量大时会产生大量Shuffle开销；
• 优化方案仅需一次全量数据扫描、一次分组聚合和一次pivot操作，Shuffle次数大幅减少，性能提升显著；
• 代码扩展性强：新增/删除日期列时，仅需修改dateCols列表，无需调整其他逻辑。

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