How to enrich dataframe by adding columns in specific condition in pyspark?

user_13161972 0
Viewed 105

I have a two different dataframes:

users:

+-------+---------+--------+
|user_id| movie_id|timestep|
+-------+---------+--------+
|   100 |   1000  |20200728|
|   101 |   1001  |20200727|
|   101 |   1002  |20200726|
+-------+---------+--------+

movies:

+--------+---------+--------------------------+
|movie_id|  title  |         genre            |
+--------+---------+--------------------------+
|   1000 |Toy Story|Adventure|Animation|Chil..|
|   1001 | Jumanji |Adventure|Children|Fantasy|
|   1002 | Iron Man|Action|Adventure|Sci-Fi   |
+--------+---------+--------------------------+

How to get a dataframe in the following format? So I can get user's taste profile for comparing different users by their similarity score?

+-------+---------+---------+---------+--------+-----+
|user_id|  Action |Adventure|Animation|Children|Drama|
+-------+---------+---------+---------+---------+----+
|   100 |    0    |    1    |    1    |   1    |  0  |
|   101 |    1    |    2    |    0    |   1    |  0  |
+-------+---------+---------+---------+--------+-----+
pyspark
1 Answers

First, you need to split your "genre" column.

from pyspark.sql import functions as F

movies = movies.withColumn("genre", F.explode(F.split("genre", '\|')))
# use \ in front of | because split use regex

then you join

user_movie = users.join(movies, on='movie_id')

and you pivot

user_movie.groupBy("user_id").pivot("genre").agg(F.count("*")).fillna(0).show()

+-------+------+---------+---------+--------+-------+------+
|user_id|Action|Adventure|Animation|Children|Fantasy|Sci-Fi|
+-------+------+---------+---------+--------+-------+------+
|    100|     0|        1|        1|       1|      0|     0|
|    101|     1|        2|        0|       1|      1|     1|
+-------+------+---------+---------+--------+-------+------+

FYI : Drama column does not appear because there is no drama "genre" in the movies dataframe. But with your full data, you will have one column for each genre.

user_50137520
Related