spark dataframe (scala)中tf idf输出的余弦相似度

import org.apache.spark.ml.feature.{HashingTF, IDF, Tokenizer}
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.Row

val spark = SparkSession.builder().appName("TF-IDF Example").getOrCreate()
val sentenceData = spark.createDataFrame(Seq(
  (0, "I love Spark"),
  (1, "I love Scala"),
  (2, "I love Spark and Scala")
)).toDF("id", "sentence")

val tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words")
val wordsData = tokenizer.transform(sentenceData)

val hashingTF = new HashingTF().setInputCol("words").setOutputCol("rawFeatures").setNumFeatures(20)
val featurizedData = hashingTF.transform(wordsData)

val idf = new IDF().setInputCol("rawFeatures").setOutputCol("features")
val idfModel = idf.fit(featurizedData)
val rescaledData = idfModel.transform(featurizedData)

val vecPairRdd = rescaledData.select("id", "features").rdd.map { case Row(id: Int, features: Vector) => (id, features) }
val vecPairRddCartesian = vecPairRdd.cartesian(vecPairRdd)
val cosineSimilarityRdd = vecPairRddCartesian.map { case ((id1, vec1), (id2, vec2)) =>
  val cosineSimilarity = vec1.dot(vec2) / (vec1.norm(2) * vec2.norm(2))
  (id1, id2, cosineSimilarity)
}

cosineSimilarityRdd.collect().foreach { case (id1, id2, cosineSimilarity) =>
  println(s"($id1, $id2) -> similarity: $cosineSimilarity")
}