计算大 DataFrame 成对余弦相似度的最有效方法

Question

我有一个由多列组成的 300.000 行pd.DataFrame ，其中一个是 50 维numpy形状数组(1,50) ，如下所示：

ID                     Array1    
1         [2.4252 ... 5.6363] 
2         [3.1242 ... 9.0091] 
3         [6.6775 ... 12.958]  
...
300000    [0.1260 ... 5.3323]    

然后，我生成一个具有相同形状的新numpy数组（我们称之为array2 ），并计算dataframe 的每一行与生成的数组之间的余弦相似度。 为此，我目前正在使用sklearn.metrics.pairwise.cosine_similarity并将结果保存在新列中：

from sklearn.metrics.pairwise import cosine_similarity
df['Cosine'] = cosine_similarity(df['Array1].tolist(), array2)

它按预期工作，平均需要2.5 秒才能执行。 我目前正试图将这个时间降低到 1 秒以下，只是为了减少我正在构建的系统中的等待时间。

我开始了解Vaex和Dask作为pandas的替代品，但未能将我提供的代码转换为同样更快的工作等效代码。

最好使用我提到的一种技术，我如何才能 go 更快地对大型数据集进行成对余弦计算？

Answer 1

您可以在此处使用 Faiss 并应用 knn 操作。 为此，您可以将 dataframe 放入 Faiss 索引中，然后使用 k=3000000（或数据帧的总行数）的数组对其进行搜索。

import faiss

dimension = 100

array1 = np.random.random((n, dimension)).astype('float32')


index = faiss.IndexFlatIP(d)
#add the rows of the dataframe into Faiss
for index, row in df.iterrows():  
    index.add(row)

k= len(df)
D, I = index.search(array1, k) 

请注意，您需要对向量进行归一化才能使其正常工作（因为上述解决方案是基于内积的）。