Scala Spark基于数据帧中的另一列增加一列而不使用for循环

Question

我有一个如下所示的数据框。 我想要一个名为 cutofftype 的新列 - 每次 ID 列更改时，它应该重置为 1，而不是当前单调递增的数字。

df = df.orderBy("ID","date").withColumn("cutofftype",monotonically_increasing_id()+1)

+------+---------------+----------+
|   ID |    date       |cutofftype|
+------+---------------+----------+
| 54441|     2016-06-20|         1|
| 54441|     2016-06-27|         2|
| 54441|     2016-07-04|         3|
| 54441|     2016-07-11|         4|
| 54500|     2016-05-02|         5|
| 54500|     2016-05-09|         6|
| 54500|     2016-05-16|         7|
| 54500|     2016-05-23|         8|
| 54500|     2016-06-06|         9|
| 54500|     2016-06-13|        10|
+------+---------------+----------+

目标如下：

+------+---------------+----------+
|   ID |    date       |cutofftype|
+------+---------------+----------+
| 54441|     2016-06-20|         1|
| 54441|     2016-06-27|         2|
| 54441|     2016-07-04|         3|
| 54441|     2016-07-11|         4|
| 54500|     2016-05-02|         1|
| 54500|     2016-05-09|         2|
| 54500|     2016-05-16|         3|
| 54500|     2016-05-23|         4|
| 54500|     2016-06-06|         5|
| 54500|     2016-06-13|         6|
+------+---------------+----------+

我知道这可以用 for 循环来完成 - 我想不用 for 循环 >> 有出路吗？

Answer 1

按问题简单划分。 你应该使用window 。

import org.apache.spark.sql.expressions.Window
val w = Window.partitionBy("ID").orderBy("date")

df.withColumn("cutofftype", row_number().over(w)).show()


+-----+----------+----------+
|   ID|      date|cutofftype|
+-----+----------+----------+
|54500|2016-05-02|         1|
|54500|2016-05-09|         2|
|54500|2016-05-16|         3|
|54500|2016-05-23|         4|
|54500|2016-06-06|         5|
|54500|2016-06-13|         6|
|54441|2016-06-20|         1|
|54441|2016-06-27|         2|
|54441|2016-07-04|         3|
|54441|2016-07-11|         4|
+-----+----------+----------+