Pyspark - How do I Flatten Nested Struct Column perserving parent name - pyspark

Question

If I have a dataframe with a struct column named structA, and in it we have 3 columns named a,b and c
if I want to flat the struct I can easily do that with df.select("structA.*") and it will display
a
b
c
1
2
3
3
5
6
What I wanted is
structA.a
structA.b
structA.c
1
2
3
3
5
6
How can I do this?

Answer 1

I'm afraid it's not straightforward as it should. You'll need to loop through the schema to get and build your desired column names, then rename columns in a bulk. Something like this
Sample dataset
df = spark.createDataFrame([
((1, 2, 3),),
((4, 5, 6),),
], 'structA struct<a:int, b:int, c:int>')
df.show()
df.printSchema()
+---------+
| structA|
+---------+
|{1, 2, 3}|
|{4, 5, 6}|
+---------+
root
|-- structA: struct (nullable = true)
| |-- a: integer (nullable = true)
| |-- b: integer (nullable = true)
| |-- c: integer (nullable = true)
from pyspark.sql import functions as F
struct_col = 'structA'
struct_cols = [[F.col(b.name).alias(f'{a.name}_{b.name}') for b in a.dataType.fields] for a in df.schema if a.name == struct_col][0]
# [Column<'a AS structA_a'>, Column<'b AS structA_b'>, Column<'c AS structA_c'>]
df.select(f'{struct_col}.*').select(struct_cols).show()
+---------+---------+---------+
|structA_a|structA_b|structA_c|
+---------+---------+---------+
| 1| 2| 3|
| 4| 5| 6|
+---------+---------+---------+

Answer 2

I did the following in order to do this:
from pyspark.sql import functions as F
from pyspark.sql.types import StructType,StructField, StringType, IntegerType
#create dataframe ----
df_data=[((1,2,3),(-1,-2,-3)),((4,5,6),(-4,-5,-6))]
structureSchema = StructType([
StructField('structA', StructType([
StructField('a', IntegerType(), True),
StructField('b', IntegerType(), True),
StructField('c', IntegerType(), True)
])),
StructField('structB', StructType([
StructField('a', IntegerType(), True),
StructField('b', IntegerType(), True),
StructField('c', IntegerType(), True)
]))
])
df=spark.createDataFrame(df_data,structureSchema)
df.show()
+---------+------------+
| structA| structB|
+---------+------------+
|[1, 2, 3]|[-1, -2, -3]|
|[4, 5, 6]|[-4, -5, -6]|
+---------+------------+
If we have multiple struct like columns we need to find them like this:
nested_cols = [c[0] for c in df.dtypes if c[1][:6] == 'struct']
nested_cols
['structA', 'structB']
Now I will create a json like object with the new structure
struct_columns={}
for struct_column in nested_cols:
struct_columns[struct_column]=df.select(struct_column+".*").columns
struct_columns
{'structA': ['a', 'b', 'c'], 'structB': ['a', 'b', 'c']}
With the structure, I will create the flattened data frame
flatten_df=df
for key in struct_columns:
print(key)
for column in struct_columns[key]:
flatten_df=flatten_df.withColumn(key+"_"+column,F.expr(f"{key}.{column}"))
flatten_df.drop(*df.columns).show()
+---------+---------+---------+---------+---------+---------+
|structA_a|structA_b|structA_c|structB_a|structB_b|structB_c|
+---------+---------+---------+---------+---------+---------+
| 1| 2| 3| -1| -2| -3|
| 4| 5| 6| -4| -5| -6|
+---------+---------+---------+---------+---------+---------+