How to create a dataframe based on the lookup dataframe and create mulitple column on dynamic and map values in the specific columns

user_8978191 0
Viewed 34

I have two dataframes one is the main and another one is the lookup dataframe. I need to achieve the third one in the customized form using pyspark. I need check the values in the column list_ids and check the match in the lookup dataframe and mark the count in the final dataframe. I have tried array intersect and array lookup but it is not working.

Main dataframe:

df = spark.createDataFrame([(123, [75319, 75317]), (212, [136438, 25274]), (215, [136438, 75317])], ("ID", "list_IDs"))
df.show()

+---+---------------+
| ID|       list_IDs|
+---+---------------+
|123| [75319, 75317]|
|212|[136438, 25274]|
|215|[136438, 75317]|
+---+---------------+

Lookup Dataframe:

df_2 = spark.createDataFrame([(75319, "Wheat", 20), (75317, "Rice", 10), (136438, "Jowar", 30), (25274, "Rajma", 40)], ("ID", "Material", "Count"))
df_2.show()

+------+--------+-----+
|    ID|Material|Count|
+------+--------+-----+
| 75319|   Wheat|   A|
| 75317|    Rice|   B|
|136438|   Jowar|   C|
| 25274|   Rajma|   D|
+------+--------+-----+

Need Resultant dataframe as

+---+---------------+------+------+-------+------+
| ID|       list_IDs|Wheat | Rice | Jowar | Rajma|
+---+---------------+------+------+-------+------+
|123| [75319, 75317]|    A|    B|    0 |     0|
|212|[136438, 25274]|    0|    0|    C |     D|
|215|[136438, 75317]|    0|    B|    C |    0 |
+---+---------------+------+------+-------+------+
pyspark apache-spark-sql apache-spark arrays python
1 Answers

You can join the two dataframes and then pivot:

import pyspark.sql.functions as F

df2 = df.join(
    df_2, 
    F.array_contains(df.list_IDs, df_2.ID)
).groupBy(df.ID, 'list_IDs').pivot('Material').agg(F.first('Count')).fillna(0)

result.show()
+---+---------------+-----+-----+----+-----+
| ID|       list_IDs|Jowar|Rajma|Rice|Wheat|
+---+---------------+-----+-----+----+-----+
|212|[136438, 25274]|   30|   40|   0|    0|
|215|[136438, 75317]|   30|    0|  10|    0|
|123| [75319, 75317]|    0|    0|  10|   20|
+---+---------------+-----+-----+----+-----+
user_141657300
Related