PySpark中实现R语言group+ungroup分组求和的等效方法

需求说明 #

需要按metric字段（取值为I或M）分组，对x字段的值求和，并将求和结果添加到对应分组的每一行中。在R中通常用group_by+ungroup实现，想了解PySpark的等效方法。

数据与代码准备 #

初始化代码 #

from pyspark.sql.functions import *
from pyspark.sql.types import *

data = [["1", "Amit", "DU", "I", "8", "6"],
        ["2", "Mohit", "DU", "I", "4", "2"],
        ["3", "rohith", "BHU", "I", "5", "3"],
        ["4", "sridevi", "LPU", "I", "1", "6"],
        ["1", "sravan", "KLMP", "M", "2", "4"],
        ["5", "gnanesh", "IIT", "M", "6", "8"],
        ["6", "gnadesh", "KLM", "M","0", "9"]]

columns = ['ID', 'NAME', 'college', 'metric', 'x', 'y']

dataframe = spark.createDataFrame(data, columns)
dataframe = dataframe.withColumn("x", dataframe.x.cast(DoubleType()))

原始数据样例 #

+---+-------+-------+------+----+---+
| ID|   NAME|college|metric|   x|  y|
+---+-------+-------+------+----+---+
|  1|   Amit|     DU|     I|   8|  6|
|  2|  Mohit|     DU|     I|   4|  2|
|  3| rohith|    BHU|     I|   5|  3|
|  4|sridevi|    LPU|     I|   1|  6|
|  1| sravan|   KLMP|     M|   2|  4|
|  5|gnanesh|    IIT|     M|   6|  8|
|  6|gnadesh|    KLM|     M|0   |  9|
+---+-------+-------+------+----+---+

期望输出 #

+---+-------+-------+------+----+---+------+
| ID|   NAME|college|metric|   x|  y| total|
+---+-------+-------+------+----+---+------+
|  1|   Amit|     DU|     I|   8|  6|  18  |
|  2|  Mohit|     DU|     I|   4|  2|  18  |  
|  3| rohith|    BHU|     I|   5|  3|  18  |
|  4|sridevi|    LPU|     I|   1|  6|  18  |
|  1| sravan|   KLMP|     M|   2|  4|  8   |
|  5|gnanesh|    IIT|     M|   6|  8|  8   |
|  6|gnadesh|    KLM|     M|   0|  9|  8   |
+---+-------+-------+------+----+---+------+

失败尝试 #

dataframe.withColumn("total", dataframe.groupBy("metric").sum("x"))

解决方案 #

方法1：使用窗口函数（推荐，最贴近R的实现逻辑） #

窗口函数可直接在原DataFrame上计算分组聚合值，并将结果映射到每一行，无需额外关联操作，逻辑和R中group_by后mutate添加聚合列完全一致。

# 定义窗口：按metric分组
window_spec = Window.partitionBy("metric")

# 添加total列，计算每个metric分组的x字段求和
result_df = dataframe.withColumn("total", sum("x").over(window_spec))

# 查看结果
result_df.show()

方法2：分组聚合后关联原表 #

如果习惯先聚合再关联的方式，可以先按metric分组计算求和，再通过metric字段将结果关联回原DataFrame。

# 分组计算每个metric的x总和
sum_df = dataframe.groupBy("metric").agg(sum("x").alias("total"))

# 将聚合结果关联回原表
result_df = dataframe.join(sum_df, on="metric", how="left")

# 查看结果
result_df.show()

说明 #

• 方法1的窗口函数更高效，尤其是大数据量场景，避免了关联操作带来的性能开销。
• 失败尝试的问题在于withColumn的参数不能直接传入分组聚合后的DataFrame，PySpark不支持这种嵌套操作，必须用窗口函数或关联来实现。

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