如何高效计算两个稀疏Pandas DataFrame的全行余弦相似度？

问题概述 #

需要计算两个稀疏Pandas DataFrame的所有行对余弦相似度：

• DF1：192184行×256列
• DF2：7739行×256列
  当前嵌套循环+scipy.spatial.distance.cosine的实现耗时长达数天，需优化性能，最终输出CSV结果。

示例输入 #

DF1:
A    1    10    7    8    3
B    1    0    1    0    1
C    6    1    0    0    0

DF2:
D    4    0    0    0    1
E    1    0    0    0    2
F    2    8    8    5    2

预期输出 #

D         E         F
A  0.113690  0.209633  0.971074
B  0.700140  0.774597  0.546019
C  0.956943  0.441129  0.259129

当前低效代码 #

import pandas as pd
from scipy import spatial

# 测试用例
a = [1,10,7,8,3]
b = [1,0,1,0,1]
c = [6,1,0,0,0]
d = [4,0,0,0,1]
e = [1,0,0,0,2]
f = [2,8,8,5,2]

DF1 = pd.DataFrame([a,b,c])
DF1.index = ["A","B","C"]

DF2 = pd.DataFrame([d,e,f])
DF2.index = ["D","E","F"]

a_dict = {}
rows_list = list(DF2.index.values)

# 性能瓶颈：嵌套循环逐行计算
for line in DF1.iterrows():
    temp = []
    for row in DF2.iterrows():
        temp.append(1- spatial.distance.cosine(line[1],row[1]))
    a_dict[line[0]] = temp

final_output = pd.DataFrame.from_dict(a_dict, orient="index", columns = rows_list)

优化方案：利用稀疏矩阵批量运算 #

核心思路是避免Python级别的嵌套循环，改用稀疏矩阵的向量/矩阵运算，依托底层优化的数值计算库实现批量计算，性能提升可达几个数量级。

优化代码实现 #

import pandas as pd
import scipy.sparse as sp
import numpy as np

# 加载你的稀疏DataFrame（示例用测试数据）
a = [1,10,7,8,3]
b = [1,0,1,0,1]
c = [6,1,0,0,0]
d = [4,0,0,0,1]
e = [1,0,0,0,2]
f = [2,8,8,5,2]

DF1 = pd.DataFrame([a,b,c])
DF1.index = ["A","B","C"]
DF2 = pd.DataFrame([d,e,f])
DF2.index = ["D","E","F"]

# 将稀疏DataFrame转换为CSR稀疏矩阵（适合行运算）
df1_sparse = sp.csr_matrix(DF1.values)
df2_sparse = sp.csr_matrix(DF2.values)

# 1. 计算所有行对的点积
dot_product = df1_sparse.dot(df2_sparse.T)

# 2. 计算每行的L2范数（模长）
df1_norms = np.sqrt(df1_sparse.power(2).sum(axis=1))
df2_norms = np.sqrt(df2_sparse.power(2).sum(axis=1))

# 3. 计算余弦相似度矩阵（点积 / 模长乘积）
# 处理全零行避免除以零（将NaN设为0）
cos_sim_matrix = dot_product / (df1_norms @ df2_norms.T)
cos_sim_matrix = np.nan_to_num(cos_sim_matrix)

# 4. 转换为DataFrame并保存为CSV
result = pd.DataFrame(cos_sim_matrix, index=DF1.index, columns=DF2.index)
result.to_csv("cosine_similarity_results.csv")

优化原理说明 #

• 稀疏矩阵运算：CSR格式的稀疏矩阵仅存储非零元素，大幅减少内存占用，且点积运算由底层优化的C代码实现，远快于Python循环。
• 批量计算：一次性计算所有行对的点积和模长，避免逐行迭代的开销。
• 避免重复计算：spatial.distance.cosine内部会重复计算模长，这里提前批量计算所有行的模长，减少冗余运算。

额外性能建议 #

• 确保你的环境安装了优化的线性代数库（如Intel MKL或OpenBLAS），scipy会自动利用这些库加速矩阵运算。
• 如果内存允许，可将稀疏矩阵转为密集矩阵（toarray()）后计算，但对于大尺寸稀疏数据，稀疏矩阵更高效。
• 若只需要按相似度排序后的DF2行，可在得到相似度矩阵后，对每行进行排序并直接输出，无需保留完整矩阵格式，进一步节省内存。

内容的提问来源于stack exchange，提问作者J_J