通过分组数据向数据框添加索引

Question

我在分组数据和添加索引时遇到问题。

一个新的列索引将从 1 开始添加到 n（例如：5）并再次从 1 迭代。 该值可以是任何值，因此基本上在每 n 条记录之后，索引应该以 1 重新开始。

原始数据框

city	id
NYC	101
Jersey City	102
Hoboken	103
Buffalo	104
Philly	105
Edison	106

输出数据框应该是这样的

city	id	index
NYC	101	1
Jersey City	102	2
Hoboken	103	3
Buffalo	104	4
Philly	105	5
Edison	106	1
trenton	107	2

Answer 1

试试这个：

data
  .withColumn("groupId", ceil(col("id") / lit(5)))
  .withColumn("index", row_number() over Window.partitionBy("groupId").orderBy("id"))
  .drop(col("groupId"))

输出（测试）：

+-----------+---+-----+
|       city| id|index|
+-----------+---+-----+
|        NYC|101|    1|
|Jersey City|102|    2|
|    Hoboken|103|    3|
|    Buffalo|104|    4|
|     Philly|105|    5|
|     Edison|106|    1|
|    Trenton|107|    2|
+-----------+---+-----+

Answer 2

你试过 ntile 吗？它将确保您始终将 id 放置在 5 个桶中，排名从 1 到 5。

WindowSpec window = Window.orderBy(col("id").asc());
dataset.withColumn("index", ntile(5).over(window)).show();