千万级行大矩阵各列欧氏距离计算优化方案咨询

你首先搞反了两个关键逻辑：

1. 你描述的需求是计算矩阵各列的欧式距离，但你的矩阵列数仅为250，列两两距离总共只有31125个值，完全不会触发内存溢出。你之前调用scipy接口出现溢出，是因为你直接传入原矩阵，接口会按行维度计算两两距离，才会尝试生成(10000000, 10000000)的巨型矩阵。
2. 如果你真的需要计算1000万行的全量两两欧式距离，总共有约5e13个距离值，即使用单精度浮点数存储也要200TB以上空间，物理上不可能全量载入内存，必须根据下游实际需求调整计算逻辑。

场景1：需求为计算250列之间的两两欧式距离 #

直接调用scipy内置的压缩格式距离计算接口即可，毫秒级出结果，内存占用可以忽略：

import numpy as np
from scipy.spatial.distance import pdist

Q = np.random.random([10000000, 250])
# 转置矩阵将列转为pdist要求的“样本为行”格式，pdist默认只返回上三角压缩结果，不会生成全量方阵
col_euclidean_dist = pdist(Q.T, metric="euclidean")

场景2：需求为计算1000万行的两两距离 #

你当前手写的代码慢的核心原因是用了两层Python原生循环，每次循环单独调用单次距离计算函数，完全没有利用numpy的向量化加速，Python循环的调度开销比底层C实现的向量化运算慢上万倍。根据你的实际用途选以下方案：

• 如果只需要距离的统计特征（均值、最大值、分位数、直方图等）：不要存储全量距离，分块计算过程中实时更新统计值即可，没有任何存储压力。
• 如果需要分块调用距离结果：基于欧式距离的数学展开式做分块矩阵运算，避免Python层循环，计算速度提升4~5个数量级。参考实现如下：

import numpy as np

def calc_block_dist(mat, block_size=2000):
    n_rows, n_feats = mat.shape
    # 预计算所有行的L2范数平方，避免重复计算
    row_norm_sq = np.sum(mat ** 2, axis=1)
    for i_start in range(0, n_rows, block_size):
        i_end = min(i_start + block_size, n_rows)
        current_block = mat[i_start:i_end]
        # 矩阵乘法批量计算点积
        dot_product = current_block @ mat.T
        # 展开式计算距离平方，浮点误差可能产生极小负值，用max截断到0再开方
        dist_sq = row_norm_sq[i_start:i_end, np.newaxis] + row_norm_sq[np.newaxis, :] - 2 * dot_product
        dist_block = np.sqrt(np.maximum(dist_sq, 0))
        # 此处按需处理当前块的距离：比如写入磁盘、计算统计值、提取近邻，不要把所有块存在内存里
        yield dist_block

• 如果只需要每个向量的TopK近邻距离：不要计算全量两两距离，直接用FAISS、Annoy这类专门针对大规模向量近邻搜索优化的库，千万级250维向量的TopK搜索可以在普通消费级显卡上分钟级完成。

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