Program RMSE formula using Pyspark

Question

I'm struggling with programming the following formula using Pyspark:

I tried to do that in python it looked the following (still not 100% sure it's correct):

win_lean = 5000

df['U_rms'] = df.Current1.pow(2).rolling(min_periods=1, window=win_len).apply(lambda x: np.sqrt(x.mean()))

In pyspark I was trying several attitudes:

1. This one didn't work at all:

win_len = 5000
    
df = df.select("Current1").withColumn("U_rms", (F.pow("Current1", 2).Window( windowDuration=win_len).apply(lambda x: F.sqrt(x.mean()))))

2. Using @udf still seems like not a giving me a correct result:

import numpy as np
from pyspark.sql.types import FloatType
from pyspark.sql import functions as F

    @udf(returnType=FloatType())
    def MyFunc(value):
      win_len=5000
      return float(
        np.sqrt(
          np.sum([np.square(value) for i in range(1, win_len+1)]) / win_len
        )
      )
    df = df.withColumn("u_rms1", MyFunc(F.col("Current1")))

In this case I also prefere not to have null values in the beginning since windowing starts after the first 5000 samples. Thank you in advance.

Answer 1

I think you can use the window function to get rolling rmse.

from pyspark.sql import functions as f
from pyspark.sql import Window

df = spark.createDataFrame([[1], [4], [3], [5]], ['value'])
df.show()
+-----+
|value|
+-----+
|    1|
|    4|
|    3|
|    5|
+-----+

# from this to the next row, you can adjust like 5000.
w = Window.rowsBetween(0, 1)
​
df.withColumn('count', f.count(f.lit(1)).over(w)) \
  .withColumn('square', f.col('value') * f.col('value')) \
  .withColumn('mse', f.sum(f.col('square')).over(w) / f.col('count')) \
  .withColumn('rmse', f.sqrt(f.col('mse'))) \
  .show(truncate=False)
+-----+-----+------+----+------------------+
|value|count|square|mse |rmse              |
+-----+-----+------+----+------------------+
|1    |2    |1     |8.5 |2.9154759474226504|
|4    |2    |16    |12.5|3.5355339059327378|
|3    |2    |9     |17.0|4.123105625617661 |
|5    |1    |25    |25.0|5.0               |
+-----+-----+------+----+------------------+

If you want to do it for the all rows,

df.withColumn('count', f.count(f.lit(1)).over(w)) \
  .withColumn('square', f.col('value') * f.col('value')) \
  .agg(f.sum(f.col('square')).alias('mse')) \
  .withColumn('rmse', f.sqrt(f.col('mse'))) \
  .show(truncate=False)
+---+----------------+
|mse|rmse            |
+---+----------------+
|51 |7.14142842854285|
+---+----------------+

Answer 2

A variation to the given answers - we can create arrays of the required amount of data and use that for calculations.

Note that I've added an id column that has the dataframe's row order (spark does not retain row order) to keep the sequence intact. You can remove it if not required, or create it using monotonically_increasing_id function.

spark.createDataFrame([[1, 1], [2, 4], [3, 3], [4, 5], [5, 9]], ['id', 'data']). \
    withColumn('data_sq', func.pow('data', 2).cast('double')). \
    withColumn('list_of_points', 
               func.collect_list('data_sq').over(wd.orderBy('id').rowsBetween(0, 1))
               ). \
    withColumn('list_of_points_sum', 
               func.expr('aggregate(list_of_points, cast(0 as double), (x, y) -> x + y)')
               ). \
    withColumn('rmse', 
               func.sqrt(func.col('list_of_points_sum') / func.size('list_of_points'))
               ). \
    show(5, False)

# +---+----+-------+--------------+------------------+------------------+
# |id |data|data_sq|list_of_points|list_of_points_sum|rmse              |
# +---+----+-------+--------------+------------------+------------------+
# |1  |1   |1.0    |[1.0, 16.0]   |17.0              |2.9154759474226504|
# |2  |4   |16.0   |[16.0, 9.0]   |25.0              |3.5355339059327378|
# |3  |3   |9.0    |[9.0, 25.0]   |34.0              |4.123105625617661 |
# |4  |5   |25.0   |[25.0, 81.0]  |106.0             |7.280109889280518 |
# |5  |9   |81.0   |[81.0]        |81.0              |9.0               |
# +---+----+-------+--------------+------------------+------------------+

I've used .rowsBetween(0, 1) to consider current row (0) and next row (1). You can update the 1 to 4999 for 5k rows.

Answer 3

The best solution is still to use pandas API which provides the exact result that I want to:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from math import *
from pyspark.pandas.config import set_option
import pyspark.pandas as ps
%matplotlib inline

ps.set_option('compute.ops_on_diff_frames', True)

win_len = 5000

dfp = dfp.pandas_api()

dfp['u_rmse1'] = dfp.Current1.pow(2).rolling(min_periods=1, window=win_len).mean().apply(np.sqrt)

dfp.to_spark()

This code is completely understandable for a person who's programming in python. With pandas API it's also pretty fast.