Search code examples
apache-sparkpyspark

How to find weighted sum on top of groupby in pyspark dataframe?

I have a dataframe where i need to first apply dataframe and then get weighted average as shown in the output calculation below. What is an efficient way in pyspark to do that? 

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

Output:

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)
Solution

  • You can multiply columns weight and val, then aggregate:

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