PySpark中为相同item_id生成对应唯一值列的问题

我有一个包含item_id列的PySpark DataFrame，示例数据如下：

+-----------+
|item_id    |
+-----------+
|     BA2C31|
|     BA2C31|
|     B4D456|
|     B4D456|
|     EDJJ88|
+-----------+

创建该DataFrame的代码：

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [(0, 'BA2C31'),
     (1, 'BA2C31'),
     (2, 'B4D456'),
     (3, 'B4D456'),
     (4, 'EDJJ88')],
    ['id', 'item_id'])

需要新增一列u_id，让每个item_id对应唯一的整数标识，期望输出如下：

+-----------+-----+
|item_id    |u_id |
+-----------+-----+
|     BA2C31|101  |
|     BA2C31|101  |
|     B4D456|102  |
|     B4D456|102  |
|     EDJJ88|103  |
+-----------+-----+

我尝试用以下代码实现，但没得到预期结果：

from pyspark.sql.functions import col, sha2, concat

df.withColumn("u_id", sha2(col("item_id")), 256)).show(10, False)

你用sha2的思路不对：首先代码本身有语法错误（多了一个右括号），其次sha2生成的是长哈希字符串，和你要的连续整数标识完全不符。下面是两种可行的实现方法：

方法1：用dense_rank()生成连续标识 #

这种方法直接通过窗口函数为不同的item_id分配连续排名，加上偏移量就能得到你要的101、102格式：

from pyspark.sql.window import Window

# 按item_id排序生成排名，加100让起始值为101
df_with_uid = df.withColumn(
    "u_id",
    F.dense_rank().over(Window.orderBy("item_id")) + 100
)

df_with_uid.show()

输出结果：

+---+--------+-----+
| id| item_id| u_id|
+---+--------+-----+
|  0|  BA2C31|  101|
|  1|  BA2C31|  101|
|  2|  B4D456|  102|
|  3|  B4D456|  102|
|  4|  EDJJ88|  103|
+---+--------+-----+

方法2：先映射唯一item_id再关联 #

如果不想用窗口函数，可以先提取所有唯一的item_id并分配整数id，再关联回原表：

from pyspark.sql.window import Window

# 提取唯一item_id并排序
unique_items = df.select("item_id").distinct().orderBy("item_id")
# 为每个item_id分配从101开始的连续id
unique_items_with_id = unique_items.withColumn(
    "u_id",
    F.row_number().over(Window.orderBy("item_id")) + 100
)

# 关联回原表，保证每个item_id对应相同的u_id
df_with_uid = df.join(unique_items_with_id, on="item_id", how="left")

df_with_uid.show()

这里用row_number()代替monotonically_increasing_id()，能确保id是严格连续的，避免后者生成非连续值的问题。

内容的提问来源于stack exchange，提问作者merkle