Spark中的区间联合答案

【问题标题】：Union of intervals in SparkSpark中的区间联合
【发布时间】：2021-11-01 03:19:58
【问题描述】：

我在 Spark 中有以下数据框。

id	start	end
140396	2002-06-18	2003-06-18
140396	2007-07-29	2015-07-29
140396	2008-02-05	2010-02-05
140396	2009-01-18	2010-01-18
140396	2013-01-19	2021-08-30
140396	2017-05-15	2021-08-30

我必须分析日期范围以获取其他日期范围，它们之间没有交集，但保留完整的日期范围。结果：

id	start	end
140396	2002-06-18	2003-06-18
140396	2007-07-29	2021-08-30

其他例子可以来自：

id	start	end
140396	2002-06-18	2003-06-18
140396	2007-07-29	2015-07-29
140396	2014-02-05	2016-02-05
140396	2017-05-15	2021-08-30

到

id	start	end
140396	2002-06-18	2003-06-18
140396	2007-07-29	2016-02-05
140396	2017-05-15	2021-08-30

请记住，还有其他用户有他们的日期，所以问题是按 id 用一个窗口划分的。

有人知道如何解决这个问题吗？

非常感谢您

【问题讨论】：

标签： scala apache-spark apache-spark-sql

【解决方案1】：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window
    
val df = Seq(
(140396, "2002-06-18", "2003-06-18"),
(140396, "2007-07-29", "2015-07-29"),
(140396, "2008-02-05", "2010-02-05"),
(140396, "2009-01-18", "2010-01-18"),
(140396, "2013-01-19", "2021-08-30"),
(140396, "2017-05-15", "2021-08-30"),
(140397, "2002-06-18", "2003-06-18"),
(140397, "2007-07-29", "2015-07-29"),
(140397, "2014-02-05", "2016-02-05"),
(140397, "2017-05-15", "2021-08-30")
) toDF ("id", "start", "end")

val windowSpec1 = Window.partitionBy(col("id"))
                        .orderBy(col("start"), col("end"))
                        .rowsBetween(Window.unboundedPreceding, Window.currentRow - 1)

val windowSpec2 = Window.partitionBy(col("id"))
                        .orderBy(col("start"), col("end"))
                        .rowsBetween(Window.unboundedPreceding, Window.currentRow)

val result = df.withColumn("left_edge", max(when(col("start") < max(col("end")).over(windowSpec1), null).otherwise(col("start"))).over(windowSpec2))
               .groupBy(col("id"), col("left_edge"))
               .agg(min(col("start")).alias("start"), max(col("end")).alias("end"))
               .select("id", "start", "end")
               .orderBy("id", "start")

display(result)

结果：

id	start	end
140396	2002-06-18	2003-06-18
140396	2007-07-29	2021-08-30
140397	2002-06-18	2003-06-18
140397	2007-07-29	2016-02-05
140397	2017-05-15	2021-08-30

参考： https://wiki.postgresql.org/wiki/Range_aggregation

【讨论】：

【解决方案2】：

您可以通过 id 按数据框分组，将结果转换为 KeyValueGroupedDataset，然后执行 merge overlapping intervals 算法

您首先需要定义case class 代表数据框的一行：

case class Line(id: Int, start: String, end: String)

然后在主体部分使用：

import sparkSession.implicits._

dataframe
  .groupBy("id")
  .as[Int, Line]
  .flatMapGroups((id, grouped) => grouped.toSeq.sortBy(_.start).foldLeft(Seq.empty[Line])((acc, line) => (acc, line) match {
    case (Nil, line) => Seq(line)
    case (x::xs, line) if x.end >= line.end => x::xs
    case (x::xs, line) if x.end < line.end && x.end >= line.start => Line(id, x.start, line.end) +: xs
    case (xs, line) => line +: xs
  }))
  .orderBy("id", "start")

【讨论】：