PySpark coalesce struct fields inside a struct - pyspark

Question

I have a struct coming from a data source where the struct fields have multiple possible data types like the following:
|-- priority: struct (nullable = true)
| |-- priority_a: struct (nullable = true)
| | |-- union: boolean (nullable = true)
| | |-- int32: integer (nullable = true)
| | |-- double: double (nullable = true)
| |-- priority_b: integer (nullable = true)
| |-- priority_c: struct (nullable = true)
| | |-- union: boolean (nullable = true)
| | |-- double: double (nullable = true)
| | |-- int32: integer (nullable = true)
| |-- priority_d: struct (nullable = true)
| | |-- union: boolean (nullable = true)
| | |-- double: double (nullable = true)
| | |-- int32: integer (nullable = true)
| |-- priority_e: double (nullable = true)
I want to coalesce the struct fields and cast them to a data type which makes the most sense, for instance:
|-- priority: struct (nullable = true)
| |-- priority_a: integer (nullable = true)
| |-- priority_b: integer (nullable = true)
| |-- priority_c: double (nullable = true)
| |-- priority_d: double (nullable = true)
| |-- priority_e: double (nullable = true)
If a column is not a struct field inside a struct, the following code works perfectly for what I need:
try:
cols = [f'{c}.{col}' for col in source.select(f'{c}.*').columns]
if f'{struct_path}.union' in cols:
cols.remove(f'{struct_path}.union')
source = source.withColumn(pc, f.coalesce(*cols).cast(t)) # t is the type I want to cast to
except:
source = source.withColumn(c, f.col(c).cast(t))
I would like to the do the same recursively for a struct where the nested struct fields can have multiple data types. Is it possible to do so?

Answer 1

StructField's fields are accessible by fields property, so what you can do is you can make a loop go through the schema and check every field to see if it's StructType
from pyspark.sql import types as T
for field in schema.fields:
if isinstance(field.dataType, T.StructType):
print(field.dataType.fields)
Or if you want to read it recursively
def flatten(schema, prefix=None):
fields = []
for field in schema.fields:
name = prefix + '.' + field.name if prefix else field.name
dtype = field.dataType
if isinstance(dtype, T.ArrayType):
dtype = dtype.elementType
if isinstance(dtype, T.StructType):
print(dtype)
fields += flatten(dtype, prefix=name)
else:
fields.append((dtype, name))
return fields