Добавление средств групп обратно в большой массив данных без объединения в PySpark

Question

У меня есть этот кусок кода:

shifted_pd = account_level_pd_shifts.filter(account_level_pd_shifts['WITHIN_PD_EXCLUSION'] == True).groupBy(['FORWARD_LOOK_MODEL', 'FOR_PD_TYPE']).agg(f.avg('PD_SHIFT').alias('SHIFTED_PD'))

account_level_pd_shifts = account_level_pd_shifts.drop('SHIFTED_PD').join(f.broadcast(shifted_pd), on=['FORWARD_LOOK_MODEL', 'FOR_PD_TYPE'], how='left')

Я не могу сделать среднее значение для groupby, просто взять первое значение и затем с помощью f.lit () добавить его в новый столбец, потому что сгруппированы у данных есть 4 элемента, это не 1 средство, которое я получаю, я получаю 4 средства для каждого.

Эти 2 строки кода находятся в итерации 23 l oop, и 23 соединения не очень хороши. Есть ли способ избежать этого и добавить средства для каждой группы напрямую или без объединения с большим фреймом данных?

Я могу предоставить дополнительную информацию, если неясно :)

Спасибо за помощь

Answer 1

То, что вы ищете, это функция управления окнами. Вы хотите рассчитать условное среднее значение за Window:

from pyspark.sql import Window
from pyspark.sql.functions import col, when, avg


w = Window.partitionBy('FORWARD_LOOK_MODEL', 'FOR_PD_TYPE')
account_level_pd_shifts = account_level_pd_shifts.withColumn("SHIFTED_PD",
                                                              avg(when(col("WITHIN_PD_EXCLUSION"),
                                                                       col("PD_SHIFT")
                                                                   )
                                                              ).over(w)
                                                          )