Replacing null with average in pyspark I have a problem during upsampling operation in PySpark. My dataframe is:

df_upsampled.show()
+-------------------+------------------+                                        
|                 et|           average|
+-------------------+------------------+
|2018-08-15 00:10:00| 4.165999948978424|
|2018-08-15 00:15:00|              null|
|2018-08-15 00:20:00|3.6580000072717667|
|2018-08-15 00:25:00|              null|
|2018-08-15 00:30:00|0.9999999925494194|


What I want to do is that by using Spark functions, replace the nulls in the "sum" column with the mean value of the previous and next variable in the "sum" column.

Wherever there is a null in column "sum", it should be replaced with the mean of the previous and next value in the same column "sum".

In this case, first null should be replaced by (4.16599 + 3.658)/2 = 3.91 and so on for the rest nulls.. 

What would be a good way to do this?
That way I found is to add two columns to the same dataframe; one lagging and one leading. The idea is that the two variables of which average is to be computed can this way be placed in one row. Like this:

+-------------------+-------------------+-------------------+-------------------+
|                 et|            average|         prev_value|         next_value|
+-------------------+-------------------+-------------------+-------------------+
|2018-08-14 04:10:00|0.11070000156760215|               null|               null|
|2018-08-14 04:15:00|               null|0.11070000156760215|0.08800000175833703|
|2018-08-14 04:20:00|0.08800000175833703|               null|               null|
|2018-08-14 04:25:00|               null|0.08800000175833703|0.10970000103116036|
|2018-08-14 04:30:00|0.10970000103116036|               null|               null|


Now, we can create a new dataframe from this such as wherever there is a null in column "average", it should take the average of the values from the same row of the next two columns.

After this, output will be like:

+-------------------+-------------------+-------------------+-------------------+
|                 et|            average|         prev_value|         next_value|
+-------------------+-------------------+-------------------+-------------------+
|2018-08-14 04:10:00|0.11070000156760215|               null|               null|
|2018-08-14 04:15:00|0.09935000166296959|0.11070000156760215|0.08800000175833703|
|2018-08-14 04:20:00|0.08800000175833703|               null|               null|
|2018-08-14 04:25:00| 0.0988500013947487|0.08800000175833703|0.10970000103116036|
|2018-08-14 04:30:00|0.10970000103116036|               null|               null|
|2018-08-14 04:35:00|0.11205000076442957|0.10970000103116036|0.11440000049769879|
|2018-08-14 04:40:00|0.11440000049769879|               null|               null|


You can now .drop() the columns prev_value and next_value to get clean output dataframe.

The code to this is:

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

my_window = Window.partitionBy().orderBy("et")
df = df.withColumn("prev_value", F.lag(df.average).over(my_window)).withColumn("next_value",F.lead(df.average).over(my_window))
df = df.withColumn("average", F.when(F.isnull(df.average),((F.col('prev_value')+F.col('next_value'))/2)).otherwise(df.average))
df.show()