如何在spark数据框架中获得给定项目向量的前k个类似项目?

Question

我得到火花数据框如下，result是id的向量:

+--------------------+--------------------+ 
| id                 |              result| 
+--------------------+--------------------+ 
|000ab862128e11eab...|[-0.46, 0.31, 0.2]  | 
|0026f306128e11eab...|[-0.46, 0.31, 0.2]  | 
|00313b10d11b11ea9...|[-0.25, 0.70, 0.36] | 
|00337629128e11eab...|[-0.46, 0.31, 0.51] | 
|005492e4128e11eab...|[0.55, 0.66, 0.85]  | 
+--------------------+--------------------+

如何有效地获得前5个最相似的项目?我定义了一个cosineSimility函数，它取两个向量在result中。作为参数。

Answer 1

您可以使用withColumn调用' cosinessimility '函数并将其结果存储为新列。然后根据这个新列对数据框进行排序，并取top(n)行。

<dataframe>.
withColumn("rank", cosineSimility(col("result"))).
sort(col("rank").desc).
top(5)