Pyspark - Как создать столбец, воспроизводящий определенное c число из ячейки

Question

Мне нужно преобразовать следующий фрейм данных:

Group - Risk Group - State - Value
A           yes        1      1000
A           no         2      100
A           no         3      100
A           no         4      100
A           no         5      100
B           yes        1      2000
B           no         2      200
B           no         3      100
B           no         4      100
B           no         5      200
C           yes        1      3000
C           no         2      100
C           no         3      100
C           no         4      200
C           no         5      200

В это, где я переношу группу риска «да» в столбец:

Group - Risk Group - State - Value - Risk Group Value
A           no         2      100         1000
A           no         3      100         1000
A           no         4      100         1000
A           no         5      100         1000 
B           no         2      200         2000
B           no         3      100         2000
B           no         4      100         2000
B           no         5      200         2000
C           no         2      100         3000
C           no         3      100         3000
C           no         4      200         3000
C           no         5      200         3000

Мой окончательный результат будет чтобы создать новый столбец, устанавливающий соотношение между (значение группы риска / значение), но это легко сделать.

Спасибо!

Answer 1

Вы можете суммировать только when Группа риска: yes, более window partitioned by Group, а затем filter удалить yes строк, чтобы получить желаемый результат.

from pyspark.sql import functions as F
from pyspark.sql.window import Window

w=Window().partitionBy("Group")
df.withColumn("Risk Group Value", F.sum(F.when(F.col("Risk Group")=='yes',F.col("Value"))).over(w))\
  .filter(F.col("Risk Group")!='yes')\
  .orderBy("Group").show()

#+-----+----------+-----+-----+----------------+
#|Group|Risk Group|State|Value|Risk Group Value|
#+-----+----------+-----+-----+----------------+
#|    A|        no|    2|  100|            1000|
#|    A|        no|    3|  100|            1000|
#|    A|        no|    4|  100|            1000|
#|    A|        no|    5|  100|            1000|
#|    B|        no|    2|  200|            2000|
#|    B|        no|    3|  100|            2000|
#|    B|        no|    4|  100|            2000|
#|    B|        no|    5|  200|            2000|
#|    C|        no|    2|  100|            3000|
#|    C|        no|    5|  200|            3000|
#|    C|        no|    3|  100|            3000|
#|    C|        no|    4|  200|            3000|
#+-----+----------+-----+-----+----------------+