如何在PySpark中按指定列分组并保留原DataFrame所有行？

要实现按指定列分组聚合（比如按id分组求value的最大值），同时保留原始DataFrame的所有行，最直接高效的方式是使用窗口函数（Window Functions），无需额外的join操作。

具体实现步骤： #

1. 导入PySpark的窗口工具和聚合函数
2. 定义分组窗口（按目标列分区）
3. 通过withColumn为每行添加对应分组的聚合结果

代码示例： #

首先创建原始DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import max

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

# 构造原始数据
data = [(1, 1), (1, 2), (2, 3), (2, 4)]
df = spark.createDataFrame(data, ["id", "value"])
df.show()

然后执行窗口函数计算：

# 定义窗口规则：按id列分区
window_spec = Window.partitionBy("id")

# 新增max列，计算每个id分组内的value最大值
result_df = df.withColumn("max", max("value").over(window_spec))
result_df.show()

执行结果： #

+---+-----+---+
| id|value|max|
+---+-----+---+
|  1|    1|  2|
|  1|    2|  2|
|  2|    3|  4|
|  2|    4|  4|
+---+-----+---+

补充说明： #

• 窗口函数会针对每个分区（这里是每个id组）计算聚合值，再将结果匹配到该分区的每一行，既保留原始行结构，又完成分组聚合。
• 除了max，你可以替换成其他聚合函数（如min、avg、sum等），只需修改max("value")为对应函数即可。

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