PySpark如何按id分组、按score降序生成从0开始的递增整数排名列

你可以通过PySpark的窗口函数实现需求，核心逻辑是按id分组后在组内排序，用行号减1得到从0开始的排名：

实现步骤 #

• 导入所需的窗口类和排名函数
• 定义窗口规则：分区键为id，分区内排序规则为score降序
• 调用row_number()计算组内行号，行号结果减1即可得到从0开始的rank列
• 可选：如果需要将原score列重命名为value，可在计算后调用withColumnRenamed处理

完整代码示例 #

# 导入依赖
from pyspark.sql import SparkSession
from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

# 初始化SparkSession
spark = SparkSession.builder.appName("rank_demo").getOrCreate()

# 构造测试数据
test_data = [
    (1, 0.5),
    (1, 2.5),
    (2, 4.45),
    (3, 8.5),
    (3, 3.25),
    (3, 5.55)
]
df = spark.createDataFrame(test_data, schema=["id", "score"])

# 定义窗口
window_spec = Window.partitionBy("id").orderBy(df["score"].desc())

# 计算rank列，并重命名score为value
result_df = df.withColumn("rank", row_number().over(window_spec) - 1) \
              .withColumnRenamed("score", "value") \
              .orderBy("id", "rank")

# 输出结果验证
result_df.show()

输出结果验证 #

运行后输出和期望结果完全一致：

+---+-----+----+
| id|value|rank|
+---+-----+----+
|  1|  2.5|   0|
|  1|  0.5|   1|
|  2| 4.45|   0|
|  3|  8.5|   0|
|  3| 5.55|   1|
|  3| 3.25|   2|
+---+-----+----+

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