Pyspark-通过 ID 分配组中的每个组

Question

我想使用 pyspark 为组中的每个组分配一个从 0 或 1 开始的唯一 ID 号，并为每个组递增 1。

我之前已经使用 pandas 和 python 完成了此操作，命令如下：

df['id_num'] = (df
                .groupby('column_name')
                .grouper
                .group_info[0])

输入和所需输出的玩具示例是：

输入

+------+
|object|
+------+
|apple |
|orange|
|pear  |
|berry |
|apple |
|pear  |
|berry |
+------+

输出：

+------+--+
|object|id|
+------+--+
|apple |1 |
|orange|2 |
|pear  |3 |
|berry |4 |
|apple |1 |
|pear  |3 |
|berry |4 |
+------+--+

Answer 1

我不确定顺序是否重要。如果没有，您可以在这种情况下使用dense_rank 窗口函数

>>> from pyspark.sql.window import Window
>>> import pyspark.sql.functions as F
>>> 
>>> df.show()
+------+
|object|
+------+
| apple|
|orange|
|  pear|
| berry|
| apple|
|  pear|
| berry|
+------+
>>> 
>>> df.withColumn("id", F.dense_rank().over(Window.orderBy(df.object))).show()
+------+---+
|object| id|
+------+---+
| apple|  1|
| apple|  1|
| berry|  2|
| berry|  2|
|orange|  3|
|  pear|  4|
|  pear|  4|
+------+---+