Как сделать косинусное сходство между двумя dataFrame эффективно

Question

У меня есть один DataFrame (d1) с (index, features) и второй (d2) с теми же столбцами.

features является Seq[Double], а index является String.

d1 будет составлять около миллиона строк, а d2 может составлять от 40 до 10000.

Я хотел бы получить в конце DataFrame с (index, CosineSimilarities) CosineSimilarities: для каждой d1 строк значение Seq[Double] равно косинусному подобию между этой строкой и каждой строкой d2. Так что CosineSimilarities длина должна быть равна d2 числу строк.

Моим первым подходом было использование DenseMatrix и IndexedRowMatrix и d1.multiply(d2.transpose). Но действительно трудно сопоставить результаты с index, и когда d2 станет большим, задача прекратится.

Мой второй подход был сделать:

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

Но это мучительно медленно.

Мой третий подход - Broadcast d2, и делайте построчно таким образом:

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

И это работает, это более масштабируемо и быстрее, чем подход2, но не так быстро, как подход1.

Есть другой, лучший способ?

Edit:

У меня была идея вычислить центр d2, а затем вычислить расстояние между каждым d1 до этого центра. Будет ли это работать? Есть ли способ получить центр тяжести данных?