如何在Spark DataFrame的窗口分区中分配唯一组ID？

我来帮你搞定这个需求！你想要给每个相同col值的行分配连续的组ID，用窗口函数完全可以实现，不过你之前的思路有两个小问题：一是Spark里并没有id()这个用来生成组ID的内置函数，二是单纯按col分区的窗口只能在每个分区内生成序号，没法得到全局连续的组编号。

下面给你两种靠谱的实现方法：

方法一：用dense_rank()窗口函数（推荐） #

dense_rank()会根据指定列的顺序生成连续的排名，相同值的行排名一致，而且不会跳过数字，正好匹配你的需求：

from pyspark.sql import Window
from pyspark.sql.functions import dense_rank

# 定义窗口：按col列排序，生成全局连续排名
group_window = Window.orderBy("col")
# 添加group列，用dense_rank()生成组ID
df = df.withColumn("group", dense_rank().over(group_window))
# 只保留需要的id和group列
result_df = df.select("id", "group")
result_df.show()

运行后就能得到你想要的结果：

+---+-----+
| id|group|
+---+-----+
|  1|    1|
|  2|    1|
|  3|    2|
|  4|    3|
|  5|    3|
+---+-----+

方法二：先给唯一列值分配ID再关联 #

如果你不依赖col的排序顺序，也可以先提取所有唯一的col值，给它们分配连续ID后再关联回原表：

from pyspark.sql import Window
from pyspark.sql.functions import row_number

# 提取唯一col值，按col排序后分配连续组ID
col_group_mapping = df.select("col")\
                      .distinct()\
                      .orderBy("col")\
                      .withColumn("group", row_number().over(Window.orderBy("col")))

# 关联回原DataFrame，得到最终结果
result_df = df.join(col_group_mapping, on="col", how="left")\
              .select("id", "group")
result_df.show()

这种方法的好处是可以灵活控制组ID的生成逻辑，比如调整排序规则或者自定义ID起始值。

为什么你的原代码不行？ #

你之前用Window.partitionBy('col')是把相同col的行分成不同分区，这时候如果用row_number()只会在每个分区内生成1、2...的序号，而不是全局连续的组ID；另外Spark中没有id()这个函数，你可能是混淆了其他工具的语法~

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