Взятие суммы ini spark- scala по условию

Question

У меня есть такой фрейм данных. Как я могу взять сумму продаж столбца, где рейтинг больше 3, за 'M'

+---+-----+----+
|  M|Sales|Rank|
+---+-----+----+
| M1|  200|   1|
| M1|  175|   2|
| M1|  150|   3|
| M1|  125|   4|
| M1|   90|   5|
| M1|   85|   6|
| M2| 1001|   1|
| M2|  500|   2|
| M2|  456|   3|
| M2|  345|   4|
| M2|  231|   5|
| M2|  123|   6|
+---+-----+----+

Ожидаемый результат -

+---+-----+----+---------------+
|  M|Sales|Rank|SumGreaterThan3|
+---+-----+----+---------------+
| M1|  200|   1|            300|
| M1|  175|   2|            300|
| M1|  150|   3|            300|
| M1|  125|   4|            300|
| M1|   90|   5|            300|
| M1|   85|   6|            300|
| M2| 1001|   1|            699|
| M2|  500|   2|            699|
| M2|  456|   3|            699|
| M2|  345|   4|            699|
| M2|  231|   5|            699|
| M2|  123|   6|            699|
+---+-----+----+---------------+

Я сделал сумму по ROwnumber как это

df.withColumn("SumGreaterThan3",sum("Sales").over(Window.partitionBy(col("M"))))` //But this will provide total sum of sales.

Чтобы повторить тот же DF-

val df = Seq(
("M1",200,1),
("M1",175,2),
("M1",150,3),
("M1",125,4),
("M1",90,5),
("M1",85,6),
("M2",1001,1),
("M2",500,2),
("M2",456,3),
("M2",345,4),
("M2",231,5),
("M2",123,6)
).toDF("M","Sales","Rank")

Answer 1

Что ж, раздела достаточно для установки функции window. Конечно, вы также должны использовать условное суммирование, смешивая sum и when.

import org.apache.spark.sql.expressions.Window
val w = Window.partitionBy("M")
df.withColumn("SumGreaterThan3", sum(when('Rank > 3, 'Sales).otherwise(0)).over(w).alias("sum")).show

Это даст вам ожидаемые результаты.