pyspark check whether each name has 3 data

user_16477433 0
Viewed 40

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.

pyspark apache-spark-sql
1 Answers

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")
    )
)
user_50137520
Related