Flatten nested dataframe and reconvert to original form

user_14171615 0
Viewed 45

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

  1. So is there a better way to do this in pyspark?
  2. Can we dynamically convert the dataframe to its original form after transformation?
  3. Can we separate the rows without flattening them?

I am new to spark and finding it very difficult to get this working

Thanks

parquet pyspark
1 Answers

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]]          |
+----------------------------------------------+
user_21298010
Related