Flatten Spark Dataframe

flatten_spark_dataframe.py

# source: https://stackoverflow.com/a/50156142/2833774

def flatten(schema, prefix=None):
    fields = []
    for field in schema.fields:
        name = prefix + '.' + field.name if prefix else field.name
        alias_name = name.replace(".", "__")
        dtype = field.dataType
        if isinstance(dtype, pst.ArrayType):
            dtype = dtype.elementType

        if isinstance(dtype, pst.StructType):
            fields += flatten(dtype, prefix=name)
        else:
            fields.append(col(name).alias(alias_name))

    return fields