PySpark实现分组内连续行state 2累计时长计算

原代码的核心错误 #

1. 日期间隔计算错误：PySpark的datediff函数仅支持计算天数差，不接受'month'作为参数，计算月份差必须使用months_between函数。
2. 字段引用错误：原代码中使用了不存在的interval字段，实际应为report_interval。
3. 累计逻辑错误：仅通过lag函数无法实现连续累计，必须结合窗口累加函数完成累计计算。

正确实现步骤 #

1. 修正报告间隔计算

首先修正report_interval的计算逻辑，使用months_between得到准确的月份差，再转为以年为单位的时长：

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

# 定义按实体分区、按报告日期排序的窗口
window_entity = Window.partitionBy('entity_id').orderBy('report_date')

# 计算报告间隔：从上次报告到本次报告的时长（年）
df = df.withColumn(
    'report_interval',
    F.months_between(F.col('report_date'), F.lag(F.col('report_date')).over(window_entity)) / 12.0
).fillna(0, subset='report_interval')

2. 标记state 2的有效时长

report_interval代表的是上次报告到本次报告的时间间隔，这段时间内实体的状态以上次报告的state为准。因此我们需要标记出这段间隔是否属于state 2的有效时长：

# 标记当前间隔是否属于state 2的贡献：上次状态为2时，计入时长
df = df.withColumn(
    'state_2_duration',
    F.when(F.lag(F.col('state')).over(window_entity) == 2, F.col('report_interval')).otherwise(0.0)
)

3. 计算累计持续时长

使用累加窗口对每个实体的有效时长进行累计，得到每次报告时state 2的累计持续时长：

# 定义累加窗口：从当前实体的第一条记录到当前行
window_cumulative = Window.partitionBy('entity_id').orderBy('report_date').rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 计算累计时长
df = df.withColumn(
    'state_2_cumulative_duration',
    F.sum(F.col('state_2_duration')).over(window_cumulative)
)

逻辑说明 #

• 首次报告的report_interval为0（因为没有上一次报告），因此对应的state_2_duration也为0，累计时长从0开始。
• 后续每条记录的state_2_duration仅在上次报告状态为2时，才计入当前的report_interval，否则为0。
• 累加窗口会自动对每个实体的有效时长进行连续累计，最终得到每次报告时state 2的累计持续时长。

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