Scala Spark中基于滚动窗口计算事件共现矩阵的最高效实现方法

针对超大数据集的共现矩阵计算，核心优化方向是避免全量自连接产生的海量笛卡尔积，采用按用户分组后内存内窗口配对的方案，shuffle量仅等于原始数据量，性能远高于自连接实现。

实现步骤 #

1. 数据预处理：将字符串日期转换为纪元天整数（方便快速计算日期差），同时过滤无效空值。
2. 按user_id分组，对每个用户的事件按日期升序排序。
3. 对每个用户的排序后事件列表，仅生成符合时间窗口要求的事件对，同时要求event_id_1 < event_id_2避免重复计数（例如(100,101)和(101,100)视为同一共现对，仅统计一次）。
4. 全局分组统计每个事件对的共现次数。

代码示例 #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.Row
import java.time.LocalDate
import java.time.format.DateTimeFormatter

// 配置参数
val N_DAYS = 15
val dateFormatter = DateTimeFormatter.ofPattern("dd/MM/yyyy")

// 1. 预处理数据，转换日期为纪元天
val rawDF = spark.read.xxx // 读取原始数据
val processedDF = rawDF
  .withColumn("event_date", to_date(col("date"), "dd/MM/yyyy"))
  .withColumn("day_num", datediff(col("event_date"), lit("1970-01-01")))
  .select("user_id", "event_id", "day_num")
  .na.drop() // 过滤无效数据

// 2. 按user_id分组后生成合法共现对
val cooccurPairs = processedDF
  .rdd
  .groupBy(row => row.getAs[Int]("user_id"))
  .flatMap { case (userId, events) =>
    // 按日期升序排序
    val sortedEvents = events.toSeq.sortBy(_.getAs[Int]("day_num"))
    val pairs = collection.mutable.ListBuffer[(Int, Int)]()
    // 遍历每个事件，仅配对后面N天内的事件
    for (i <- sortedEvents.indices) {
      val e1 = sortedEvents(i)
      val e1Id = e1.getAs[Int]("event_id")
      val e1Day = e1.getAs[Int]("day_num")
      // 最多遍历到超过N天的事件就停止
      for (j <- i + 1 until sortedEvents.length) {
        val e2 = sortedEvents(j)
        val e2Day = e2.getAs[Int]("day_num")
        if (e2Day - e1Day <= N_DAYS) {
          val e2Id = e2.getAs[Int]("event_id")
          // 保证小的id在前，避免重复配对
          if (e1Id < e2Id) {
            pairs.append((e1Id, e2Id))
          } else if (e2Id < e1Id) {
            pairs.append((e2Id, e1Id))
          }
          // 相同event_id不需要统计可以直接跳过
        } else {
          // 已经排序，后面的事件日期更大，直接跳出循环
          break
        }
      }
    }
    pairs
  }

// 3. 统计共现次数，转为DataFrame输出
val resultDF = cooccurPairs
  .map(pair => (pair, 1))
  .reduceByKey(_ + _)
  .map { case ((e1, e2), cnt) => (e1, e2, cnt) }
  .toDF("event_id_1", "event_id_2", "count")

resultDF.show()

额外性能优化点 #

• 若单用户事件量超过1000条，可将内层循环替换为二分查找定位N天窗口的右边界，时间复杂度从O(k²)降到O(klogk)，k为单用户事件数。
• 数据倾斜优化：若存在热点用户（单用户事件数超过10万），可对热点user_id加盐拆分后再配对，避免单个任务执行时间过长。
• 可优先使用DataFrame的高阶函数实现，避免RDD的序列化开销：用collect_list over窗口排序，再用filter配合transform生成合法配对。

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