Scala DataFrame按ID分组，识别条件行的前后n行并标记

需求说明 #

给定Scala DataFrame，需按id字段分组，每组内按datetime字段排序。当cond列值为1时，为该条件行的前n行、当前行、后n行标记ans列：

• 条件行标记为at
• 条件行前n行标记为before
• 条件行后n行标记为after
• 若某行同时属于两个条件行的before和after范围（重叠场景），标记为between
• 不在任何范围内的行标记为0

普通场景示例（n=3） #

输入DataFrame #

val df = Seq(
  (101, "2020-01-01", 0),
  (101, "2020-02-01", 0),
  (101, "2020-03-01", 0),
  (101, "2020-04-01", 0),
  (220, "2020-01-01", 0),
  (220, "2020-01-05", 0),
  (220, "2020-01-10", 0),
  (220, "2020-01-15", 0),
  (220, "2020-02-01", 0),
  (220, "2020-02-05", 1),
  (220, "2020-02-10", 0),
  (220, "2020-02-15", 0),
  (220, "2020-02-20", 0),
  (220, "2020-02-25", 0)
).toDF("id", "datetime", "cond")

期望结果 #

val ans = Seq(
  (101, "2020-01-01", 0, "0"),
  (101, "2020-02-01", 0, "0"),
  (101, "2020-03-01", 0, "0"),
  (101, "2020-04-01", 0, "0"),
  (220, "2020-01-01", 0, "0"),
  (220, "2020-01-05", 0, "0"),
  (220, "2020-01-10", 0, "before"),
  (220, "2020-01-15", 0, "before"),
  (220, "2020-02-01", 0, "before"),
  (220, "2020-02-05", 1, "at"),
  (220, "2020-02-10", 0, "after"),
  (220, "2020-02-15", 0, "after"),
  (220, "2020-02-20", 0, "after"),
  (220, "2020-02-25", 0,  "0")
).toDF("id", "datetime", "cond", "ans")

边缘场景示例（n=4，范围重叠） #

输入DataFrame #

val df = Seq(
  (101, "2020-01-01", 0),
  (101, "2020-02-01", 0),
  (101, "2020-03-01", 0),
  (101, "2020-04-01", 0),
  (220, "2020-01-01", 0),
  (220, "2020-01-05", 0),
  (220, "2020-01-10", 0),
  (220, "2020-01-15", 0),
  (220, "2020-02-01", 0),
  (220, "2020-02-05", 1),
  (220, "2020-02-10", 0),
  (220, "2020-02-15", 1),
  (220, "2020-02-20", 0),
  (220, "2020-02-25", 0),
  (220, "2020-02-26", 0)
).toDF("id", "datetime", "cond")

期望结果 #

val edge = Seq(
  (101, "2020-01-01", 0, "0"),
  (101, "2020-02-01", 0, "0"),
  (101, "2020-03-01", 0, "0"),
  (101, "2020-04-01", 0, "0"),
  (220, "2020-01-01", 0, "0"),
  (220, "2020-01-05", 0, "before"),
  (220, "2020-01-10", 0, "before"),
  (220, "2020-01-15", 0, "before"),
  (220, "2020-02-01", 0, "before"),
  (220, "2020-02-05", 1, "at"),
  (220, "2020-02-10", 0, "between"),
  (220, "2020-02-15", 1, "at"),
  (220, "2020-02-20", 0, "after"),
  (220, "2020-02-25", 0, "after"),
  (220, "2020-02-26", 0, "after")
).toDF("id", "datetime", "cond", "ans")

实现方案 #

代码实现 #

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

// 定义参数n，根据场景修改
val n = 3 // 普通场景用3，边缘场景替换为4

// 1. 为每个id分组内的行按datetime排序并添加行号
val windowById = Window.partitionBy("id").orderBy("datetime")
val dfWithRowNum = df.withColumn("row_num", row_number().over(windowById))

// 2. 收集每个id内所有cond=1的行号，生成对应的前后范围区间
val condRowRanges = dfWithRowNum.filter(col("cond") === 1)
  .groupBy("id")
  .agg(collect_list("row_num").alias("cond_row_nums"))
  .withColumn("ranges", expr(
    "transform(cond_row_nums, x -> struct(x - n as start, x as at, x + n as end))"
  ))

// 3. 关联原数据和范围数据
val joinedDf = dfWithRowNum.join(condRowRanges, Seq("id"), "left_outer")

// 4. 定义UDF判断每行所属的标记类别
val determineAns = udf((rowNum: Int, ranges: Seq[Row]) => {
  if (ranges == null || ranges.isEmpty) "0"
  else {
    // 匹配当前行所在的所有范围类别
    val matchedTypes = ranges.map(r => {
      val start = r.getInt(0)
      val at = r.getInt(1)
      val end = r.getInt(2)
      rowNum match {
        case num if num == at => "at"
        case num if num >= start && num < at => "before"
        case num if num > at && num <= end => "after"
        case _ => null
      }
    }).filter(_ != null).distinct
    
    // 处理重叠情况：同时属于before和after则标记为between
    matchedTypes match {
      case Seq("before", "after") => "between"
      case Seq(singleType) => singleType
      case _ => "0"
    }
  }
})

// 5. 生成ans列并清理中间列
val result = joinedDf.withColumn("ans", determineAns(col("row_num"), col("ranges")))
  .drop("row_num", "cond_row_nums", "ranges")

// 查看结果
result.show()

逻辑说明 #

1. 行号标记：为每个id分组内的行按datetime排序后添加行号，用于快速定位行的相对位置。
2. 范围生成：收集每个分组中cond=1的行号，计算出每个条件行对应的前n行、当前行、后n行的区间范围。
3. 关联判断：通过UDF检查每行是否落在某个条件行的范围内，同时处理重叠场景（当一行同时属于两个条件行的before和after范围时，标记为between）。
4. 清理输出：移除中间生成的辅助列，得到最终结果。

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