在PySpark DataFrame中计算用户间余弦相似度（不使用sklearn）

直接用PySpark内置函数实现，完全适配大数据场景，无需依赖sklearn。具体步骤如下：

1. 转换特征为向量 #

先把用户的数值特征列合并成向量格式，方便后续计算：

from pyspark.ml.feature import VectorAssembler

# 将val1、val2、val3合并为特征向量列
assembler = VectorAssembler(inputCols=["val1", "val2", "val3"], outputCol="features")
df_features = assembler.transform(df).select("id", "features")

2. 生成不重复用户配对 #

通过自连接生成所有非重复的用户对，排除自身配对和双向重复配对：

from pyspark.sql import functions as F

# 用a.id < b.id的条件避免重复计算同一对用户
df_pairs = df_features.alias("a").join(df_features.alias("b"), on=(F.col("a.id") < F.col("b.id")), how="inner")

3. 计算余弦相似度 #

利用PySpark的向量内置函数计算点积、模长，最终得到余弦得分：

# 计算点积、模长并推导余弦相似度
df_similarity = df_pairs.select(
    F.col("a.id").alias("user1"),
    F.col("b.id").alias("user2"),
    F.expr("dot(a.features, b.features)").alias("dot_product"),
    F.expr("norm(a.features, 2)").alias("norm_a"),
    F.expr("norm(b.features, 2)").alias("norm_b")
).withColumn(
    "cosine_similarity",
    F.col("dot_product") / (F.col("norm_a") * F.col("norm_b"))
).select("user1", "user2", "cosine_similarity")

结果示例 #

执行后输出的结果格式如下（保留三位小数示例）：

+------+------+-------------------+
|user1 |user2 |cosine_similarity  |
+------+------+-------------------+
|user1 |user2 |0.998              |
|user1 |user3 |0.997              |
|user1 |user4 |0.999              |
|user2 |user3 |0.999              |
|user2 |user4 |0.999              |
|user3 |user4 |0.999              |
+------+------+-------------------+

关键优化点 #

• 用a.id < b.id的连接条件，将计算量减少一半，避免重复计算双向用户对
• 所有计算基于Spark分布式引擎执行，适配大数据量场景，不会出现单机内存溢出问题

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