Как эффективно разделить каждое значение на сумму значений в одной группе?

Question

У меня есть Spark DataFrame следующим образом:

F1  F2  F3
x   y   n1
x   z   n2
x   a   n3
x   b   n4
t   y   n5
t   y2  n6
t   y3  n7
t   y4  n8

Для каждой строки я хочу создать еще один кадр данных с другим полем F4 следующим образом:

F1  F2  F3  F4
x   y   n1  n1/(n1+2+n3+n4)
x   z   n2  n2/(n1+2+n3+n4)
x   a   n3  n3/(n1+2+n3+n4)
x   b   n4  n4/(n1+2+n3+n4)
t   y   n5  n5/(n5+n6+n7+n8)
t   y2  n6  n6/(n5+n6+n7+n8)
t   y3  n7  n7/(n5+n6+n7+n8)
t   y4  n8  n8/(n5+n6+n7+n8)

То есть каждый F3 должен быть разделен на сумму F3 с одинаковым значением F1.

Как выполнить эти агрегации в Spark Scala?

Answer 1

Оконные функции удобны при добавлении дополнительных столбцов, которые зависят от группы / раздела данных. В этом случае вы можете использовать оконную функцию для вычисления суммы всех значений F3, разделенных на столбец F1.

Это может быть решено следующим образом:

val w = Window.partitionBy("F1")
val df2 = df.withColumn("F4", $"F3" / sum($"F3").over(w))