合并pyspark中的重叠日期范围答案

【问题标题】：Merge overlapping date ranges in pyspark合并pyspark中的重叠日期范围
【发布时间】：2021-09-28 14:18:23
【问题描述】：

我有一个带有架构（ID、start_date、end_date）的数据框

示例数据框：

id	start_date	end_date
1	03/05/2020	07/05/2020
1	05/05/2020	07/05/2020
1	01/05/2020	07/05/2020
1	01/05/2020	13/05/2020
2	06/05/2020	09/05/2020
1	18/05/2020	20/05/2020
2	09/05/2020	17/05/2020

我想要以下格式的结果

id	start_dates	days
1	[03/05/2020, 05/05/2020, 01/05/2020, 01/05/2020, 18/05/2020]	16
2	[06/05/2020, 09/05/2020]	12

days 是非重叠日期间隔的总和 (start_date - end_date)。

我尝试的方法是按 id 分组，将 start_date 和 end_date 收集为列表，然后编写一个 UDF 来获取天数列。我想知道是否有更好的方法来实现这一点。

【问题讨论】：

如何计算 id 1 的 16 天？
@koiralo 13 天（01/05/2020 至 13/05/2020）+ 3 天（18/05/2020 至 20/05/2020）
@koiralo：希望您的问题得到解答。如果没有，请告诉我。

标签： python apache-spark pyspark

【解决方案1】：

这是一个可能的解决方案，它生成 sequence 并在 lag 函数的帮助下找到 array_overlaps。然后分隔 id 列中的每个组并对值求和。然后再次相加得到最终结果：

from pyspark.sql import functions as F, Window

w = Window.partitionBy("id").orderBy(F.monotonically_increasing_id())

out = (df.withColumn("start_date",F.to_date("start_date","dd/MM/yyyy"))
       .withColumn("end_date",F.to_date("end_date","dd/MM/yyyy"))
       .withColumn("Seq",F.sequence("start_date","end_date"))
       .withColumn("Lag",F.arrays_overlap("Seq",F.lag("Seq").over(w)))
       .withColumn("Flag",F.when(F.col("Lag")|F.col("Lag").isNull(),1).otherwise(0))
        
        .groupBy("Id","Flag").agg(F.collect_list("start_date").alias("start_date"),
                (F.datediff(F.max("end_date"),F.min("start_date"))+1).alias("Days"))
        .groupBy("id").agg(F.flatten(F.collect_list("start_date")).alias("start_date"),
                           F.sum("Days").alias("Days"))
       )

out.show(truncate=False)

+---+------------------------------------------------------------+----+
|id |start_date                                                  |Days|
+---+------------------------------------------------------------+----+
|1  |[2020-05-03, 2020-05-05, 2020-05-01, 2020-05-01, 2020-05-18]|16  |
|2  |[2020-05-06, 2020-05-09]                                    |12  |
+---+------------------------------------------------------------+----+

【讨论】：