PySpark DataFrame分组后获取最大值对应行标签的实现方法

在PySpark中可以通过窗口函数实现类似Pandas里idxmax的需求，具体步骤如下：

1. 导入所需的窗口工具和函数：

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, desc

2. 定义窗口规则：按group列分组，每组内按len_text降序排列

window_spec = Window.partitionBy("group").orderBy(desc("len_text"))

3. 添加组内排名并筛选目标行：

# 给每行添加组内排名，筛选排名第一的行后删除临时rank列
result_df = df.withColumn("rank", row_number().over(window_spec)) \
              .filter("rank == 1") \
              .drop("rank")

# 查看结果
result_df.show(truncate=False)

执行后会得到预期输出：

+--------------------+----------+-----+--------+
|text                |word_count|group|len_text|
+--------------------+----------+-----+--------+
|I went home         |3         |1    |11      |
|I looked at the cat |4         |2    |19      |
|The cat looked at me|5         |3    |20      |
|I went homes        |3         |4    |12      |
+--------------------+----------+-----+--------+

补充说明 #

• 如果同一组内有多个行的len_text值同为最大值，row_number()会只保留其中一行；若要保留所有最大值的行，可以替换为rank()或dense_rank()函数。

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