PySpark按ID分组统计status为0的连续累计天数实现方法

问题说明 #

现有存储客户每日状态的PySpark DataFrame，结构如下：

+---+------+-----------+
|id |status|date       |
+---+------+-----------+
|1  |1     |01-01-2022 |
|1  |0     |02-01-2022 |
|1  |0     |03-01-2022 |
|1  |0     |04-01-2022 |
|1  |1     |05-01-2022 |
|1  |0     |06-01-2022 |
|2  |1     |01-01-2022 |
|2  |0     |02-01-2022 |
|2  |0     |03-01-2022 |
|2  |1     |04-01-2022 |
|2  |0     |05-01-2022 |
+---+------+-----------+

需要统计每个客户连续处于status=0状态的天数，期望输出如下：

+---+------+-----------+------------+
|id |status|date       |count_status|
+---+------+-----------+------------+
|1  |1     |01-01-2022 | 0          |
|1  |0     |02-01-2022 | 1          |
|1  |0     |03-01-2022 | 2          |
|1  |0     |04-01-2022 | 3          |
|1  |1     |05-01-2022 | 0          |
|1  |0     |06-01-2022 | 1          |
|2  |1     |01-01-2022 | 0          |
|2  |0     |02-01-2022 | 1          |
|2  |0     |03-01-2022 | 2          |
|2  |1     |04-01-2022 | 0          |
|2  |0     |05-01-2022 | 1          |
+---+------+-----------+-----------+

pandas环境下已通过以下逻辑实现，需要改写为PySpark版本：

df['count_status'] = np.where(df['status'] == 0, 
                             df.groupby(['id', 
                                        (df['status'] != df['status'].shift(1)).cumsum()]).cumcount()+1,
                                  0)

实现代码 #

核心逻辑和pandas实现完全对齐，通过窗口函数实现偏移取值、累计求和、分组计数三个步骤：

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

# 1. 定义基础窗口：按客户ID分区，按日期升序排列
base_win = Window.partitionBy("id").orderBy("date")

# 2. 标记状态变更点：取上一行status，和当前行不一致则标记为1
df = df.withColumn("prev_status", F.lag("status").over(base_win))
df = df.withColumn(
    "is_change",
    F.when(
        F.col("prev_status").isNull() | (F.col("status") != F.col("prev_status")),
        1
    ).otherwise(0)
)

# 3. 对变更标记累计求和，生成每个连续状态段的唯一ID
df = df.withColumn("segment_id", F.sum("is_change").over(base_win))

# 4. 按客户+连续状态段分组，计算段内序号，status为0时取序号否则为0
segment_win = Window.partitionBy("id", "segment_id").orderBy("date")
df = df.withColumn(
    "count_status",
    F.when(F.col("status") == 0, F.row_number().over(segment_win)).otherwise(0)
)

# 5. 删除中间计算列，得到最终结果
df = df.drop("prev_status", "is_change", "segment_id")

执行后输出和期望结果完全一致。

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