Calcular similitud de coseno en Apache Spark

Tengo un DataFrame con IDF de ciertas palabras calculadas. Por ejemplo

(10,[0,1,2,3,4,5],[0.413734499590671,0.4244680552337798,0.4761400657781007, 1.4004620708967006,0.37876590175292424,0.48374466516332])



 .... and so on

Ahora proporcione una consulta Q, puedo calcular el TF-IDF de esta consulta. ¿Cómo calculo la similitud de coseno de la consulta con todos los documentos en el marco de datos (hay cerca de un millón de documentos)

Podría hacerlo manualmente en un trabajo de reducción de mapas usando la multiplicación vectorial

Coseno Similitud (Q, documento) = Producto de punto (Q, documento) / || Q || * || documento ||

pero seguramente Spark ML debe admitir de forma nativa el cálculo de la similitud coseno de un texto?

En otras palabras, dada una consulta de búsqueda, ¿cómo encuentro los cosenos más cercanos del documento TF-IDF del DataFrame?

Hay un ejemplo relacionado con su problema en el repositorio de Spark aquí . La estrategia es representar los documentos como RowMatrix y luego usar su método columnSimilarities (). Eso te dará una matriz de todas las similitudes de coseno. Extraiga la fila que corresponde a su documento de consulta y ordene. Eso dará los índices de los documentos más similares.

Dependiendo de su aplicación, todo este trabajo puede realizarse antes de la consulta.