基于列条件筛选最新记录的Spark代码优化问题

嘿，我来帮你搞定这个Spark DataFrame的分组筛选问题！你遇到的核心问题是原代码的Window分区逻辑有误，导致同分组下同时间戳的记录被错误过滤掉了，咱们一步步来修正。

需求明确 #

咱们先再理一遍需求，确保对齐：

• 按OrganizationId和SegmentId作为分组键
• 每组内的处理规则：
  • 如果该组只有唯一的TimeStamp：保留该组的所有记录
  • 如果该组有多个不同的TimeStamp：仅保留最新时间戳对应的所有记录

原代码问题诊断 #

你原来的代码是这样的：

val windowSpec3 = Window.partitionBy("OrganizationId", "SegmentId", "TimeStamp").orderBy(unix_timestamp($"TimeStamp", "yyyy-MM-dd'T'HH:mm:ss").cast("timestamp").desc) 
val latestForEachKey = latestForEachKey2.withColumn("rank", row_number().over(windowSpec3)).filter($"rank" === 1).drop("rank")

这里的问题在于：你把TimeStamp也加入了分区键，这会把同一个(OrganizationId, SegmentId)下的不同时间戳拆成多个子分区，然后用row_number()给每个子分区内的记录排名，只保留每条的第一条——这就直接导致同时间戳的其他记录被过滤掉了，完全不符合咱们的需求。

修正后的代码实现 #

咱们需要换个思路：先找出每个(OrganizationId, SegmentId)分组内的最大时间戳，然后筛选出该组内时间戳等于这个最大值的所有记录。这里提供两种可行的实现方式：

方式一：使用Window函数计算组内最大时间戳 #

// 定义窗口：仅按OrganizationId和SegmentId分区，计算每组的最大TimeStamp
val windowSpec = Window.partitionBy("OrganizationId", "SegmentId")

// 给每条记录添加组内最大时间戳字段
val dfWithMaxTs = yourOriginalDf.withColumn(
  "max_group_timestamp",
  max(unix_timestamp($"TimeStamp", "yyyy-MM-dd'T'HH:mm:ss").cast("timestamp")).over(windowSpec)
)

// 筛选出TimeStamp等于组内最大时间戳的所有记录，然后移除辅助字段
val latestRecords = dfWithMaxTs
  .filter(unix_timestamp($"TimeStamp", "yyyy-MM-dd'T'HH:mm:ss").cast("timestamp") === $"max_group_timestamp")
  .drop("max_group_timestamp")

方式二：分组聚合后关联原表 #

如果你的数据量较大，这种方式可能会有更好的性能：

// 先聚合得到每个(OrganizationId, SegmentId)对应的最新时间戳
val maxTimestampPerGroup = yourOriginalDf
  .groupBy("OrganizationId", "SegmentId")
  .agg(
    max(unix_timestamp($"TimeStamp", "yyyy-MM-dd'T'HH:mm:ss").cast("timestamp")).alias("max_group_timestamp")
  )

// 关联原表，筛选出符合条件的记录
val latestRecords = yourOriginalDf
  .join(maxTimestampPerGroup, Seq("OrganizationId", "SegmentId"), "inner")
  .filter(unix_timestamp($"TimeStamp", "yyyy-MM-dd'T'HH:mm:ss").cast("timestamp") === $"max_group_timestamp")
  .drop("max_group_timestamp")

结果验证 #

拿你给出的示例来验证：

• SegmentId=27：原数据有3条记录，其中2条是最新时间戳2018-05-29T09:17:18+00:00，1条是更早的时间戳。修正后的代码会保留那2条最新时间戳的记录，完全符合预期。
• SegmentId=6：两条记录都是同一个最新时间戳，代码会全部保留，不会丢失任何一条。

最终生成的DataFrame会完全匹配你给出的预期输出结果。

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