Pyspark：添加具有 groupby 平均值的列

Posted 2023-04-15

【中文标题】Pyspark：添加具有 groupby 平均值的列

【英文标题】：Pyspark: Add column with average of groupby

【发布时间】：2021-02-02 23:04:55

【问题描述】：

我有一个这样的数据框：

test = spark.createDataFrame(
    [
        (1, 0, 100), 
        (2, 0, 200),
        (3, 1, 150), 
        (4, 1, 250),
    ],
    ['id', 'flag', 'col1'] 
)

我想创建另一列并输入标志的 groupby 的平均值

test.groupBy(f.col('flag')).agg(f.avg(f.col("col1"))).show()

+----+---------+
|flag|avg(col1)|
+----+---------+
|   0|    150.0|
|   1|    200.0|
+----+---------+

最终产品：

+---+----+----+---+
| id|flag|col1|avg|
+---+----+----+---+
|  1|   0| 100|150|
|  2|   0| 200|150|
|  3|   1| 150|200|
|  4|   1| 250|200|
+---+----+----+---+

【问题讨论】：

这能回答你的问题吗？ Adding a group count column to a PySpark dataframe

【参考方案1】：

你可以使用window函数：

from pyspark.sql.window import Window
from pyspark.sql import functions as F

w = Window.partitionBy('flag')
test.withColumn("avg", F.avg("col1").over(w)).show()

+---+----+----+-----+                                                           
| id|flag|col1|  avg|
+---+----+----+-----+
|  1|   0| 100|150.0|
|  2|   0| 200|150.0|
|  3|   1| 150|200.0|
|  4|   1| 250|200.0|
+---+----+----+-----+