PySpark：如何基于CoordinateMatrix的相似度获取Top-K物品ID

你已经完成了余弦相似度的核心计算，得到了包含物品对相似度的dfsim DataFrame，接下来提取每个物品的Top-K相似物品，用Spark的窗口函数就能高效搞定——毕竟你的数据量不小，得尽量用分布式操作避免性能瓶颈。

步骤1：过滤物品自身的相似度记录 #

首先要排除物品和自身的相似结果（自己和自己的相似度为1，但这不是我们要的"相似物品"）：

from pyspark.sql.functions import col

# 过滤掉物品与自身的相似度记录
dfsim_filtered = dfsim.filter(col("from") != col("to"))

步骤2：用窗口函数分组取Top-K #

利用Spark窗口函数按物品ID分组，组内按相似度降序排序后取前K条记录，这是大数据场景下最高效的方式：

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

# 定义窗口规则：按物品ID(from)分组，相似度(sim)降序排序
window_spec = Window.partitionBy("from").orderBy(col("sim").desc())

K = 5  # 替换成你需要的Top-K数值
top_k_sims = dfsim_filtered.withColumn("rank", row_number().over(window_spec)) \
                          .filter(col("rank") <= K) \
                          .drop("rank")  # 移除辅助用的rank列

步骤3：（可选）整理成聚合格式 #

如果需要把每个物品的Top-K结果整理成列表形式，方便后续使用，可以用聚合函数打包：

from pyspark.sql.functions import collect_list, struct

# 将每个物品的Top-K相似物品打包成（物品ID, 相似度）的结构体列表
top_k_aggregated = top_k_sims.groupBy("from") \
                             .agg(collect_list(struct("to", "sim")).alias("top_k_similar_items"))

此时top_k_aggregated的每一行中，from是目标物品ID，top_k_similar_items是一个列表，每个元素包含相似物品的ID和对应的余弦相似度。

单个物品的Top-K查询 #

如果只需要查询某一个特定物品的Top-K相似物品，直接过滤即可：

target_item_id = 0  # 替换成你要查询的物品ID
top_k_for_target = top_k_sims.filter(col("from") == target_item_id) \
                             .select("to", "sim") \
                             .orderBy(col("sim").desc())
top_k_for_target.show()

重要提示 #

因为你的数据规模极大（309235 x 81566），绝对不要用collect()把全量数据拉到本地，所有操作尽量在Spark集群上完成，最后按需将结果写入存储（比如HDFS、数据库）或只提取单个物品的结果。

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