PySpark DataFrame按双列分组排序后生成组内重置行索引的实现

你可以使用 PySpark 内置的窗口函数实现分组内排序生成行索引的需求，完整代码如下：

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

# 定义窗口规则：按user_id分组，分组内按var1升序排序
# 如果你需要按var1降序排序，可修改为 orderBy(col("var1").desc())
window_spec = Window.partitionBy("user_id").orderBy("var1")

# 新增索引列order
df_with_order = df1.withColumn("order", row_number().over(window_spec))

# 打印结果
df_with_order.show(truncate=False)

运行后输出结果和你预期的完全一致：

+-------+----+-----+
|user_id|var1|order|
+-------+----+-----+
|u1     |10  |1    |
|u1     |20  |2    |
|u2     |10  |1    |
|u2     |10  |2    |
|u2     |30  |3    |
+-------+----+-----+

核心逻辑说明 #

• partitionBy("user_id")：指定行索引按user_id分组独立计算，每个分组的行号都从1开始
• orderBy("var1")：指定分组内的排序规则，此处为升序，可根据需求调整排序字段和排序方向
• row_number()：为窗口范围内的每行生成连续不重复的整数行号

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