如何在Spark DataFrame中合并连续/重叠的时间间隔？

输入数据 #

需求说明 #

• 对相同ID的记录，合并连续或重叠的时间范围：
  • ID=1的三条记录时间范围连续/重叠，合并为一条全局时间范围
  • ID=2的前两条记录连续，合并为一条；第三条与前一条间隔一年（2019年缺失），保留独立记录

期望输出 #

实现步骤与代码 #

核心思路 #

1. 将字符串格式的日期转换为Spark日期类型，支持时间运算
2. 按ID分组、按开始日期排序，用窗口函数识别连续/重叠的时间区间
3. 为连续/重叠的区间生成统一分组标识，最后按分组聚合得到合并后的时间范围

Python 实现代码 #

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

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

# 加载输入数据
data = [
    (1, "01/01/1993", "31/12/2018"),
    (1, "01/06/2015", "31/12/2023"),
    (1, "01/01/2024", "31/12/2024"),
    (2, "01/01/1993", "31/12/2017"),
    (2, "01/01/2018", "31/12/2018"),
    (2, "01/01/2020", "31/12/2023")
]
df = spark.createDataFrame(data, ["ID", "start_date", "end_date"])

# 1. 转换字符串日期为Spark日期类型（dd/MM/yyyy格式）
df = df.withColumn("start_dt", F.to_date("start_date", "dd/MM/yyyy")) \
       .withColumn("end_dt", F.to_date("end_date", "dd/MM/yyyy"))

# 2. 定义窗口：按ID分组，按开始日期升序排列
window = Window.partitionBy("ID").orderBy("start_dt")

# 3. 生成连续分组标识
# - 取上一条记录的结束日期，判断当前记录是否与上一条连续/重叠
# - 累加标识值，为连续区间生成统一group_id
df = df.withColumn("prev_end", F.lag("end_dt").over(window)) \
       .withColumn("is_new_group", F.when(F.col("start_dt") <= F.date_add(F.col("prev_end"), 1), 0).otherwise(1)) \
       .withColumn("group_id", F.sum("is_new_group").over(window.rowsBetween(Window.unboundedPreceding, Window.currentRow)))

# 4. 按ID和group_id聚合，合并时间范围并转回原日期格式
result_df = df.groupBy("ID", "group_id") \
              .agg(F.min("start_dt").alias("start_dt"), F.max("end_dt").alias("end_dt")) \
              .withColumn("start_date", F.date_format("start_dt", "dd/MM/yyyy")) \
              .withColumn("end_date", F.date_format("end_dt", "dd/MM/yyyy")) \
              .select("ID", "start_date", "end_date") \
              .orderBy("ID", "start_date")

# 查看结果
result_df.show()

Scala 实现代码 #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window

object MergeTimeRanges {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder.appName("MergeTimeRanges").getOrCreate()
    import spark.implicits._

    // 加载输入数据
    val data = Seq(
      (1, "01/01/1993", "31/12/2018"),
      (1, "01/06/2015", "31/12/2023"),
      (1, "01/01/2024", "31/12/2024"),
      (2, "01/01/1993", "31/12/2017"),
      (2, "01/01/2018", "31/12/2018"),
      (2, "01/01/2020", "31/12/2023")
    ).toDF("ID", "start_date", "end_date")

    // 1. 转换字符串日期为Spark日期类型
    val df = data.withColumn("start_dt", to_date(col("start_date"), "dd/MM/yyyy"))
                 .withColumn("end_dt", to_date(col("end_date"), "dd/MM/yyyy"))

    // 2. 定义窗口
    val window = Window.partitionBy("ID").orderBy("start_dt")

    // 3. 生成连续分组标识
    val groupedDF = df.withColumn("prev_end", lag(col("end_dt"), 1).over(window))
                      .withColumn("is_new_group", when(col("start_dt") <= date_add(col("prev_end"), 1), 0).otherwise(1))
                      .withColumn("group_id", sum(col("is_new_group")).over(window.rowsBetween(Window.unboundedPreceding, Window.currentRow)))

    // 4. 聚合并转换回原日期格式
    val resultDF = groupedDF.groupBy("ID", "group_id")
                            .agg(min(col("start_dt")).alias("start_dt"), max(col("end_dt")).alias("end_dt"))
                            .withColumn("start_date", date_format(col("start_dt"), "dd/MM/yyyy"))
                            .withColumn("end_date", date_format(col("end_dt"), "dd/MM/yyyy"))
                            .select("ID", "start_date", "end_date")
                            .orderBy("ID", "start_date")

    // 查看结果
    resultDF.show()
  }
}

关键逻辑说明 #

• 日期转换：必须将字符串转为DateType才能正确进行时间比较和运算，避免字符串排序的错误
• 连续判断：用date_add(prev_end, 1)处理"次日连续"的场景（比如上一条结束于2017-12-31，下一条开始于2018-01-01，视为连续）
• 分组标识：通过累加is_new_group的值，将连续/重叠的区间归为同一个group_id，确保聚合时能合并成一条记录

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