PySpark按A列GroupBy后保留分组内B列最大值所在行的问题

嘿，刚入门PySpark遇到这种问题太正常啦！你遇到的核心痛点就是普通的groupBy+agg操作只能返回分组键和聚合后的结果列，没法保留原始数据里的其他字段。这时候**窗口函数（Window Functions）**就是你的救星，我给你详细讲怎么实现：

步骤1：导入需要的模块 #

首先得导入窗口函数相关的工具：

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, col, rank

步骤2：定义窗口规则 #

我们需要按"A"列分组，然后在每个分组内按"B"列降序排序——这样每个分组里"B"值最大的行会排在最前面：

# 定义窗口：按A分组，B列降序排序
window_spec = Window.partitionBy("A").orderBy(col("B").desc())

步骤3：添加行号并过滤目标行 #

给每行添加一个行号标记，然后只保留每个分组里行号为1的行（也就是B值最大的那一行）：

# 给每行添加行号
df_with_row_num = df.withColumn("row_num", row_number().over(window_spec))

# 过滤出每个分组中B最大的行，然后删除临时的行号列
df_cleaned = df_with_row_num.filter(col("row_num") == 1).drop("row_num")

特殊情况处理：多个行有相同的最大值 #

如果你的数据里，某个分组内有多行的B值都是最大值，用row_number()会随机保留其中一行。要是你想把所有最大值的行都保留下来，把row_number()换成rank()就行：

# 用rank()替代row_number()，保留所有B值为最大值的行
df_with_rank = df.withColumn("rank", rank().over(window_spec))
df_cleaned = df_with_rank.filter(col("rank") == 1).drop("rank")

举个实际例子，假设你的原始DataFrame是这样的：

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  X| 5 |  a|
|  X| 3 |  b|
|  Y| 7 |  c|
|  Y| 7 |  d|
+---+---+---+

用row_number()会得到：

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  X| 5 |  a|
|  Y| 7 |  c|  # 或随机保留d那一行
+---+---+---+

用rank()则会保留Y组的两行：

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  X| 5 |  a|
|  Y| 7 |  c|
|  Y| 7 |  d|
+---+---+---+

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