PySpark withColumn that uses column data from another row - pyspark

Question

I have a dataframe like this:
order_type
customer_id
order_id
related_order_id
purchase
123
abc
null
return
123
bcd
null
purchase
234
xyz
null
return
234
zzz
null
Where I want to fill in the related_order_id column as the order_id of the related purchase, only for rows where order_type is return. A return and a purchase row can be related by their customer_id.
I've tried to use withColumn(), but I haven't figured out a way that would allow me to also look at other rows and their column data.
The end result should look something like
order_type
customer_id
order_id
related_order_id
purchase
123
abc
null
return
123
bcd
abc
purchase
234
xyz
null
return
234
zzz
xyz

Answer 1

You can use the lag() function to use data from the previous row.
Assuming a return is always preceded by a purchase, you can do:
from pyspark.sql.window import Window
from pyspark.sql import functions as F
from pyspark.sql.functions import col
w = Window().partitionBy("customer_id").orderBy("order_type")
df.withColumn("related_order_id", F.when(col("order_type")=="return", \
F.lag(col("order_id")).over(w)) \
.otherwise(col("related_order_id"))).show()
Output:
+----------+-----------+--------+----------------+
|order_type|customer_id|order_id|related_order_id|
+----------+-----------+--------+----------------+
| purchase| 123| abc| null|
| return| 123| bcd| abc|
| purchase| 234| xyz| null|
| return| 234| zzz| xyz|
+----------+-----------+--------+----------------+