Pyspark使用Window.partitionBy双列分组筛选2-core用户商品评分数据集

要同时满足「用户至少为2件商品打分」「商品至少被2名用户评分」两个约束，不能仅做一次单维度过滤，因为过滤其中一个维度后，另一个维度的计数会发生变化，可能导致原本满足条件的条目不再符合要求。你可以通过多轮交替过滤的方式实现，200万条的数据量很小，不管用groupby还是窗口函数都不会有性能问题。

方法1：基于分组过滤的实现（更直观） #

from pyspark.sql import functions as F

# 迭代过滤直到数据稳定，避免极端场景下不收敛
prev_count = -1
current_df = df
while prev_count != current_df.count():
    prev_count = current_df.count()
    # 过滤符合条件的用户
    valid_user = current_df.groupBy("user").agg(F.count("*").alias("cnt")).filter(F.col("cnt") >= 2).select("user")
    current_df = current_df.join(valid_user, on="user", how="inner")
    # 过滤符合条件的商品
    valid_product = current_df.groupBy("product").agg(F.count("*").alias("cnt")).filter(F.col("cnt") >= 2).select("product")
    current_df = current_df.join(valid_product, on="product", how="inner")

# 输出结果
current_df.orderBy("user", "product").show()

跑样例数据的话，最终输出就是你要的结果：

+----+-------+------+
|user|product|rating|
+----+-------+------+
|   B|     p2|     3|
|   B|     p3|     3|
|   J|     p1|     3|
|   J|     p2|     4|
|   M|     p1|     4|
|   M|     p3|     3|
+----+-------+------+

方法2：基于窗口函数的实现 #

如果你习惯用窗口函数，也可以调整为多轮窗口过滤：

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

prev_count = -1
current_df = df
user_window = Window.partitionBy("user")
prod_window = Window.partitionBy("product")

while prev_count != current_df.count():
    prev_count = current_df.count()
    # 过滤用户
    current_df = current_df.withColumn("user_cnt", F.count("rating").over(user_window))\
                            .filter(F.col("user_cnt") >=2).drop("user_cnt")
    # 过滤商品
    current_df = current_df.withColumn("prod_cnt", F.count("rating").over(prod_window))\
                            .filter(F.col("prod_cnt") >=2).drop("prod_cnt")

current_df.orderBy("user", "product").show()

两种方法的效率没有明显差异，你可以根据自己的习惯选择。

内容的提问来源于stack exchange，提问作者krishthw