Flatten nested dataframe and reconvert to original form - pyspark

Question

I have a parquet file with the following schema
root
|-- listOfMetrics: array (nullable = true)
| |-- element: struct (containsNull = true)
| | |-- Action: string (nullable = true)
| | |-- id: long (nullable = true)
| | |-- date: date (nullable = true)
| | |-- female_executives: double (nullable = true)
| | |-- male_executives: double (nullable = true)
| | |-- female_directors: double (nullable = true)
| | |-- male_directors: double (nullable = true)
| | |-- female_executives_and_directors: double (nullable = true)
| | |-- male_executives_and_directors: double (nullable = true)
| | |-- flag: integer (nullable = true)
and df.show() returns something like the following
+----------------------------------+
| listOfMetrics |
+----------------------------------+
| [[ADD, 5394, 2...|
| [[ADD, 527, 20...|
| [[ADD, 714, 20...|
| [[ADD, 765, 20...|
| [[ADD, 996, 20...|
| [[ADD, 146, 20...|
| [[ADD, 947, 20...|
+----------------------------------+
The 'Action' column is what I am targeting. This column can contain 'DELETE' or 'ADD' so based on this, I need to separate the rows.
The approach that I took was to flatten and then separate using pyspark.sql and reconvert back to its original form but failed at the conversion step.
I have the following questions
So is there a better way to do this in pyspark?
Can we dynamically convert the dataframe to its original form after transformation?
Can we separate the rows without flattening them?
I am new to spark and finding it very difficult to get this working
Thanks

Answer 1

It is possible to split each array into two parts using filter: one part containing only elements with ADD and one part containing only elements with DELETE. After that each part will become a separate line using stack. The original structure can be kept and flattening the dataframe is not necessary.
I am using a slightly simplified set of test data:
root
|-- listOfMetrics: array (nullable = true)
| |-- element: struct (containsNull = true)
| | |-- Action: string (nullable = true)
| | |-- flag: long (nullable = true)
| | |-- id: long (nullable = true)
+----------------------------------------------------------------------------------+
|listOfMetrics |
+----------------------------------------------------------------------------------+
|[[ADD, 2, 1], [ADD, 4, 3], [DELETE, 6, 5], [DELETE, 8, 7], [ADD, 10, 9]] |
|[[ADD, 12, 11], [ADD, 14, 13], [DELETE, 16, 15], [DELETE, 18, 17], [ADD, 110, 19]]|
+----------------------------------------------------------------------------------+
The code:
df.withColumn("listOfMetrics", F.expr("""stack(2,
filter(listOfMetrics, e -> e.Action = 'ADD'),
filter(listOfMetrics, e -> e.Action = 'DELETE'))""")) \
.show(truncate=False)
Output:
+----------------------------------------------+
|listOfMetrics |
+----------------------------------------------+
|[[ADD, 2, 1], [ADD, 4, 3], [ADD, 10, 9]] |
|[[DELETE, 6, 5], [DELETE, 8, 7]] |
|[[ADD, 12, 11], [ADD, 14, 13], [ADD, 110, 19]]|
|[[DELETE, 16, 15], [DELETE, 18, 17]] |
+----------------------------------------------+