How to sum diagonal values in Pyspark dataframe

This would be fairly trivial in pandas using shift so if your data is small enough it might just be easier to write a UDF and pass the entire dataset to it.

In any case, you can use lag in pyspark. You'll need to create a window function, which requires a grouping and ordering, which you really want to be a single group, and the order to be basically the same order the data is already in. To solve this you can assign everything to group 1 and order by monotonically_increasing_id.

Then create lags of your columns, starting at 0 and decreasing one per column. Then you can sum these up.

import pandas as pd
import numpy as np
import pyspark.sql.functions as F
from pyspark.sql.window import Window
from functools import reduce
from operator import add

df = pd.DataFrame(pd.DataFrame(np.random.randint(0,100,size=(5, 5)), columns=['col1','col2','col3','col4','col5']))
df = spark.createDataFrame(df)



df = df.withColumn('group',F.lit(1))
df = df.withColumn('order', F.monotonically_increasing_id())
w = Window.partitionBy(F.col('group')).orderBy(F.col('order'))


df = df.withColumn('lag1',F.lag(F.col('col1'),0).over(w)) \
       .withColumn('lag2',F.lag(F.col('col2'),-1).over(w)) \
       .withColumn('lag3',F.lag(F.col('col3'),-2).over(w)) \
       .withColumn('lag4',F.lag(F.col('col4'),-3).over(w)) \
       .withColumn('lag5',F.lag(F.col('col5'),-4).over(w))


display(df.withColumn('diagonal_sum',reduce(add, [F.col(x) for x in df.columns if 'lag' in x])))

Output

        col1    col2    col3    col4    col5    group         order  lag1   lag2    lag3    lag4    lag5    diagonal_sum
0         94      44      91      91      39        1   25769803776    94   96.0    19.0    56.0    28.0    293.0
1         77      96       0      25      70        1   51539607552    77   94.0    76.0    7.0 NaN NaN
2         87      94      19      76      63        1   77309411328    87   26.0    67.0    NaN NaN NaN
3         68      26      76      56      34        1   103079215104   68   32.0    NaN NaN NaN NaN
4         78      32      67       7      28        1   128849018880   78   NaN NaN NaN NaN NaN

Similar to Chris's answer , you can use reduce with lead after coalescing the nulls as below:

from pyspark.sql import functions as F, Window as W
from functools import reduce
w = W.orderBy(order_col)

order_col = 'w'
othercols = df.columns[1:]    
out = df.withColumn("DiagSum",reduce(lambda x,y: F.round(x+y,2),
[F.coalesce(F.lead(F.col(col),e).over(w),F.lit(0)) for e,col in enumerate(othercols)]))

out.show()

+---+-----+-----+-----+-----+-----+-------+
|  w| Col1| Col2| Col3| Col4| Col5|DiagSum|
+---+-----+-----+-----+-----+-----+-------+
| w1|161.8|161.8|161.8|161.8|161.8|  771.0|
| w2|112.0|112.0|112.0|112.0|112.0|  653.9|
| w3|157.7|157.7|157.7|157.7|157.7|  647.2|
| w4|146.1|146.1|146.1|146.1|146.1|  489.5|
| w5|193.4|193.4|193.4|193.4|193.4|  343.4|
| w6| 44.7| 44.7| 44.7| 44.7| 44.7|  150.0|
| w7|105.3|105.3|105.3|105.3|105.3|  105.3|
+---+-----+-----+-----+-----+-----+-------+