Как получить идентификатор элемента из матрицы сходства косинусов?

Question

Я использую Spark Scala для вычисления косинусного сходства между строками Dataframe.

Схема данных ниже:

root
    |-- itemId: string (nullable = true)
    |-- features: vector (nullable = true)

Пример кадра данных ниже

    +-------+--------------------+
    | itemId|            features|
    +-------+--------------------+
    | ab    |[4.7143,0.0,5.785...|
    | cd    |[5.5,0.0,6.4286,4...|
    | ef    |[4.7143,1.4286,6....|
    ........
    +-------+--------------------+

Код для вычисления косинуса сходства:

val irm = new IndexedRowMatrix(myDataframe.rdd.zipWithIndex().map {
      case (row, index) => IndexedRow(row.getAs[Vector]("features"), index)
}).toCoordinateMatrix.transpose.toRowMatrix.columnSimilarities

В матрице irm у меня есть (i, j, оценка), где i, j - индексы элемента i и j моего исходного кадра данных. То, что я хотел бы получить (itemIdA, itemIdB, Score), где itemIdA и itemIdB являются идентификаторы индекса i и j соответственно, путем объединения этого irm с начальным кадром данных или, если есть какой-либо лучший вариант?

Answer 1

Создайте индекс строки перед преобразованием кадра данных в матрицу и создайте отображение между индексом и идентификатором. После вычисления используйте созданный Map для преобразования индекса столбца (ранее индекс строки, но измененный с помощью transpose) в идентификатор.

val rdd = myDataframe.as[(String, org.apache.spark.mllib.linalg.Vector)].rdd.zipWithIndex()
val indexMap = rdd.map{case ((id, vec), index) => (index, id)}.collectAsMap()

Рассчитайте косинус сходства, как и прежде, используя:

val irm = new IndexedRowMatrix(rdd.map{case ((id, vec), index) => IndexedRow(index, vec)})
  .toCoordinateMatrix().transpose().toRowMatrix().columnSimilarities()

Преобразовать индексы столбцов обратно в идентификаторы:

irm.entries.map(e => (indexMap(e.i), indexMap(e.j), e.value)) 

Это должно дать вам то, что вы ищете.