Spark Scala中使用lag函数生成活跃标记连续链的起始周

我有一个包含week、person、person_active_flag列的DataFrame，person列包含多个不同用户。需求是新增chain_beginning列：当person_active_flag为1时，填充该连续1状态链的起始week；为0时填充0。

期望结果 #

val df_beginning = Seq(("2022-06-06", "person1", 1, "2022-06-06"),
             ("2022-06-13", "person1", 1, "2022-06-06"),
             ("2022-06-20", "person1", 1, "2022-06-06"),
             ("2022-06-27", "person1", 0, "0"),
             ("2022-07-04", "person1", 0, "0"),
             ("2022-07-11", "person1", 1, "2022-07-11"),
             ("2022-07-18", "person1", 1, "2022-07-11"),
             ("2022-07-25", "person1", 0, "0"),
             ("2022-08-01", "person1", 0, "0"),
             ("2022-08-08", "person1", 1, "2022-08-08"),
             ("2022-08-15", "person1", 1, "2022-08-08"),
             ("2022-08-22", "person1", 1, "2022-08-08"),
             ("2022-08-29", "person1", 1, "2022-08-08"))
.toDF("week", "person", "person_active_flag", "chain_beginning")
.orderBy($"week")

尝试的代码（未得到正确结果） #

val w = Window.partitionBy($"person").orderBy($"week".asc)

df_beginning
.withColumn("beginning_chain", 
    when($"person_active_flag" === 1 && (lag($"person_active_flag", 1).over(w) === 0 || lag($"person_active_flag", 1).over(w).isNull), 1).otherwise(0)
)

.withColumn("first_week", when($"beginning_chain" === 1, $"week"))

.withColumn("beginning_chain_week", 
    when($"person_active_flag" === 1 && lag($"person_active_flag", 1).over(w).isNull, $"first_week")
   .when($"person_active_flag" === 1 && lag($"person_active_flag", 1).over(w) === 0, $"first_week")
   .when($"person_active_flag" === 1 && lag($"person_active_flag", 1).over(w) === 1, lag($"first_week", 1).over(w))
   .otherwise(0)
)

正确实现方法 #

核心思路是先为每个连续的活跃状态链分配唯一分组ID，再在每个分组内取最小的week作为链的起始值。具体步骤如下：

1. 定义基础窗口：按person分区，按week升序排序，用于状态变化的判断。
2. 标记活跃链的起始点：当当前行person_active_flag为1，且前一行状态不是1（或为第一行）时，标记为1，否则为0。
3. 计算累计和生成分组ID：对起始点标记做累计求和，这样同一连续活跃链的所有行会拥有相同的分组ID，非活跃行的分组ID为0。
4. 定义分组窗口：按person和分组ID分区，用于提取每个活跃链的起始周。
5. 生成最终的chain_beginning列：活跃状态下取分组内最小的week，非活跃状态填0。

完整代码：

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

// 1. 基础窗口：按用户分区，按周排序
val baseWindow = Window.partitionBy($"person").orderBy($"week".asc)

// 2. 标记活跃链的起始点
val dfWithStartFlag = df_beginning
  .withColumn("is_chain_start", 
    when(
      $"person_active_flag" === 1 && 
      (lag($"person_active_flag", 1).over(baseWindow) =!= 1 || lag($"person_active_flag", 1).over(baseWindow).isNull),
      1
    ).otherwise(0)
  )

// 3. 生成分组ID：同一连续活跃链的分组ID相同
val dfWithGroupId = dfWithStartFlag
  .withColumn("chain_group_id", sum($"is_chain_start").over(baseWindow.rangeBetween(Window.unboundedPreceding, 0)))

// 4. 分组窗口：按用户和分组ID分区，提取每个链的起始周
val groupWindow = Window.partitionBy($"person", $"chain_group_id").orderBy($"week".asc)

// 5. 生成最终的chain_beginning列
val resultDf = dfWithGroupId
  .withColumn("chain_beginning", 
    when(
      $"person_active_flag" === 1,
      first($"week").over(groupWindow)
    ).otherwise("0")
  )
  .drop("is_chain_start", "chain_group_id") // 移除中间列
  .orderBy($"week")

resultDf.show()

运行上述代码后，即可得到符合预期的结果。

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