Окончание датирования записей с использованием оконных функций в Spark SQL

Question

У меня есть кадр данных, как показано ниже

+----+----+----------+----------+
|colA|colB|      colC|      colD|
+----+----+----------+----------+
|   a|   2|2013-12-12|2999-12-31|
|   b|   3|2011-12-14|2999-12-31|
|   a|   4|2013-12-17|2999-12-31|
|   b|   8|2011-12-19|2999-12-31|
|   a|   6|2013-12-23|2999-12-31|
+----+----+----------+----------+

Мне нужно сгруппировать записи на основе ColA и ранжировать записи на основе colC (последняя дата получает больший ранг), а затем обновить даты в colD повычитание дня из записи colC соседнего ранга.

Окончательный кадр данных должен выглядеть следующим образом:

+----+----+----------+----------+
|colA|colB|      colC|      colD|
+----+----+----------+----------+
|   a|   2|2013-12-12|2013-12-16|
|   a|   4|2013-12-17|2013-12-22|
|   a|   6|2013-12-23|2999-12-31|
|   b|   3|2011-12-14|2011-12-18|
|   b|   8|2011-12-29|2999-12-31|
+----+----+----------+----------+

Answer 1

Вы можете получить его, используя функции окна

scala> val df = Seq(("a",2,"2013-12-12","2999-12-31"),("b",3,"2011-12-14","2999-12-31"),("a",4,"2013-12-17","2999-12-31"),("b",8,"2011-12-19","2999-12-31"),("a",6,"2013-12-23","2999-12-31")).toDF("colA","colB","colC","colD")
df: org.apache.spark.sql.DataFrame = [colA: string, colB: int ... 2 more fields]

scala> val df2 = df.withColumn("colc",'colc.cast("date")).withColumn("cold",'cold.cast("date"))
df2: org.apache.spark.sql.DataFrame = [colA: string, colB: int ... 2 more fields]

scala> df2.createOrReplaceTempView("yash")

scala> spark.sql(""" select cola,colb,colc,cold, rank() over(partition by cola order by colc) c1, coalesce(date_sub(lead(colc) over(partition by cola order by colc),1),cold) as cold2 from yash """).show
+----+----+----------+----------+---+----------+
|cola|colb|      colc|      cold| c1|     cold2|
+----+----+----------+----------+---+----------+
|   b|   3|2011-12-14|2999-12-31|  1|2011-12-18|
|   b|   8|2011-12-19|2999-12-31|  2|2999-12-31|
|   a|   2|2013-12-12|2999-12-31|  1|2013-12-16|
|   a|   4|2013-12-17|2999-12-31|  2|2013-12-22|
|   a|   6|2013-12-23|2999-12-31|  3|2999-12-31|
+----+----+----------+----------+---+----------+


scala> 

Удаление ненужных столбцов

scala> spark.sql(""" select cola,colb,colc, coalesce(date_sub(lead(colc) over(partition by cola order by colc),1),cold) as cold from yash """).show
+----+----+----------+----------+
|cola|colb|      colc|      cold|
+----+----+----------+----------+
|   b|   3|2011-12-14|2011-12-18|
|   b|   8|2011-12-19|2999-12-31|
|   a|   2|2013-12-12|2013-12-16|
|   a|   4|2013-12-17|2013-12-22|
|   a|   6|2013-12-23|2999-12-31|
+----+----+----------+----------+


scala>

Answer 2

Вы можете создать row_number для раздела по colA и упорядочить по colC, а затем самостоятельно объединить на фрейме данных.Код должен выглядеть следующим образом.

val rnkDF = df.withColumn("rnk", row_number().over(Window.partitionBy("colA").orderBy($"colC".asc)))
  .withColumn("rnkminusone", $"rnk" - lit(1))

val joinDF = rnkDF.alias('A).join(rnkDF.alias('B), ($"A.colA" ===  $"B.colA").and($"A.rnk" === $"B.rnkminusone"),"left")
    .select($"A.colA".as("colA")
    , $"A.colB".as("colB")
    , $"A.colC".as("colC")
    , when($"B.colC".isNull, $"A.colD").otherwise(date_sub($"B.colC", 1)).as("colD"))

Результаты приведены ниже.Надеюсь, это поможет.

+----+----+----------+----------+ |colA|colB| colC| colD| +----+----+----------+----------+ | a| 2|2013-12-12|2013-12-16| | a| 4|2013-12-17|2013-12-22| | a| 6|2013-12-23|2999-12-31| | b| 3|2011-12-14|2011-12-18| | b| 8|2011-12-19|2999-12-31| +----+----+----------+----------+