Pyspark - функция UDAF для группы по двум столбцам даты, UDAF для расчета RMSE между фактическими значениями и прогнозами.

Question

У меня есть данные в фрейме данных pyspark за следующие несколько лет. Week_start_dt - это когда я начал предсказание. И start_month - предыдущие 12 месяцев.

+--------------------+------------------+----------------------+----------------+
|     start_month    |     week_start_dt|           predictions|       actuals  |
+--------------------+------------------+----------------------+----------------+
|             2019-01|        2019-11-11|                    12|              11|
|             2018-12|        2019-11-11|                    13|              11|
|             2019-08|        2019-11-11|                     9|              11|
|             2019-11|        2019-11-11|                    12|              11|
|             2019-11|        2019-11-11|                  1970|            1440|
|             2019-11|        2019-11-11|                   478|             501|
+--------------------+------------------+----------------------+----------------+

Я хотел бы рассчитать RMSE с groupby на start_month и week_start_dt. Я думаю, что для этого потребуется определенная пользователем агрегатная функция. Что-то вроде этого в pandas: Python Кадр данных: вычисление R ^ 2 и RMSE с использованием Groupby для одного столбца

Я использую следующий код, чтобы просто подсчитать число для фактических данных и прогнозов по группам.

df_startmonth_week = actuals_compare.groupby('start_month', 'week_start_dt').agg(f.sum('predictions'), f.sum('actuals'))

Что я изменяю на своем этапе агрегации для вычисления СКО между прогнозами и фактическими данными? Нужен ли UDF для этого?

Вот пример конечной цели, которую я поставил в excel

| week_start_dt | start_month | RMSE |
|---------------|-------------|------|
| 20-01-2020    | 2019-02     | 2345 |
| 20-01-2020    | 2019-03     | 2343 |
| 20-01-2020    | 2019-04     | 2341 |
| 20-01-2020    | 2019-05     | 2100 |
| 20-01-2020    | 2019-06     | 1234 |

Answer 1

Я не вижу разницы с проблемой здесь , поэтому я адаптировал решение к именам переменных, которые немного отличаются:

import pyspark.sql.functions as psf

def compute_RMSE(expected_col, actual_col):

  rmse = old_df.withColumn("squarederror",
                           psf.pow(psf.col(actual_col) - psf.col(expected_col),
                                   psf.lit(2)
                           ))
  .groupby('start_month', 'week_start_dt')
  .agg(psf.avg(psf.col("squarederror")).alias("mse"))
  .withColumn("rmse", psf.sqrt(psf.col("mse")))

  return(rmse)


compute_RMSE("predictions", "actuals")

Скажите, пропустил ли я небольшое Разница в проблеме