pyspark - 如何在pyspark数据帧中的groupby之上找到加权总和？

Question

我有一个数据框，我需要首先应用数据框，然后获得加权平均值，如下面的输出计算所示。 pyspark 中这样做的有效方法是什么？

data = sc.parallelize([
[111,3,0.4],
[111,4,0.3],
[222,2,0.2],
[222,3,0.2],
[222,4,0.5]]
).toDF(['id', 'val','weight'])
data.show()


+---+---+------+
| id|val|weight|
+---+---+------+
|111|  3|   0.4|
|111|  4|   0.3|
|222|  2|   0.2|
|222|  3|   0.2|
|222|  4|   0.5|
+---+---+------+

输出:

id  weigthed_val
111 (3*0.4 + 4*0.3)/(0.4 + 0.3)
222 (2*0.2 + 3*0.2+4*0.5)/(0.2+0.2+0.5)

Answer 1

您可以将列权重和 val 相乘，然后聚合:

import pyspark.sql.functions as F
data.groupBy("id").agg((F.sum(data.val * data.weight)/F.sum(data.weight)).alias("weighted_val")).show()

+---+------------------+
| id|      weighted_val|
+---+------------------+
|222|3.3333333333333335|
|111|3.4285714285714293|
+---+------------------+