在Pyspark中将UDF余弦相似度应用于分组ML向量的问题

from pyspark.sql.functions import udf
from pyspark.ml.linalg import Vectors
from pyspark.ml.feature import VectorAssembler

data = [(1, Vectors.dense([0.5, 0.2, 0.1])),
        (1, Vectors.dense([0.9, 0.7, 0.4])),
        (2, Vectors.dense([0.2, 0.1, 0.3])),
        (2, Vectors.dense([0.1, 0.5, 0.6]))]
df = spark.createDataFrame(data, ["group", "features"])

def cosine_similarity(v1, v2):
    dot_product = float(v1.dot(v2))
    norm_v1 = float(v1.norm(2))
    norm_v2 = float(v2.norm(2))
    similarity = dot_product / (norm_v1 * norm_v2)
    return similarity

cosine_similarity_udf = udf(cosine_similarity)

assembler = VectorAssembler(inputCols=["features"], outputCol="vector_features")
df = assembler.transform(df)

df.groupBy("group").agg(cosine_similarity_udf("vector_features", "vector_features").alias("similarity"))