不使用DataFrame.apply高效计算每行对应嵌入向量余弦距离的方法

确实，df.apply()在处理大型数据集时会因为逐行执行Python代码而变得非常慢——它本质上是在遍历每一行并运行你的lambda函数，完全没有利用numpy/scipy的向量化优化能力。下面是几种高效的替代方案，核心思路是把逐行操作转换成批量的矩阵运算：

方法一：手动向量化计算（推荐，内存和速度最优） #

余弦距离的公式是 1 - (A·B)/(||A|| * ||B||)，其中A·B是向量点积，||A||是向量的L2范数。我们可以直接用numpy的向量化运算来批量计算所有行的结果：

import numpy as np
import pandas as pd
import scipy.spatial

# 你的原始数据（加随机种子方便验证结果一致性）
np.random.seed(42)
df = pd.DataFrame(
    {
        "item_1": np.random.randint(low=0, high=10, size=1000),
        "item_2": np.random.randint(low=0, high=10, size=1000),
    }
)
embeddings = {item_id: np.random.randn(100) for item_id in range(0, 10)}

# 1. 将embeddings字典转换为numpy矩阵，行索引对应item_id
embedding_matrix = np.array([embeddings[item_id] for item_id in range(10)])

# 2. 提取所有item_1和item_2对应的嵌入向量矩阵
A = embedding_matrix[df['item_1'].values]  # 形状：(1000, 100)
B = embedding_matrix[df['item_2'].values]  # 形状：(1000, 100)

# 3. 批量计算点积、范数和余弦距离
dot_product = (A * B).sum(axis=1)
norm_A = np.linalg.norm(A, axis=1)
norm_B = np.linalg.norm(B, axis=1)
cosine_distance = 1 - (dot_product / (norm_A * norm_B))

为什么这个方法更快？ #

所有运算都是numpy的底层C实现，没有Python层面的循环开销，速度会比apply()快几十甚至上百倍，尤其是当你的数据集达到十万/百万行级别时，差距会非常明显。

你可以用下面的代码验证结果和原始apply()方法完全一致：

# 原始apply方法的结果
original_result = df.apply(lambda row: scipy.spatial.distance.cosine(embeddings[row["item_1"]], embeddings[row["item_2"]]), axis=1)

# 验证一致性
print(np.allclose(cosine_distance, original_result))  # 输出True

方法二：使用scipy的cdist（不推荐大数据集） #

如果你想直接用scipy的距离函数，可以用scipy.spatial.distance.cdist，但注意这个方法会生成一个len(A) × len(B)的距离矩阵，当数据集很大时（比如10万行），这个矩阵会占用巨大的内存（10万×10万=1e10个元素），所以只适合小数据集：

from scipy.spatial.distance import cdist

cosine_distance_matrix = cdist(A, B, metric='cosine')
cosine_distance = np.diag(cosine_distance_matrix)  # 取对角线得到对应行的距离

总结 #

处理大型DataFrame时，一定要优先考虑向量化运算，避免使用apply()这类逐行循环的方法。上面的第一种方法在速度和内存效率上都是最优的，完全可以替代你的原始实现。

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