如何在Pyspark DataFrame中使用groupby找到加权和？

Question

我有一个DataFrame，首先需要应用DataFrame，然后按照下面的计算输出获取加权平均值。在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)

- learner

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|
+---+------------------+

- Psidom

---

谢谢。完美地工作了。 - learner