Spark技术咨询：如何为每个EventId获取IncDate晚于EventDate的3个Incident

嘿，我来帮你解决这个Spark的问题！你说要给每个eventId找出满足eventDate < IncDate条件的后续3个incId，之前用窗口函数没得到正确结果，大概率是窗口的范围、过滤逻辑或者排序方向没处理对，咱们一步步来搞定它。

首先咱们先构建一个贴近需求的示例数据，方便后续演示。假设我们有两个数据集：

• events表：存储每个事件的ID和发生日期
• incs表：存储每个关联记录的ID、所属事件ID以及发生日期

// 示例Events数据
val eventsDF = spark.createDataFrame(Seq(
  ("E1", "2023-01-01"),
  ("E2", "2023-02-15")
)).toDF("eventId", "eventDate")

// 示例Incidents数据
val incsDF = spark.createDataFrame(Seq(
  ("I1", "E1", "2023-01-02"),
  ("I2", "E1", "2023-01-05"),
  ("I3", "E1", "2023-01-08"),
  ("I4", "E1", "2023-01-10"),
  ("I5", "E2", "2023-02-16"),
  ("I6", "E2", "2023-02-20"),
  ("I7", "E2", "2023-02-25")
)).toDF("incId", "eventId", "incDate")

预期输出应该是每个eventId对应符合日期条件的前3个incId，比如E1对应[I1,I2,I3]，E2对应[I5,I6,I7]。

核心思路是：先关联数据并过滤出符合eventDate < incDate的记录，再用窗口函数按事件分区、按inc日期排序，最后取前3条。

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

// 1. 关联两张表，过滤出日期符合条件的记录
val joinedDF = eventsDF.join(incsDF, Seq("eventId"), "inner")
  .where(col("incDate") > col("eventDate"))

// 2. 定义窗口：按eventId分区，按incDate升序排序（保证是"后续"的顺序）
val windowSpec = Window.partitionBy("eventId").orderBy("incDate")

// 3. 添加行号，筛选前3条，最后将同一event的incId聚合为数组
val resultDF = joinedDF
  .withColumn("row_num", row_number().over(windowSpec))
  .where(col("row_num") <= 3)
  .groupBy("eventId", "eventDate")
  .agg(collect_list("incId").alias("top_3_incIds"))

// 查看结果
resultDF.show(false)

执行后你会得到预期的输出：

+-------+----------+---------------+
|eventId|eventDate |top_3_incIds   |
+-------+----------+---------------+
|E1     |2023-01-01|[I1, I2, I3]   |
|E2     |2023-02-15|[I5, I6, I7]   |
+-------+----------+---------------+

如果你习惯用SQL来实现，逻辑是完全一致的：

WITH joined_data AS (
  -- 关联并过滤符合日期条件的记录
  SELECT e.eventId, e.eventDate, i.incId, i.incDate
  FROM events e
  INNER JOIN incs i ON e.eventId = i.eventId
  WHERE i.incDate > e.eventDate
),
ranked_data AS (
  -- 按eventId分区，给符合条件的inc按日期排序并加行号
  SELECT *,
         ROW_NUMBER() OVER (PARTITION BY eventId ORDER BY incDate ASC) AS row_num
  FROM joined_data
)
-- 筛选前3个，聚合为数组
SELECT eventId, eventDate, COLLECT_LIST(incId) AS top_3_incIds
FROM ranked_data
WHERE row_num <= 3
GROUP BY eventId, eventDate;

你之前用窗口函数没得到正确结果，可能是这几个原因：

• 没提前过滤不符合日期条件的记录：窗口函数会包含所有分区内的记录，包括incDate <= eventDate的，导致行号计算错误。
• 排序方向错了：如果用了ORDER BY incDate DESC，会取最新的3个而不是事件发生后的前3个。
• 用了rank()而非row_number()：如果存在相同incDate的记录，rank()会生成重复的行号，可能导致最终结果超过3条；而row_number()会给每条记录唯一行号，更符合"取3个"的需求。
• 窗口范围设置错误：比如误加了rows between ...的范围，导致窗口包含了不需要的记录。

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