pyspark parent child recursive on same dataframe

Viewed 260

I have the following two Dataframes that stores diagnostic and part change for helicopter parts. diagnostic dataframe stores the maintenance activities carried out date. The part change dataframe stores all part removals for all the helicopter parts, parent(rotor), and child (turbofan, axle, module).

What I am trying to achieve is quite complex, based on the diagnostic df I want to provide me the first removal for the same part along with its parent roll all the way up to so that I get the helicopter serial no at that maintenance date.

Here the initial code to generate the sample datasets:

from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.sql import functions as F
from pyspark.sql import Window as W

diagnostic_list = [
  {"SN":"tur26666","PN": "turbofan", "maintenance":"23/05/2016"}
]
part_change_list= [
  {"SN":"tur26666","PN": "turbofan", "removal":"30/03/2019", "Parent_SN" : "com26666","Parent_PN": "compressor"},
{"SN":"tur26666","PN": "turbofan", "removal":"30/06/2016", "Parent_SN" : "com26666","Parent_PN": "compressor"},

 {"SN":"com26666","PN": "compressor", "removal":"13/03/2019", "Parent_SN" : "rot71777","Parent_PN": "rotorcraft"},
 {"SN":"com26666","PN": "compressor", "removal":"30/06/2016", "Parent_SN" : "rot26666","Parent_PN": "rotorcraft"},

 {"SN":"rot26666","PN": "rotorcraft", "removal":"31/12/2019", "Parent_SN" : "OYAAA","Parent_PN": "helicopter"},
  {"SN":"rot26666","PN": "rotorcraft", "removal":"24/06/2016", "Parent_SN" : "OYHZZ","Parent_PN": "helicopter"},    
]
spark = SparkSession.builder.getOrCreate()
diagnostic_df = spark.createDataFrame(Row(**x) for x in diagnostic_list)
part_change_df = spark.createDataFrame(Row(**x) for x in part_change_list)

diagnostic_df.show()


+--------+--------+-----------+
|      SN|      PN|maintenance|
+--------+--------+-----------+
|tur26666|turbofan| 23/05/2016|
+--------+--------+-----------+

part_change_df.show()


+--------+----------+----------+---------+----------+
|      SN|        PN|   removal|Parent_SN| Parent_PN|
+--------+----------+----------+---------+----------+
|tur26666|  turbofan|30/03/2019| com26666|compressor|
|tur26666|  turbofan|30/06/2016| com26666|compressor|
|com26666|compressor|13/03/2019| rot71777|rotorcraft|
|com26666|compressor|29/06/2016| rot26666|rotorcraft|
|rot26666|rotorcraft|31/12/2019|    OYAAA|helicopter|
|rot26666|rotorcraft|24/06/2016|    OYHZZ|helicopter|
+--------+----------+----------+---------+----------+

I was able to get the first removal for the child turbofan with the below code :

working_df = (
    diagnostic_df.join(part_change_df, ["SN", "PN"], how="inner")
    .filter(F.col("removal") >= F.col("maintenance"))
    .withColumn(
        "rank",
        F.rank().over(
            W.partitionBy([F.col(col) for col in ["SN", "PN", "maintenance"]]).orderBy(
                F.col("removal")
            )
        ),
    )
    .filter(F.col("rank") == 1)
    .drop("rank")
)
working_df.show()

+--------+--------+-----------+----------+---------+----------+
|      SN|      PN|maintenance|   removal|Parent_SN| Parent_PN|
+--------+--------+-----------+----------+---------+----------+
|tur26666|turbofan| 23/05/2016|30/06/2016| com26666|compressor|
+--------+--------+-----------+----------+---------+----------+

How can I create a for loop or a recursive loop within the part_change_df to get the results like this that takes each parent of the first child and makes it the next child and get the first removal information after the first child(turbofan)'s maintenance date)?

+--------+--------+-----------+----------+---------+----------+--------------+--------------+--------------+-------------------+-------------------+-------------------+
|      SN|      PN|maintenance|   removal|Parent_SN| Parent_PN|Parent_removal|next_Parent_SN|next_Parent_PN|next_Parent_removal|next_next_Parent_PN|next_next_Parent_SN|
+--------+--------+-----------+----------+---------+----------+--------------+--------------+--------------+-------------------+-------------------+-------------------+
|tur26666|turbofan| 23/05/2016|30/06/2016| com26666|compressor|    29/06/2016|      rot26666|    rotorcraft|         24/06/2016|         helicopter|              OYHZZ|
+--------+--------+-----------+----------+---------+----------+--------------+--------------+--------------+-------------------+-------------------+-------------------+

I could hardcode each parent and join working dataframe with the part change dataframe, but the problem i don't know exactly how high the number of parents a child will have . The ultimate goal is like to get the child maintenance date and roll up all the way to the final parent removal date and the helicopter serial no:

+--------+--------+-----------+----------+-------------------+-------------------+-------------------+
|      SN|      PN|maintenance|   removal|next_Parent_removal|next_next_Parent_PN|next_next_Parent_SN|
+--------+--------+-----------+----------+-------------------+-------------------+-------------------+
|tur26666|turbofan| 23/05/2016|30/06/2016|         24/06/2016|         helicopter|              OYHZZ|
+--------+--------+-----------+----------+-------------------+-------------------+-------------------+
0 Answers
Related