如何使用Scala在Spark数据框中为每个不同的列值创建行ID的列

Question

我在scala spark中有一个数据框，

类别|得分|

A | 0.2

A | 0.3

A | 0.3

B | 0.9

B | 0.8

B | 1

我想 将行ID列添加为

类别|得分|行ID

A | 0.2 | 0

A | 0.3 | 1

A | 0.3 | 2

B | 0.9 | 0

B | 0.8 | 1

B | 1 | 2

基本上，我希望列类别中每个不同的值的行ID单调增加。我已经有一个排序的数据框，因此所有具有相同类别的行都分组在一起。但是，我仍然不知道如何生成在出现新类别时重新启动的row_id。请帮忙！

Answer 1

这是Window聚合函数的好用例

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.row_number
import df.sparkSession.implicits._

val window = Window.partitionBy('category).orderBy('score)
df.withColumn("row-id", row_number.over(window))

Window函数的工作方式类似于groupBy，除了窗口组中的每个 row 返回单个值而不是每个 group 返回单个值。在这种情况下，该值是行的位置在同一类别的行组中。另外，如果这是您要实现的效果，则无需预先对列category进行预排序。