Shifting a PySpark dataframe column by a variable value in another column - apache-spark

Question

I have a PySpark dataframe that looks like this
Date
Value
Shift_Index
2021/02/11
50.12
0
2021/02/12
72.30
4
2021/02/15
81.87
1
2021/02/16
90.12
2
2021/02/17
91.31
1
2021/02/18
81.23
2
2021/02/19
73.45
1
2021/02/22
87.17
0
I want to lead the offset (On the basis of the values in the Shift_Index column here) which I have to pass depends on a particular Column of type Integer.
Can we somehow use an offset value that depends on the column value in lead/lag function in spark SQL ?
I wanted somewhat like this, which works fine in SQL server, but unfortunately throws exception in Spark SQL.
Create table test_table(ID int identity(1,1), Value float, shift_col int, New_Value float)
SELECT Value, shift_col,
ISNULL(LEAD(Value, shift_col) OVER(ORDER BY ID ASC), Value) AS New_Value
FROM test_table
The final result that I need looks something similar to this :
Date
Value
Shift_Index
New_Value
2021/02/11
50.12
0
50.12
2021/02/12
72.30
4
81.23
2021/02/15
81.87
1
90.12
2021/02/16
90.12
2
81.23
2021/02/17
91.31
1
81.23
2021/02/18
81.23
2
87.17
2021/02/19
73.45
1
87.17
2021/02/22
87.17
0
87.17
The following exceptions are encountered
Py4JJavaError: An error occurred while calling o77.sql.
: org.apache.spark.sql.AnalysisException: cannot resolve 'lead(sample_data_temp.shift_col, NULL)' due to data type mismatch: Offset expression 'shift_col#2835' must be a literal
Any help will be really appreciated.
Thanks in advance.

Answer 1

You could do so with a window and lead. If you have very spread out values for Shift_index you could do a select distinct to determine which shifts you need instead of everything up to the max shift.
Ideally you have something to partition your window by otherwise this can be very heavy for large datasets. Spark provides a warning:
WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.
Edit: solution without join, still not partitioning means this doesn't parallelize well.
from pyspark.sql import functions as f
from pyspark.sql.window import Window
w = Window().orderBy(f.col('Date'))
max_shift = df.agg(f.max('Shift_index')).collect()[0][0]
for shift in range(1, max_shift+1):
df = df.withColumn('Value' + str(shift), f.lead(f.col('Value'), shift).over(w))
case_shift = 'CASE Shift_index WHEN 0 THEN Value ' + ' '.join([f'WHEN {i} THEN Value{i}' for i in range(1, max_shift + 1)]) + ' ELSE NULL END'
df = df.select(
f.col('Date'),
f.col('Shift_index'),
f.col('Value'),
f.expr(case_shift).alias('New_Value')
)
df.show()
+----------+-----+-----------+---------+
| Date|Value|Shift_index|New_Value|
+----------+-----+-----------+---------+
|2021/02/11|50.12| 0| 50.12|
|2021/02/12| 72.3| 4| 81.23|
|2021/02/15|81.87| 1| 90.12|
|2021/02/16|90.12| 2| 81.23|
|2021/02/17|91.31| 1| 81.23|
|2021/02/18|81.23| 2| 87.17|
|2021/02/19|73.45| 1| 87.17|
|2021/02/22|87.17| 0| 87.17|
+----------+-----+-----------+---------+