Find 3 closest dates in Spark dataframe per some criteria using pyspark

Question

I have the following spark dataframes:-

df1

id  dia_date
1   2/12/17
1   4/25/16
2   12/8/17
2   6/12/11

df2

id  obs_date    obs_value
1   2/16/17 4
1   2/20/17 2
1   2/9/17  4
1   12/12/18    5
1   4/18/16 1
1   4/18/16 6
1   4/30/16 7
1   5/25/16 9
2   12/12/17    10
2   12/6/17 11
2   12/14/17    4
2   6/11/11 5
2   6/11/11 6

I want the resulting dataframe to be as follows:-

1) Find the three nearest date in by comparing the date in df1 in df2.
2) If it does not have 3 nearest dates insert null for them.
3) The nearest dates should be found grouped by the id field only. Like for dia_date for id '1' we must look at the obs_date field in the df2 for id1 '1' only.

Example of resulting dataframe:-

id  dia_date    obs_date1   obs_val1    obs_date2   obs_val2    obs_date3   obs_val3
1   2/12/17 2/9/17  4   2/16/17 4   2/20/17 2
1   4/25/16 4/18/16 1   4/18/16 6   4/30/16 7
2   12/8/17 12/6/17 11  12/12/17    10  12/14/17    4
2   6/12/11 6/11/11 5   6/11/11 6   null    null

I want to do it using pyspark. Have tried some ways but finding it really difficult since I am just starting with pyspark.

Need up to 6 observations potentially. If on same date, is that nearest? — thebluephantom
Yes, but if both in your second DF, it means you could have more than 3 results. The question is unclear. If you have dia = 4/12 and you have obs = 4/12, 4/12/6/12, 3/12, 2/12, what is your answer expected? I have 2 possibilities. It;s a bit like the csion here green 0, red, black either side. I think you would want 4/12, 3/12, 5/12. Please advise — thebluephantom

thebluephantom thebluephantom · Accepted Answer · 2018-11-04T18:48:26

Here is a Scala answer as the problem has nothing to do with pyspark. You may convert.

Your final output I could not get, but an alternative should suffice.

// Assuming we could also optimize this further but not doing so.
// Assuming distinct values to compare against. If not, then some further logic required.
// A bug found in Ranking it looks like - or me??? Worked around this and dropped the logic
// Pivoting does not help here. Grouping by in SQL to vet specific columns names not 
elegant, used the more Scala approach.

import org.apache.spark.sql.functions._
import spark.implicits._
import java.time._
import org.apache.spark.sql.functions.{rank}
import org.apache.spark.sql.expressions.Window

def toEpochDay(s: String) = LocalDate.parse(s).toEpochDay
val toEpochDayUdf = udf(toEpochDay(_: String))

// Our input.
val df0 = Seq(
              ("1","2018-09-05"), ("1","2018-09-14"), 
              ("2","2018-12-23"), ("5","2015-12-20"),
              ("6","2018-12-23") 
             ).toDF("id", "dia_dt")
val df1 = Seq(
              ("1","2018-09-06", 5), ("1","2018-09-07", 6), ("6","2023-09-07", 7), 
              ("2","2018-12-23", 4), ("2","2018-12-24", 5), ("2","2018-10-23", 5),
              ("1","2017-09-06", 5), ("1","2017-09-07", 6),
              ("5","2015-12-20", 5), ("5","2015-12-21", 6), ("5","2015-12-19", 5), ("5","2015-12-18", 7), ("5","2015-12-22", 5),
              ("5","2015-12-23", 6), ("5","2015-12-17", 6), ("5","2015-12-26", 60)
             ).toDF("id", "obs_dt", "obs_val")

val myExpression = "abs(dia_epoch - obs_epoch)"
// Hard to know how to restrict further at this point.
val df2 = df1.withColumn("obs_epoch", toEpochDayUdf($"obs_dt")) 
val df3 = df2.join(df0, Seq("id"), "inner").withColumn("dia_epoch", toEpochDayUdf($"dia_dt"))
             .withColumn("abs_diff", expr(myExpression)) 

@transient val w1 = org.apache.spark.sql.expressions.Window.partitionBy("id", "dia_epoch" ).orderBy(asc("abs_diff"))
val df4 = df3.select($"*", rank.over(w1).alias("rank")) // This is required

// Final results as collect_list. Distinct column names not so easy due to not being able to use pivot - may be a limitation on knowledge on my side.
df4.orderBy("id", "dia_dt")
   .filter($"rank" <= 3)
   .groupBy($"id", $"dia_dt") 
   .agg(collect_list(struct($"obs_dt", $"obs_val")).as("observations"))
   .show(false)

returns:

+---+----------+---------------------------------------------------+
|id |dia_dt    |observations                                       |
+---+----------+---------------------------------------------------+
|1  |2018-09-05|[[2017-09-07, 6], [2018-09-06, 5], [2018-09-07, 6]]|
|1  |2018-09-14|[[2017-09-07, 6], [2018-09-06, 5], [2018-09-07, 6]]|
|2  |2018-12-23|[[2018-10-23, 5], [2018-12-23, 4], [2018-12-24, 5]]|
|5  |2015-12-20|[[2015-12-19, 5], [2015-12-20, 5], [2015-12-21, 6]]|
|6  |2018-12-23|[[2023-09-07, 7]]                                  |
+---+----------+---------------------------------------------------+

Take it further, heavy lifting done.

Find 3 closest dates in Spark dataframe per some criteria using pyspark

1 Answers