Spark DataFrame同值同排名需求：rank函数无法满足的问题

你之前的代码错误在于窗口定义：Window.partitionBy("text").orderBy("id")是将每个text作为独立分区，然后在分区内按id排序生成rank，这会导致同一text下不同id得到不同rank，不符合需求。

要实现同一text值对应相同rank，可以采用以下两种方法：

方法一：全局窗口 + dense_rank() #

直接使用无分区的全局窗口，按text排序后用dense_rank()函数，该函数会为相同排序键的行分配相同的连续rank：

import pyspark.sql.functions as F
from pyspark.sql.window import Window

# 定义全局窗口，按text字段排序
global_window = Window.orderBy("text")
# 添加rank列
result_df = sdf.withColumn("rank", F.dense_rank().over(global_window))
# 按预期顺序展示结果
result_df.orderBy("text", "id").show()

输出结果：

+---+----+----+
| id|text|rank|
+---+----+----+
| p1|  t1|   1|
| p2|  t1|   1|
| p3|  t1|   1|
| p4|  t2|   2|
| p4|  t3|   3|
| p4|  t3|   3|
+---+----+----+

方法二：先生成text的rank映射，再关联 #

先提取唯一的text值并分配rank，再与原DataFrame关联，适合需要自定义text排序规则的场景：

import pyspark.sql.functions as F
from pyspark.sql.window import Window

# 生成text到rank的映射表
text_rank_map = sdf.select("text") \
    .distinct() \
    .orderBy("text") \
    .withColumn("rank", F.row_number().over(Window.orderBy("text")))

# 关联原DataFrame
result_df = sdf.join(text_rank_map, on="text", how="left") \
    .orderBy("text", "id")
result_df.show()

此方法同样会得到符合预期的结果，若需要调整rank的排序逻辑，只需修改orderBy的字段即可。

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