Счетчик искр

Question

У меня есть df

id, date, item

1, 20180101, A
1, 20180102, A
1, 20180103, B
1, 20180104, A
2, 20180101, C
2, 20180102, D
2, 20180103, D
2, 20180104, D

, и я хотел бы создать новый столбец streak, который содержит количество серий элементов, которые есть у каждого пользователя

id, date, item, streak

1, 20180101, A, 1
1, 20180102, A, 2
1, 20180103, B, 1
1, 20180104, A, 1
2, 20180101, C, 1
2, 20180102, D, 1
2, 20180103, D, 2
2, 20180103, D, 3

, которые я могу использовать оконная функция row_number и разделение по id и элементу для кумулятивного подсчета пары id-item, но это не перезапустит серию счетчиков после появления нового элемента.

Answer 1

Сделайте мое лучшее решение:

import org.apache.spark.sql.expressions.Window

val w1 = Window.partitionBy("id", "item").orderBy("date")
val w2 = Window.partitionBy("id", "item", "index").orderBy("date")
df.withColumn("lag_date", lag("date", 1, "").over(w1))
  .withColumn("jump", not(col("lag_date") === lit("") || date_add(to_date(col("lag_date"), "yyyyMMdd"), 1) === to_date(col("date"), "yyyyMMdd")).cast("int"))
  .withColumn("index", sum("jump").over(w1))
  .withColumn("streak", row_number.over(w2))
  .orderBy("id", "date")
  .show(false)

Столбец jump используется для вычисления index, где index означает, что индекс streak. Например, id = 1 и item = A, должно быть 2 индекса. Индексы 0 и 1 указывают первую полосу от date = 20180101 до 20180102 и вторую полосу от date = 20180104 соответственно. Если есть запись с date = 20180105, она также будет иметь index = 1 и продолжится до streak = 2.

Результат:

+---+--------+----+--------+----+-----+------+
|id |date    |item|lag_date|jump|index|streak|
+---+--------+----+--------+----+-----+------+
|1  |20180101|A   |        |0   |0    |1     |
|1  |20180102|A   |20180101|0   |0    |2     |
|1  |20180103|B   |        |0   |0    |1     |
|1  |20180104|A   |20180102|1   |1    |1     |
|2  |20180101|C   |        |0   |0    |1     |
|2  |20180102|D   |        |0   |0    |1     |
|2  |20180103|D   |20180102|0   |0    |2     |
|2  |20180104|D   |20180103|0   |0    |3     |
+---+--------+----+--------+----+-----+------+

, где я не удалены временные столбцы, чтобы показать, как работает этот код.