Spark DataFrames Scala - jump to next group during a loop

user_8310898 0
Viewed 115

I have a dataframe as below and it records the quarter and date during which same incident occurs to which IDs.

enter image description here

I would like to mark an ID and date if the incident happen at two consecutive quarters. And this is how I did it.

val arrArray = dtf.collect.map(x => (x(0).toString, x(1).toString, x(2).toString.toInt))

if (arrArray.length > 0) {

  var bufBAQDate = ArrayBuffer[Any]()

  for (a <- 1 to arrArray.length - 1) {
    val (strBAQ1, strDate1, douTime1) = arrArray(a - 1)
    val (strBAQ2, strDate2, douTime2) = arrArray(a)

    if (douTime2 - douTime1 == 15 && strBAQ1 == strBAQ2 && strDate1 == strDate2) {
      bufBAQDate = (strBAQ2, strDate2) +: bufBAQDate
      //println(strBAQ2+" "+strDate2+" "+douTime2)
    }
  }

  val vecBAQDate = bufBAQDate.distinct.toVector.reverse

Is there a better way of doing it? As the same insident can happen many times to one ID during a single day, it is better to jump to the next ID and/or date once an ID and a date is marked. I dont want to create nested loops to filter dataframe.

apache-spark grouping scala loops
1 Answers

Note that you current solution misses 20210103 as 1400 - 1345 = 55

I think this does the trick

val windowSpec = Window.partitionBy("ID")
  .orderBy("datetime_seconds")

val consecutiveIncidents = dtf.withColumn("raw_datetime", concat($"Date", $"Time"))
  .withColumn("datetime", to_timestamp($"raw_datetime", "yyyyMMddHHmm"))
  .withColumn("datetime_seconds", $"datetime".cast(LongType))
  .withColumn("time_delta", ($"datetime_seconds" - lag($"datetime_seconds", 1).over(windowSpec)) / 60)
  .filter($"time_delta" === lit(15))
  .select("ID", "Date")
  .distinct()
  .collect()
  .map { case Row(id, date) => (id, date) }
  .toList

Basically - convert the datetimes to timestamps, then look for records with the same ID and consecutive times, with their times separated by 15 minutes.

This is done by using lag over a window grouped by ID and ordered by the time.

In order to calculate the time difference The timestamp is converted to unix epoch seconds.

If you don't want to count day-crossing incidients, you can add the date to the groupyBy clause of the window

user_154333080
Related