Flatten Nested JSON using PySpark

Question

I have a nested JSON that Im able to fully flatten by using the below function

# Flatten nested df  
def flatten_df(nested_df):
    for col in nested_df.columns:
        array_cols = [ c[0]  for c in nested_df.dtypes if c[1][:5] == 'array']
    for col in array_cols:
        nested_df =nested_df.withColumn(col, F.explode_outer(nested_df[col]))

    nested_cols = [c[0] for c in nested_df.dtypes if c[1][:6] == 'struct']

    if len(nested_cols) == 0:
        return nested_df

    flat_cols = [c[0] for c in nested_df.dtypes if c[1][:6] != 'struct']

    flat_df = nested_df.select(flat_cols +
                        [F.col(nc+'.'+c).alias(nc+'_'+c)
                            for nc in nested_cols
                            for c in nested_df.select(nc+'.*').columns])

    return flatten_df(flat_df)

I want to explode the nested structure but do not want to flatten all the way. I want flatten only to the first level and keep the subsequent nested structures as it is.

Here is my schema of the dataframe Im working with.

root
 |-- module: array (nullable = true)
 |    |-- element: array (containsNull = true)
 |    |    |-- element: struct (containsNull = true)
 |    |    |    |-- chart: array (nullable = true)
 |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |-- header: array (nullable = true)
 |    |    |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |    |    |-- alt: string (nullable = true)
 |    |    |    |    |    |    |    |-- assetId: string (nullable = true)
 |    |    |    |    |    |    |    |-- header: string (nullable = true)
 |    |    |    |    |    |    |    |-- height: string (nullable = true)
 |    |    |    |    |    |    |    |-- linkedid: string (nullable = true)
 |    |    |    |    |    |    |    |-- selected: boolean (nullable = true)
 |    |    |    |    |    |    |    |-- src: string (nullable = true)
 |    |    |    |    |    |    |    |-- styleCodes: string (nullable = true)
 |    |    |    |    |    |    |    |-- viewLarger: string (nullable = true)
 |    |    |    |    |    |    |    |-- width: string (nullable = true)
 |    |    |    |    |    |-- id: string (nullable = true)
 |    |    |    |    |    |-- row: array (nullable = true)
 |    |    |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |    |    |-- alt: string (nullable = true)
 |    |    |    |    |    |    |    |-- label: string (nullable = true)
 |    |    |    |    |    |    |    |-- value: array (nullable = true)
 |    |    |    |    |    |    |    |    |-- element: string (containsNull = true)
 |    |    |    |-- header: array (nullable = true)
 |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |-- decorators: array (nullable = true)
 |    |    |    |    |    |    |-- element: string (containsNull = true)
 |    |    |    |    |    |-- id: string (nullable = true)
 |    |    |    |    |    |-- value: string (nullable = true)
 |    |    |    |-- id: string (nullable = true)
 |    |    |    |-- image: array (nullable = true)
 |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |-- alt: string (nullable = true)
 |    |    |    |    |    |-- assetId: string (nullable = true)
 |    |    |    |    |    |-- id: string (nullable = true)
 |    |    |    |    |    |-- originalSrc: string (nullable = true)
 |    |    |    |    |    |-- src: string (nullable = true)
 |    |    |    |    |    |-- styleCodes: string (nullable = true)
 |    |    |    |    |    |-- viewLarger: boolean (nullable = true)
 |    |    |    |-- paragraph: array (nullable = true)
 |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |-- decorators: array (nullable = true)
 |    |    |    |    |    |    |-- element: array (containsNull = true)
 |    |    |    |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |    |    |    |-- style.bold: struct (nullable = true)
 |    |    |    |    |    |    |    |    |    |-- length: long (nullable = true)
 |    |    |    |    |    |    |    |    |    |-- offset: long (nullable = true)
 |    |    |    |    |    |-- id: string (nullable = true)
 |    |    |    |    |    |-- value: array (nullable = true)
 |    |    |    |    |    |    |-- element: string (containsNull = true)

The final data frame I want is

module_id   | module_header |   paragraph_id |  paragraph_value | image_id | image_src| chart_id |  chart_header |  chart_row |

Here 'module_id' is the 'id' under module array, module_header is the header array under module array and so on. I need to stop the flattening at the first level and not go all the way to the end.

Flatten Nested JSON using PySpark

0 Answers