Windowing Data into Rows in Pyspark - apache-spark

Question

I'm preparing a dataset to develop a supervised model to predict a value given the 5 previous values before it. For example given the sample data below, I would predict the 6th column given columns 1:5, or the 8th column given columns 3:7.
id 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ...
a 150 110 130 80 136 150 190 110 150 110 130 136 100 150 190 110
b 100 100 130 100 136 100 160 230 122 130 15 200 100 100 136 100
c 130 122 140 140 122 130 15 200 100 100 130 100 136 100 160 230
To that end, I want to reorganize the sample data above into rows of 6 columns, taking every slice/window of 6 values possible (e.g. 1:6, 2:7, 3:8). How can I do that? Is it possible in PySpark/SQL? Example of output below, index just for clarification:
1 2 3 4 5 6
a[1:6] 150 110 130 80 136 150
a[2:7] 110 130 80 136 150 190
a[3:8] 130 80 136 150 190 110
...
c[1:6] 130 122 140 140 122 130
c[2:7] 122 140 140 122 130 15
...
c[10:16] 130 100 136 100 160 230

Answer 1

You can convert your columns into an array of arrays or array of structs and then explode, for example:
from pyspark.sql.functions import struct, explode, array, col
# all columns except the first
cols = df.columns[1:]
# size of the splits
N = 6
Use array of arrays:
df_new = df.withColumn('dta', explode(array(*[ array(*cols[i:i+N]) for i in range(len(cols)-N+1) ]))) \
.select('id', *[ col('dta')[i].alias(str(i+1)) for i in range(N) ])
df_new.show()
+---+---+---+---+---+---+---+
| id| 1| 2| 3| 4| 5| 6|
+---+---+---+---+---+---+---+
| a|150|110|130| 80|136|150|
| a|110|130| 80|136|150|190|
| a|130| 80|136|150|190|110|
| a| 80|136|150|190|110|150|
| a|136|150|190|110|150|110|
| a|150|190|110|150|110|130|
| a|190|110|150|110|130|136|
| a|110|150|110|130|136|100|
| a|150|110|130|136|100|150|
| a|110|130|136|100|150|190|
| a|130|136|100|150|190|110|
| b|100|100|130|100|136|100|
+---+---+---+---+---+---+---+
Use array of structs (spark 2.4+):
df_new = df.withColumn('dta', array(*cols)) \
.selectExpr("id", f"""
inline(transform(sequence(0,{len(cols)-N}), i -> ({','.join(f'dta[i+{j}] as `{j+1}`' for j in range(N))})))
""")
the code inside above f-string is the same as the following for N=6:
inline(transform(sequence(0,10), i -> struct(dta[i] as `1`, dta[i+1] as `2`, dta[i+2] as `3`, dta[i+3] as `4`, dta[i+4] as `5`, dta[i+5] as `6`)))

Answer 2

Yes you can use this code (and modify it to get what you need):
partitions = []
for row in df.rdd.toLocalIterator():
row_list = list(row)
num_elements = 6
for i in range(0, len(row_list) - num_elements):
partition = row[i : i+num_elements]
partitions.append(partition)
output_df = spark.createDataFrame(partitions)