PySpark如何基于前序计算值迭代DataFrame实现状态时长计算

原始数据 #

需求说明 #

按ID分组实现以下计算：

• 当前行状态与前一行相同时，沿用前序计算的起始日期
• 状态变更时，更新起始日期为当前行的date
• 计算当前date与起始日期的时间差（单位：分钟）

尝试的代码 #

用户尝试用窗口函数lag操作，但由于lag取原始数据而非计算后的值，无法满足需求：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

partition = Window.partitionBy('id').orderBy(F.col('date'))

df = (df
      .withColumn("lag_sts", F.lag("status", 1).over(partition))
      .withColumn("lag_date", F.lag("date", 1).over(partition))
      .withColumn("Expected", F.when(F.col("status") == F.col("lag_sts"), F.col("lag_date"))
                  .otherwise(F.col("date"))
      )
)

当前运行结果 #

解决方案 #

核心思路是先对连续相同状态的行分组，再取每组的最早日期作为起始日期，最后计算时长。

示例代码 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 定义窗口：按ID分组，按日期排序
w = Window.partitionBy("id").orderBy("date")

# 2. 标记状态变更行：状态与前一行不同则为1，第一行标记为1
df = df.withColumn(
    "change_flag",
    F.when(F.lag("status").over(w) != F.col("status"), 1)
     .when(F.lag("status").over(w).isNull(), 1)
     .otherwise(0)
)

# 3. 累加变更标记，生成连续状态组的ID
df = df.withColumn(
    "group_id",
    F.sum("change_flag").over(w.rowsBetween(Window.unboundedPreceding, Window.currentRow))
)

# 4. 按组取最小日期作为Expected起始日期
group_window = Window.partitionBy("id", "group_id")
df = df.withColumn(
    "Expected",
    F.min("date").over(group_window)
)

# 5. 计算状态持续时长（分钟）
df = df.withColumn(
    "Expected Status Duration in Min",
    (F.unix_timestamp("date") - F.unix_timestamp("Expected")) / 60
)

# 可选：删除中间辅助列
df = df.drop("change_flag", "group_id")

# 查看结果
df.show(truncate=False)

结果说明 #

运行后将生成符合预期的结果：

• 连续相同状态的行共用同一个Expected起始日期
• 状态变更时，Expected自动更新为当前行的日期
• 时长准确计算为当前日期与起始日期的分钟差

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