Flattening JSON records using PySpark
By: Shreyas M S
Tags: Data Engineering, Data Structures, Algorithms, Json

Step 2: The unnest_dict function unnests the dictionaries in the json_schema recursively and maps the hierarchical path to the field to the column name in the all_fields dictionary whenever it encounters a leaf node (check done in is_leaf function). Additionally, it also stored the path to the array-type fields in cols_to_explode set.

is_leaf

unnest_dict

Step 3: Next, all the leaf fields are obtained by checking if elements of all_fields start with any element in cols_to_explode and are stored in all_cols_in_explode_cols .

Step 4: Using all_cols_in_explode_cols , rest is calculated which contains fields directly accessible with or without the dot notation, using a simple set difference operation.

Step 5: Now, structure is computed using cols_to_explode that is used for step by step node traversal to get to the array-type fields.

get_structure

Step 6: Next, a BFS traversal is performed on structure to obtain the order in which the array explode has to take place and this order is stored in order class variable.

extract_order

Step 7: Finally, order and all_cols_in_explode_cols is used to get all mutually exclusive fields for each element in order . To do this, a bottom-to-top approach is used i.e. the order list is reversed and the leaf fields inside each of the fields in order are mapped and stored in bottom_to_top . Note that '.order_details' key in bottom_to_top has no elements it. If there were leaf nodes under it, those would be directly accessible and would appear in rest .

Combining all the functions, the class would look like this:

AutoFlatten class

To make use of the class variables to open/explode, this block of code is executed:

Here, the JSON records are read from the S3 path, and the global schema is computed. This schema is then passed while creating an object of the AutoFlatten class that initializes all class variables. When the compute function is called from the object of AutoFlatten class, the class variables are updated.

To open/explode, all first-level columns are selected with the columns in rest which haven’t appeared already. A counter is kept on the target names which counts the duplicate target column names. Any target column name having a count greater than 1 is renamed as <path_to_target_field> with each level separated by a > . All paths to those fields are added to the visited set of paths.

Then a check is done if order is empty or not. An empty order list means that there is no array-type field in the schema and vice-versa. In case, the order list isn’t empty, a traversal is done over each element in order and based on its type array/struct (only appears if an array-type field is a child of struct-type field), the columns are exploded/opened respectively and only those columns are selected which already haven’t appeared. Additionally, duplicate target column names are replaced by <path_to_target_field> with each level separated by a > and the paths to those fields are added to the visited set of paths.

Let's see what columns appear in final_df .

>>> final_df.columns

['city', 'country', 'house_number', 'street', 'pincode', 'state', 'email', 'id', 'name', 'discount', 'order_details>id', 'price', 'product_name', 'quantity', 'ordered_time']

Since id in the order_details field was a duplicate, it was renamed as order_details>id .

Looking at the counts of the initial dataframe df and final_df dataframe, we know that the array explode has occurred properly.

dataframe counts