Scala 火花计数列值条纹

我有一个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

，其中包含每个用户拥有的项目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

---

我可以使用窗口功能

行数

和按id和项目划分来累计计算id项目对，但这不会在出现新项目后重新开始计数。

我的最佳解决方案就是这样

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)

跳转

列用于计算

索引

，其中

索引

表示

条纹

的索引。例如，

id=1

和

item=A

，应该有两个索引。索引

0

和

1

分别表示从

date=20180101

到

20180102

的第一条条纹和从

date=20180104

的第二条条纹。如果有一条记录的日期为20180105，则该记录的索引也将为1，并继续为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     |
+---+--------+----+--------+----+-----+------+

---

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

太好了！很好。