transform a spark dataframe based on the value of a column - apache-spark

Question

Suppose I have the dataframe below:
s.id
s.first
s.last
s.age
d.first
d.last
d.age
UPDATED_FIELDS
1
AAA
BBB
10
AAA__
BBB
10
["first"]
2
CCC
DDD
20
CCC__
DDD
21
["first", "age"]
I want to transform it to the format below, so for each of the UPDATED_FIELDS in the first dataframe, I want to create a new row in my second dataframe.
id
field
s_value
d_value
1
first
AAA
AAA_
2
first
CCC
CCC_
2
age
20
21
I feel like I need to create a new dataframe, but couldn't get it working

Answer 1

This isn't straightforward because after you explode "UPDATED_FIELDS" to create the "field" column, you have the following dataframe:
+----+-------+------+-----+-------+------+-----+--------------+-----+
|s.id|s.first|s.last|s.age|d.first|d.last|d.age|UPDATED_FIELDS|field|
+----+-------+------+-----+-------+------+-----+--------------+-----+
| 1| AAA| BBB| 10| AAA__| BBB| 10| [first]|first|
| 2| CCC| DDD| 20| CCC__| DDD| 21| [first, age]|first|
| 2| CCC| DDD| 20| CCC__| DDD| 21| [first, age]| age|
+----+-------+------+-----+-------+------+-----+--------------+-----+
The new columns s_value and d_value depend on using the literal from the exploded column "field" to tell us which column's value to use. It would be nice to do something like:
df.withColumn(
"field", F.explode("UPDATED_FIELDS")
).withColumn(
"s_value", F.col(f"s.{F.col('field')}")
)
... but this will result in an error because {F.col('field')} cannot be interpreted as a literal.
Update: based on this helpful answer, you can in fact use a string literal as a column lookup. You will want to first make sure you can access df['field'] so you need to perform the explode separately first, then use F.coalesce(*F.when(...)) as a way of obtaining the value inside the column name that comes from the string inside df['field']:
df = df.withColumn(
"field", F.explode("UPDATED_FIELDS")
)
## obtain list of all unique strings in "field"
base_cols = df.select(F.collect_set("field").alias("column")).first()["column"]
df.withColumn(
"s_value", F.coalesce(*[F.when(df['field'] == c, df[f"`s.{c}`"]) for c in base_cols])
).withColumn(
"d_value", F.coalesce(*[F.when(df['field'] == c, df[f"`d.{c}`"]) for c in base_cols])
).select(
F.col("`s.id`").alias("id"),"field","s_value","d_value"
)
+---+-----+-------+-------+
| id|field|s_value|d_value|
+---+-----+-------+-------+
| 1|first| AAA| AAA__|
| 2|first| CCC| CCC__|
| 2| age| 20| 21|
+---+-----+-------+-------+