Conditionally adding a new column in pyspark - pyspark

Question

I have some example data that originated from json that looks similar to the below:
{ hero: axe, attribute: strength, active_abilities: [q, w, r], inactive_abilities: e }
{ hero: invoker, attribute: intelligence, active_abilities: [q, w, e, r, f, d], inactive_abilities: null }
{ hero: phantom assassin, attribute: agility, active_abilities: [q, w, e], inactive_abilities: r }
{ hero: life stealer, attribute: strength, active_abilities: [q, r], inactive_abilities: [w, e] }
The issue I'm having is that column 'inactive_abilities' is being read as a string due to the variability in the data types that can be present in that column. It is possible for the data to be null, a single string (if only 1 ability), an array (if multiple abilites). What I want in the end is to have several new columns based on the number of 'inactive_abilities'. If there is only 1 or null ability, I want a new column inactive_ability which will only be populated if there is one inactive ability and null if there are none or if there are multiple inactive abilities. Then I would like multiple columns like inactive_ability1, inactive_ability2, inactive_ability3, etc... in the case that the array holds > 1 value. So from the example above, the end result should look like:
{ hero: axe, attribute: strength, active_abilities: [q, w, r], inactive_abilities: e , inactive_ability: e, inactive_ability1: null, inactive_ability2: null, inactive_ability3, null, inactive_ability4: null}
{ hero: invoker, attribute: intelligence, active_abilities: [q, w, e, r, f, d], inactive_abilities: null, inactive_ability: null, inactive_ability1: null, inactive_ability2: null, inactive_ability3, null, inactive_ability4: null }
{ hero: phantom assassin, attribute: agility, active_abilities: [q, w, e], inactive_abilities: r, inactive_ability: r, inactive_ability1: null, inactive_ability2: null, inactive_ability3, null, inactive_ability4: null }
{ hero: life stealer, attribute: strength, active_abilities: [q, r], inactive_abilities: [w, e], inactive_ability: null, inactive_ability1: w, inactive_ability2: e, inactive_ability3, null, inactive_ability4: null }
I can't assume there will be a fixed number of 'inactive_abilities' but if it exceeds 4, the rest can be ignored. The part I'm having trouble is being able to cast the field into an array and reading it as such when appropriate, and then creating and populating the new columns based on the conditions mentioned above.

Answer 1

The string inactive_abilities column can be converted to array by removing the [ and ] and splitting the string by ,.
from pyspark.sql import functions as F
data = [{"hero": "axe", "attribute": "strength", "active_abilities": ["q", "w", "r"], "inactive_abilities": "e"},
{"hero": "invoker", "attribute": "intelligence", "active_abilities": ["q", "w", "e", "r", "f", "d"],
"inactive_abilities": None},
{"hero": "phantom assassin", "attribute": "agility", "active_abilities": ["q", "w", "e"],
"inactive_abilities": "r"},
{"hero": "life stealer", "attribute": "strength", "active_abilities": ["q", "r"],
"inactive_abilities": "[w, e]"}, ]
df = spark.createDataFrame(data)
array_select_expr = [
F.when(F.size("parsed_inactive_abilities") > 1, F.trim(F.col("parsed_inactive_abilities")[i])).alias(
f"inactive_ability{i}") if i > 0 else F.when(
F.size("parsed_inactive_abilities") == 1, F.trim(F.col("parsed_inactive_abilities")[0])).alias(
"inactive_abilities")
for i in range(0, 5)]
(df.withColumn("parsed_inactive_abilities", F.split(F.regexp_replace(F.col("inactive_abilities"), "[\[\]]", ""), ","))
.select("*", *array_select_expr)
.drop("parsed_inactive_abilities").show())
Output
+------------------+------------+----------------+------------------+------------------+-----------------+-----------------+-----------------+-----------------+
| active_abilities| attribute| hero|inactive_abilities|inactive_abilities|inactive_ability1|inactive_ability2|inactive_ability3|inactive_ability4|
+------------------+------------+----------------+------------------+------------------+-----------------+-----------------+-----------------+-----------------+
| [q, w, r]| strength| axe| e| e| null| null| null| null|
|[q, w, e, r, f, d]|intelligence| invoker| null| null| null| null| null| null|
| [q, w, e]| agility|phantom assassin| r| r| null| null| null| null|
| [q, r]| strength| life stealer| [w, e]| null| e| null| null| null|
+------------------+------------+----------------+------------------+------------------+-----------------+-----------------+-----------------+-----------------+
# Value in array_select_expr
[Column<'CASE WHEN (size(parsed_inactive_abilities) = 1) THEN trim(parsed_inactive_abilities[0]) END AS inactive_abilities'>,
Column<'CASE WHEN (size(parsed_inactive_abilities) > 1) THEN trim(parsed_inactive_abilities[1]) END AS inactive_ability1'>,
Column<'CASE WHEN (size(parsed_inactive_abilities) > 1) THEN trim(parsed_inactive_abilities[2]) END AS inactive_ability2'>,
Column<'CASE WHEN (size(parsed_inactive_abilities) > 1) THEN trim(parsed_inactive_abilities[3]) END AS inactive_ability3'>,
Column<'CASE WHEN (size(parsed_inactive_abilities) > 1) THEN trim(parsed_inactive_abilities[4]) END AS inactive_ability4'>]