如何在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

您可以将列 weight 和 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|
+---+------------------+