如何按Key在PySpark DataFrame中标记重叠时间范围？

我们需要给PySpark DataFrame添加一个valid列，标记同一key下的行是否存在时间范围（start_date至end_date）重叠的情况：若当前行与同组内其他行时间重叠，则标记为false，否则标记为true。

输入示例 #

原始DataFrame：

+---+-------------------+-------------------+
|key|start_date         |end_date           |
+---+-------------------+-------------------+
|A  |2022-01-11 00:00:00|8888-12-31 00:00:00|
|B  |2020-01-01 00:00:00|2022-02-10 00:00:00|
|B  |2019-02-08 00:00:00|2020-02-15 00:00:00|
|B  |2022-02-16 00:00:00|2022-12-15 00:00:00|
|C  |2018-01-01 00:00:00|2122-02-10 00:00:00|
+---+-------------------+-------------------+

期望输出 #

处理后的结果：

+---+-------------------+-------------------+-----+
|key|start_date         |end_date           |valid|
+---+-------------------+-------------------+-----+
|A  |2022-01-11 00:00:00|8888-12-31 00:00:00|true |
|B  |2020-01-01 00:00:00|2022-02-10 00:00:00|false|
|B  |2019-02-08 00:00:00|2020-02-15 00:00:00|false|
|B  |2022-02-16 00:00:00|2022-12-15 00:00:00|true |
|C  |2018-01-01 00:00:00|2122-02-10 00:00:00|true |
+---+-------------------+-------------------+-----+

解决方案代码 #

我们可以利用PySpark的窗口函数，通过对比同组内相邻行的时间范围来判断重叠：

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

# 初始化SparkSession
spark = SparkSession.builder.appName("OverlapDetection").getOrCreate()

# 创建并预处理示例数据
data = [
    ("A", "2022-01-11 00:00:00", "8888-12-31 00:00:00"),
    ("B", "2020-01-01 00:00:00", "2022-02-10 00:00:00"),
    ("B", "2019-02-08 00:00:00", "2020-02-15 00:00:00"),
    ("B", "2022-02-16 00:00:00", "2022-12-15 00:00:00"),
    ("C", "2018-01-01 00:00:00", "2122-02-10 00:00:00")
]

df = spark.createDataFrame(data, ["key", "start_date", "end_date"])
# 将字符串格式的时间转为Timestamp类型
df = df.withColumn("start_date", F.to_timestamp("start_date"))
df = df.withColumn("end_date", F.to_timestamp("end_date"))

# 定义窗口：按key分组，按start_date排序
window_spec = Window.partitionBy("key").orderBy("start_date")

# 获取相邻行的时间：上一行的end_date、下一行的start_date
df_with_neighbors = df.withColumn(
    "prev_end", F.lag("end_date").over(window_spec)
).withColumn(
    "next_start", F.lead("start_date").over(window_spec)
)

# 判断是否存在重叠并标记valid列
df_result = df_with_neighbors.withColumn(
    "valid",
    F.when(
        # 当前行与上一行重叠 或 当前行与下一行重叠
        (F.col("start_date") < F.col("prev_end")) | (F.col("end_date") > F.col("next_start")),
        F.lit(False)
    ).otherwise(F.lit(True))
).drop("prev_end", "next_start")

# 展示结果
df_result.show(truncate=False)

代码解释 #

1. 窗口定义：partitionBy("key")保证只在同一key分组内处理数据，orderBy("start_date")让组内行按开始时间排序，便于相邻行的时间对比。
2. 相邻行时间获取：lag("end_date")取上一行的结束时间，lead("start_date")取下一行的开始时间，这两个值是判断当前行是否重叠的关键。
3. 重叠判断逻辑：如果当前行的开始时间早于上一行的结束时间，说明和上一行重叠；如果当前行的结束时间晚于下一行的开始时间，说明和下一行重叠。只要满足任一条件，就标记为false，否则为true。

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