使用 Apache Spark SQL 查找间隔最长且无事故的设施

Question

我有下一个数据集：

|facility|date      |accidents|
| foo    |2019-01-01|1        |
| foo    |2019-01-02|null     |
| foo    |2019-01-03|null     |
| foo    |2019-01-04|2        |
| bar    |2019-01-01|1        |
| bar    |2019-01-02|null     |
| bar    |2019-01-03|3        |

目标是找到一个无事故连续时间最长的设施：

|facility|startDate |interval|
|foo     |2019-01-02|2       |

是否可以使用 Spark SQL 做到这一点？谢谢

PS代码示例：

case class FacilityRecord(name: String, date: java.sql.Date, accidents: Option[Int])
case class IntervalWithoutAccidents(name: String, startDate: java.sql.Date, interval: Int)

implicit val spark: SparkSession = SparkSession.builder
      .appName("Test")
      .master("local")
      .getOrCreate()

import spark.implicits._

val facilityRecords = Seq(
  FacilityRecord("foo", Date.valueOf("2019-01-01"), Some(1)),
  FacilityRecord("foo", Date.valueOf("2019-01-02"), None),
  FacilityRecord("foo", Date.valueOf("2019-01-03"), None),
  FacilityRecord("foo", Date.valueOf("2019-01-04"), Some(2)),
  FacilityRecord("bar", Date.valueOf("2019-01-01"), Some(1)),
  FacilityRecord("bar", Date.valueOf("2019-01-02"), None),
  FacilityRecord("bar", Date.valueOf("2019-01-03"), Some(3))
)

val facilityRecordsDataset = spark.createDataset(facilityRecords)

facilityRecordsDataset.show()

val intervalWithoutAccidents: IntervalWithoutAccidents = ??? // TODO: find the interval
val expectedInterval = IntervalWithoutAccidents("foo", startDate = Date.valueOf("2019-01-02"), interval = 2)
assert(expectedInterval == intervalWithoutAccidents)

println(intervalWithoutAccidents)

Answer 1

这是一个 2 步方法：

1. 使用 Window 函数为当前日期和下一个事故日期之间的每一行中的accident_date每个facility计算interval值创建列和first。
2. 使用 Window 函数计算max intervalper并过滤具有最大间隔值的行。facilitymax

示例代码如下：

import java.sql.Date
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window
import spark.implicits._

val df = Seq(
  ("foo", Date.valueOf("2019-01-01"), Some(1)),
  ("foo", Date.valueOf("2019-01-02"), None),
  ("foo", Date.valueOf("2019-01-03"), None),
  ("foo", Date.valueOf("2019-01-04"), Some(2)),
  ("bar", Date.valueOf("2019-01-01"), Some(1)),
  ("bar", Date.valueOf("2019-01-02"), None),
  ("bar", Date.valueOf("2019-01-03"), Some(3))
).toDF("facility", "date", "accidents")

val win = Window.partitionBy($"facility").orderBy($"date").
  rowsBetween(0, Window.unboundedFollowing)

第 1 步：计算 interval

val df2 = df.
  withColumn("accident_date", when($"accidents".isNotNull, $"date")).
  withColumn("interval",
    datediff(first($"accident_date", ignoreNulls=true).over(win), $"date")
  )

df2.show
// +--------+----------+---------+-------------+--------+
// |facility|      date|accidents|accident_date|interval|
// +--------+----------+---------+-------------+--------+
// |     bar|2019-01-01|        1|   2019-01-01|       0|
// |     bar|2019-01-02|     null|         null|       1|
// |     bar|2019-01-03|        3|   2019-01-03|       0|
// |     foo|2019-01-01|        1|   2019-01-01|       0|
// |     foo|2019-01-02|     null|         null|       2|
// |     foo|2019-01-03|     null|         null|       1|
// |     foo|2019-01-04|        2|   2019-01-04|       0|
// +--------+----------+---------+-------------+--------+

第 2 步：计算 max interval

df2.select($"facility", $"date".as("start_date"),
    max($"interval").over(Window.partitionBy($"facility")).as("max_interval")
  ).
  where($"interval" === $"max_interval").
  show
// +--------+----------+------------+
// |facility|start_date|max_interval|
// +--------+----------+------------+
// |     bar|2019-01-02|           1|
// |     foo|2019-01-02|           2|
// +--------+----------+------------+