Pyspark-通过ID分配组中的每个组[重复]

Posted 2023-03-31

【中文标题】Pyspark-通过ID分配组中的每个组[重复]

【英文标题】：Pyspark- Assign each group in groupBy an ID [duplicate]

【发布时间】：2018-12-18 10:17:47

【问题描述】：

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

我之前使用带有 python 的 pandas 和命令：

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 |
+------+--+

【参考方案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|
+------+---+

【讨论】：

顺序本身并不重要，只要按组递增就可以了。

如果你有数据分区会发生什么？我收到警告：WARN WindowExec:66 - No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.

【参考方案2】：

from pyspark.sql.functions import col, create_map, lit
from itertools import chain
values = [('apple',),('orange',),('pear',),('berry',),('apple',),('pear',),('berry',)]
df = sqlContext.createDataFrame(values,['object'])

#Creating a column of distinct elements and converting them into dictionary with unique indexes.
df1 = df.distinct()
distinct_list = list(df1.select('object').toPandas()['object'])
dict_with_index = distinct_list[i]:i+1 for i in range(len(distinct_list))

#Applying the mapping of dictionary.
mapping_expr = create_map([lit(x) for x in chain(*dict_with_index.items())])
df=df.withColumn("id", mapping_expr.getItem(col("object")))
df.show()
+------+---+
|object| id|
+------+---+
| apple|  2|
|orange|  1|
|  pear|  3|
| berry|  4|
| apple|  2|
|  pear|  3|
| berry|  4|
+------+---+