如何高效计算两个DataFrame间的余弦相似度？

我有两个结构相同的DataFrame：d1和d2，都包含index（String类型）和features（Seq[Double]类型）列。其中d1大概有100万行，d2的行数在40到10000之间。我最终想要得到一个包含index和CosineSimilarities（Seq[Double]类型，对应d1每一行与d2所有行的余弦相似度，长度等于d2的行数）的DataFrame。

我已经尝试过三种方法，但都有各自的问题：

• 方法1：用DenseMatrix、IndexedRowMatrix调用d1.multiply(d2.transpose)，但很难把结果映射回原index，而且当d2行数较多时任务会失败；
• 方法2：计算笛卡尔积后逐对计算：

  d1.cartesian(d2)
    .repartition(n)
    .map { case ((d1index, d1features), (_, d2features)) => 
      (d1index, myCosineSimilarityMethod(d1features, d2features)) 
    }
  

  这种方法速度极慢，完全不适合大数据量；
• 方法3：广播d2后逐行计算：

  d1.mapValues { d1features => 
    d2broadcasted.value
      .map { case (_, d2features) => myCosineSimilarityMethod(d1features, d2features) }
      .toSeq
  }
  

  这个方法可行，比方法2扩展性更好、速度更快，但还是不如方法1的计算效率高。

想请教有没有更优的解决方法？

补充问题：我想到可以计算d2的质心，再计算每个d1行到该质心的距离，这种思路可行吗？另外怎么获取DataFrame的质心？

嘿，针对你的场景，其实可以结合矩阵乘法的高效性和广播的灵活性，同时解决方法1的索引映射问题，这里有几个更优的思路：

1. 带索引追踪的优化版矩阵乘法方案 #

既然方法1的矩阵乘法效率最高，那我们重点解决它的索引映射问题就好：

• 先给d1添加一个自增整数列（比如row_id），用来关联矩阵行和原index；
• 将d1和d2的features转换成Spark ML的Vector类型，分别构建RowMatrix；
• 对d2的RowMatrix做转置后，和d1的矩阵执行乘法得到相似度矩阵；
• 把矩阵结果转换成DataFrame，通过row_id关联回原index，最后把每行的相似度值整理成Seq[Double]格式。

这种方式既保留了矩阵乘法的高效性，又解决了索引丢失的问题，针对d2行数较多的情况，可以通过调整RowMatrix的blockSize参数、增加分区数来缓解内存压力，整体效率会比方法3高不少。

2. 利用Spark ML原生CosineSimilarity组件 #

Spark ML里有现成的优化过的余弦相似度计算工具，比你自己实现的myCosineSimilarityMethod更高效：

• 先把d1和d2的features转换成DenseVector或SparseVector（如果特征稀疏的话，稀疏向量能大幅减少计算量）；
• 把d2的所有向量合并成一个Matrix并广播，然后对d1的每个向量，批量计算与d2所有向量的余弦相似度；
• 底层是优化后的矩阵运算逻辑，能充分利用Spark的分布式计算能力，比手动实现的逐行计算快很多。

关于质心思路的解答 #

这个思路是否可行，完全取决于你的业务需求：

• 如果你的目标是快速得到d1每行与d2整体的相似程度，而不是和d2每一行的相似度，那计算质心是完全可行的，而且计算量会小很多；
• 如果还是需要和d2每一行的相似度，那质心方法就无法替代原需求，只能作为额外的辅助特征。

如何计算d2的质心 #

质心就是所有特征向量的逐元素均值，你可以这样实现：

import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.functions._

// 先把Seq[Double]转换成Spark ML的Vector类型
val d2WithVector = d2.withColumn("vec", array_to_vector(col("features")))

// 对所有向量的每个维度取均值，得到质心
val centroid = d2WithVector.select(mean("vec")).first().getAs[Vector](0)

得到centroid后，就可以用它和d1的每行特征计算余弦相似度了。

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