使用dataframe Scala Spark计算余弦相似度

import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.{SparkSession, DataFrame}

object CosineSimilarityExample {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("Cosine Similarity Example")
      .master("local[*]")
      .getOrCreate()

    import spark.implicits._

    // 创建示例数据
    val data = Seq(
      (1, 1.0, 2.0, 3.0),
      (2, 4.0, 5.0, 6.0),
      (3, 7.0, 8.0, 9.0)
    ).toDF("id", "feature1", "feature2", "feature3")

    // 将数据转换为向量
    val assembler = new VectorAssembler()
      .setInputCols(Array("feature1", "feature2", "feature3"))
      .setOutputCol("features")

    val vectorizedData = assembler.transform(data).select("id", "features")

    // 计算余弦相似度
    val df1 = vectorizedData.filter($"id" === 1).select("features").rdd.map(row => row.getAs[Vector](0)).collect().head
    val df2 = vectorizedData.filter($"id" === 2).select("features").rdd.map(row => row.getAs[Vector](0)).collect().head

    val cosineSimilarity = df1.dot(df2) / (df1.norm(2) * df2.norm(2))
    println(s"Cosine Similarity between id 1 and id 2: $cosineSimilarity")

    spark.stop()
  }
}