Spark Count 列值的连续性

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 和 item 进行分区来累积计数 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的第一次连续20180102和从date = 20180101开始的第二次连续date = 20180104 。 如果有一个date = 20180105的记录，它也会有index = 1并持续到date = 20180105 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     |
+---+--------+----+--------+----+-----+------+

我没有删除临时列来显示此代码的工作原理。