Add sequence id for groupby column

Question

I have a requirement that to add a sequence id for groupby column. The original Dataset is like this:

+------------+-----------------+-------+
|   appId|                  rpc|elapsed|
+------------+-----------------+-------+
|     account|/rpc1            |      7|
|     service|/rpc4            |      0|
|     service|/rpc5            |      0|
|     account|/rpc1            |     78|
|         api|/rpc2            |     87|
|         api|/rpc2            |     52|
|     service|/rpc4            |      0|
|         api|/rpc3            |     52|
|     service|/rpc4            |      1|
|     service|/rpc4            |      0|
|     service|/rpc5            |      0|
+------------+-----------------+-------+

After doing a dataset.select("appId", "rpc","elapsed").orderby("appId", "rpc","elapsed").show

+------------+-----------------+-------+
|   appId|                  rpc|elapsed| 
+------------+-----------------+-------+
|     account|/rpc1            |      7| 
|     account|/rpc1            |     78|  
|         api|/rpc2            |     87| 
|         api|/rpc2            |     52|  
|         api|/rpc3            |     52| 
|     service|/rpc4            |      0|
|     service|/rpc4            |      1|  
|     service|/rpc4            |      0|  
|     service|/rpc4            |      0| 
|     service|/rpc5            |      0|  
|     service|/rpc5            |      0|  
+------------+-----------------+-------+

I want to add an id column for the groupbyed result as following

+------------+-----------------+-------+---+
|   appId|                  rpc|elapsed| id|
+------------+-----------------+-------+---+
|     account|/rpc1            |      7|  1|
|     account|/rpc1            |     78|  2|
|         api|/rpc2            |     87|  1|
|         api|/rpc2            |     52|  2|
|         api|/rpc3            |     52|  1|
|     service|/rpc4            |      0|  1|
|     service|/rpc4            |      1|  2|
|     service|/rpc4            |      0|  3|
|     service|/rpc4            |      0|  4|
|     service|/rpc5            |      0|  1|
|     service|/rpc5            |      0|  2|
+------------+-----------------+-------+---+

How to achieve this?

Answer 1

you can such an id using window functions:

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

dataset
.withColumn("id",row_number().over(Window.partitionBy("appId","rpc").orderBy("elapsed"))
.show