pyspark check whether each name has 3 data - pyspark

Question

In pyspark, I have a DataFrame as follows. I want to check whether each name has 3 action data (0, 1, 2). If there are missing, add a new row, the score column is set to 0, and the other columns are unchanged(ex: str1, str2, str3).
+-----+--------+--------+--------+-------+-------+
| name| str1 | str2 | str3 | action| score |
+-----+--------+--------+--------+-------+-------+
| A | str_A1 | str_A2 | str_A3 | 0| 2|
| A | str_A1 | str_A2 | str_A3 | 1| 6|
| A | str_A1 | str_A2 | str_A3 | 2| 74|
| B | str_B1 | str_B2 | str_B3 | 0| 59|
| B | str_B1 | str_B2 | str_B3 | 1| 18|
| C | str_C1 | str_C2 | str_C3 | 0| 3|
| C | str_C1 | str_C2 | str_C3 | 1| 33|
| C | str_C1 | str_C2 | str_C3 | 2| 3|
+-----+--------+--------+--------+-------+-------+
For example, name B has no action 2, add a new row data as follows
+-----+--------+--------+--------+-------+-------+
| name| str1 | str2 | str3 | action| score |
+-----+--------+--------+--------+-------+-------+
| A | str_A1 | str_A2 | str_A3 | 0| 2|
| A | str_A1 | str_A2 | str_A3 | 1| 6|
| A | str_A1 | str_A2 | str_A3 | 2| 74|
| B | str_B1 | str_B2 | str_B3 | 0| 59|
| B | str_B1 | str_B2 | str_B3 | 1| 18|
| B | str_B1 | str_B2 | str_B3 | 2| 0|<---- new row data
| C | str_C1 | str_C2 | str_C3 | 0| 3|
| C | str_C1 | str_C2 | str_C3 | 1| 33|
| C | str_C1 | str_C2 | str_C3 | 2| 3|
+-----+--------+--------+--------+-------+-------+
It is also possible that there is only one row data for one name, and two new row data need to be added.
+-----+--------+--------+--------+-------+-------+
| name| str1 | str2 | str3 | action| score |
+-----+--------+--------+--------+-------+-------+
| A | str_A1 | str_A2 | str_A3 | 0| 2|
| A | str_A1 | str_A2 | str_A3 | 1| 6|
| A | str_A1 | str_A2 | str_A3 | 2| 74|
| B | str_B1 | str_B2 | str_B3 | 0| 59|
| B | str_B1 | str_B2 | str_B3 | 1| 18|
| B | str_B1 | str_B2 | str_B3 | 2| 0|
| C | str_C1 | str_C2 | str_C3 | 0| 3|
| C | str_C1 | str_C2 | str_C3 | 1| 33|
| C | str_C1 | str_C2 | str_C3 | 2| 3|
| D | str_D1 | str_D2 | str_D3 | 0| 45|
+-----+--------+--------+--------+-------+-------+
+-----+--------+--------+--------+-------+-------+
| name| str1 | str2 | str3 | action| score |
+-----+--------+--------+--------+-------+-------+
| A | str_A1 | str_A2 | str_A3 | 0| 2|
| A | str_A1 | str_A2 | str_A3 | 1| 6|
| A | str_A1 | str_A2 | str_A3 | 2| 74|
| B | str_B1 | str_B2 | str_B3 | 0| 59|
| B | str_B1 | str_B2 | str_B3 | 1| 18|
| B | str_B1 | str_B2 | str_B3 | 2| 0|
| C | str_C1 | str_C2 | str_C3 | 0| 3|
| C | str_C1 | str_C2 | str_C3 | 1| 33|
| C | str_C1 | str_C2 | str_C3 | 2| 3|
| D | str_D1 | str_D2 | str_D3 | 0| 45|
| D | str_D1 | str_D2 | str_D3 | 1| 0|<---- new row data
| D | str_D1 | str_D2 | str_D3 | 2| 0|<---- new row data
+-----+--------+--------+--------+-------+-------+
I am new to pyspark and don't know how to do this operation.
Thank you for your help.

Answer 1

Solution with a UDF
from pyspark.sql import functions as F, types as T
#F.udf(T.MapType(T.StringType(), T.IntegerType()))
def add_missing_values(values):
return {i: values.get(i, 0) for i in range(3)}
df = (
df.groupBy("name", "str1", "str2", "str3")
.agg(
F.map_from_entries(F.collect_list(F.struct("action", "score"))).alias("values")
)
.withColumn("values", add_missing_values(F.col("values")))
.select(
"name", "str1", "str2", "str3", F.explode("values").alias("action", "score")
)
)
df.show()
+----+------+------+------+------+-----+
|name| str1| str2| str3|action|score|
+----+------+------+------+------+-----+
| A|str_A1|str_A2|str_A3| 0| 2|
| A|str_A1|str_A2|str_A3| 1| 6|
| A|str_A1|str_A2|str_A3| 2| 74|
| B|str_B1|str_B2|str_B3| 0| 59|
| B|str_B1|str_B2|str_B3| 1| 18|
| B|str_B1|str_B2|str_B3| 2| 0|<---- new row data
| C|str_C1|str_C2|str_C3| 0| 3|
| C|str_C1|str_C2|str_C3| 1| 33|
| C|str_C1|str_C2|str_C3| 2| 3|
| D|str_D1|str_D2|str_D3| 0| 45|
| D|str_D1|str_D2|str_D3| 1| 0|<---- new row data
| D|str_D1|str_D2|str_D3| 2| 0|<---- new row data
+----+------+------+------+------+-----+
Full Spark solution :
df = (
df.groupBy("name", "str1", "str2", "str3")
.agg(
F.map_from_entries(F.collect_list(F.struct("action", "score"))).alias("values")
)
.withColumn(
"values",
F.map_from_arrays(
F.array([F.lit(i) for i in range(3)]),
F.array(
[F.coalesce(F.col("values").getItem(i), F.lit(0)) for i in range(3)]
),
),
)
.select(
"name", "str1", "str2", "str3", F.explode("values").alias("action", "score")
)
)