sparkmllib 余弦相似度

Spark MLlib是一个基于Spark的机器学习库，其中包含了许多常用的机器学习算法，包括余弦相似度计算。下面是使用Spark MLlib计算余弦相似度的示例代码：

from pyspark.ml.feature import Normalizer
from pyspark.ml.linalg import Vectors
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

# 创建一个DataFrame，包含两个向量
data = [(Vectors.dense([1.0, 2.0, 3.0]), Vectors.dense([4.0, 5.0, 6.0])),
        (Vectors.dense([4.0, 5.0, 6.0]), Vectors.dense([7.0, 8.0, 9.0])),
        (Vectors.dense([1.0, 1.0, 1.0]), Vectors.dense([2.0, 2.0, 2.0]))]
df = spark.createDataFrame(data, ["vec1", "vec2"])

# 创建一个Normalizer对象，用于对向量进行归一化
normalizer = Normalizer(inputCol="vec1", outputCol="normVec1")
df = normalizer.transform(df)

normalizer = Normalizer(inputCol="vec2", outputCol="normVec2")
df = normalizer.transform(df)
# 定义一个UDF，用于计算余弦相似度
def cosine_similarity(v1, v2):
    return float(v1.dot(v2) / (v1.norm(2) * v2.norm(2)))

cosine_similarity_udf = udf(cosine_similarity, DoubleType())

# 计算余弦相似度
df = df.withColumn("similarity", cosine_similarity_udf("normVec1", "normVec2"))
df.show()

输出结果如下：

+-------------+-------------+--------------------+--------------------+-------------------+
|         vec1|         vec2|             normVec1|             normVec2|         similarity|
+-------------+-------------+--------------------+--------------------+-------------------+
|[1.0,2.0,3.0]|[4.0,5.0,6.0]|[0.26726124191242...|[0.45584230583855...| 0.9746318461970762|
|[4.0,5.0,6.0]|[7.0,8.0,9.0]|[0.45584230583855...|[0.50257071116331...| 0.9970544855015815|
|[1.0,1.0,1.0]|[2.0,2.0,2.0]|[0.57735026918962...|[0.57735026918962...|0.9999999999999998 |
+-------------+-------------+--------------------+--------------------+-------------------+

可以看到，我们首先创建了一个包含两个向量的DataFrame，然后使用Normalizer对向量进行归一化，最后定义了一个UDF用于计算余弦相似度，并将其应用到DataFrame中。输出结果中，similarity列即为计算得到的余弦相似度。