apply Window.partitionBy for two columns to get n-core dataset in pyspark

user_12055790 0
Viewed 312

I have a data set of 2M entries with user,item,rating information. I want to filter out data so that it includes items that are rated by at least 2 users and users that rated at least 2 items. I can get one constraint done using a window function but not sure how to get both done.

input:

user product rating
J p1 3
J p2 4
M p1 4
M p3 3
B p2 3
B p4 3
B p3 3
N p3 2
N p5 4

here is sample data.

from pyspark import SparkContext
from pyspark.sql import SparkSession
# Create Spark Context
sc = SparkSession.builder.master("local[*]")\
     .config("spark.jars.packages", "org.apache.spark:spark-avro_2.12:3.1.2")\
     .getOrCreate()

sampleData = (("J", "p1", 3), \
    ("J", "p2", 4),  \
    ("M", "p1", 4),   \
    ("M", "p3", 3),  \
    ("B", "p2", 3),  \
    ("B", "p4", 3),  \
    ("B", "p3", 3),  \
    ("N", "p3", 2),\
    ("N", "p5", 4) \
  )
 
columns= ["user", "product", "rating"]

df = sc.createDataFrame(data = sampleData, schema = columns)

desired output is,

user product rating
J p1 3
J p2 4
M p1 4
M p3 3
B p2 3
B p3 3

window function I used to fulfill "users that rated at least 2 items" is

from pyspark.sql import functions as F
from pyspark.sql.functions import  count, col
from pyspark.sql.window import Window

window = Window.partitionBy("user")

df.withColumn("count", F.count("rating").over(window))\
    .filter(F.col("count") >= 2).drop("count")
pyspark window-functions count group-by
3 Answers

How about the below?

df = spark.createDataFrame(data = sampleData, schema = columns)
df_p = df.groupBy('product').count().filter('count >= 2').select('product')
df = df.join(df_p, ['product'], 'inner')
df_u = df.select('user').groupBy('user').count().filter('count >= 
2').select('user')
df = df.join(df_u, ['user'], 'inner')

Gives below output:

user product rating
B p2 3
B p3 3
M p1 4
M p3 3
J p2 4
J p1 3
user_50773950

You can do that with two window functions. I am less familiar with the df syntax, here is the sql:

df.createOrReplaceTempView("ratings")

spark.sql("""
SELECT USER,
       product,
       rating,
       Count(*)OVER (partition BY USER )    num_ratings_for_user,
       Count(*)OVER (partition BY product ) num_raters_for_product
FROM   ratings 
""")

You can filter over this.

user_1806500 
from pyspark.sql import functions as F
from pyspark.sql.window import Window

window1 = Window.partitionBy("user")
window2 = Window.partitionBy("product")

df.withColumn("count_users", F.count("rating").over(window1))\
  .filter(F.col("count_users") >= 2)\
  .withColumn("count_prod", F.count("rating").over(window2))\
  .filter(F.col("count_prod") >= 2)\
  .drop("count", "count_users", "count_prod")\
  .show()

User N has rated more than 1 product so the output should be:

+----+-------+------+
|user|product|rating|
+----+-------+------+
|   J|     p1|     3|
|   M|     p1|     4|
|   B|     p2|     3|
|   J|     p2|     4|
|   B|     p3|     3|
|   M|     p3|     3|
|   N|     p3|     2|
+----+-------+------+
user_116985620
Related