Create item- item Interaction Matrix in Pyspark - pyspark

Question

I have a dataset containing two columns, user_id and item_id. The DataFrame looks like this:
index user_id item_id
0 user1 A
1 user1 B
2 user2 A
3 user3 B
4 user4 C
I'm looking for a way to transform this table into an item-item interaction matrix where we have distinct intersection of common users between items:
A B C
A 2 1 0
B 1 2 0
C 0 0 1
And another item-item interaction matrix where we have distinct union of users between items:
A B C
A 2 3 3
B 3 2 3
C 3 3 1

Answer 1

Step 0. Define the dataframe
import pyspark.sql.functions as F
data = [(0, "user1", "A"),
(1, "user1", "B"),
(2, "user2", "A"),
(3, "user3", "B"),
(4, "user4", "C")]
df = spark.createDataFrame(data, schema=["index", "user_id", "item_id"])
Step 1. Collect user data for each item in df_collect
df_collect = (df
.select("user_id", "item_id")
.groupBy("item_id")
.agg(F.collect_set("user_id").alias("users")))
Step 2. Cross join df_collect with itself to get all item-item combinations
df_crossjoin = (df_collect
.join(df_collect
.withColumnRenamed("item_id", "item_y")
.withColumnRenamed("users", "users_y")))
Step 2. Find user union and intersection and the count
df_ui = (df_crossjoin
.withColumn("users_union",
F.size((F.array_union("users", "users_y"))))
.withColumn("users_intersect",
F.size(F.array_intersect("users", "users_y"))))
Step 3. Pivot to get item-item matrix
df_matrix_union = (df_ui
.groupBy("item_id")
.pivot("item_y")
.agg(F.first("users_union"))
.orderBy("item_id"))
df_matrix_intrsct = (df_ui
.groupBy("item_id")
.pivot("item_y")
.agg(F.first("users_intersect"))
.orderBy("item_id"))