PySpark实现：计算距上次signal为true的时间差

原始DataFrame（实际包含多个id，此处仅展示单个） #

+---+-------------------+------+
| id|             timstm|signal|
+---+-------------------+------+
| X1|2022-07-01 00:00:00|  null|
| X1|2022-07-02 00:00:00|  true|
| X1|2022-07-03 00:00:00|  null|
| X1|2022-07-05 00:00:00|  null|
| X1|2022-07-09 00:00:00|  true|
+---+-------------------+------+

期望输出 #

新增time_diff列，记录当前行与上次signal为true的时间间隔（单位：天）：

+---+-------------------+------+---------+
| id|             timstm|signal|time_diff|
+---+-------------------+------+---------+
| X1|2022-07-01 00:00:00|  null|     null|
| X1|2022-07-02 00:00:00|  true|      0.0|
| X1|2022-07-03 00:00:00|  null|      1.0|
| X1|2022-07-05 00:00:00|  null|      3.0|
| X1|2022-07-09 00:00:00|  true|      0.0|
+---+-------------------+------+---------+

解决思路与实现代码 #

核心思路是利用窗口函数追踪每个id内最近一次signal为true的时间戳，再计算时间差。具体步骤如下：

1. 定义窗口：按id分区，按timstm升序排序，窗口范围包含当前行及之前所有行，确保只回溯历史数据。
2. 标记有效时间戳：仅保留signal为true时的timstm，其余设为null。
3. 提取最近有效时间戳：用last函数（忽略null）获取当前行之前最近的有效时间戳。
4. 计算时间差：针对不同场景处理结果，当前行signal为true时设为0.0，无历史有效时间戳时设为null，其余情况计算天数差并转成浮点数。

以下是Python版本的实现代码：

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

# 定义窗口规则
window_spec = Window.partitionBy("id").orderBy("timstm").rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 1. 标记signal为true的时间戳，其余为null
df = df.withColumn("last_true_timestamp", F.when(F.col("signal") == True, F.col("timstm")))

# 2. 获取当前行之前最近的signal为true的时间戳
df = df.withColumn("last_true_timestamp", F.last("last_true_timestamp", ignoreNulls=True).over(window_spec))

# 3. 计算time_diff列
df = df.withColumn(
    "time_diff",
    F.when(F.col("signal") == True, 0.0)
    .when(F.col("last_true_timestamp").isNull(), None)
    .otherwise(F.datediff(F.col("timstm"), F.col("last_true_timestamp")).cast("double"))
)

# 查看结果
df.show()

关键说明 #

• rowsBetween(Window.unboundedPreceding, Window.currentRow)：限制窗口范围为当前行及之前的所有数据，避免提前引用未来的signal=true记录。
• last(..., ignoreNulls=True)：自动跳过null值，直接取最近的有效时间戳，无需额外过滤。
• datediff函数直接计算两个日期的天数差，转成double类型与示例输出格式一致。

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