MongoDB聚合中实现余弦距离计算（替换欧氏距离）

我在MongoDB中存储了以numpy数组形式保存的人脸embedding，当前使用聚合管道通过欧氏算法计算数组间距离，现需修改为余弦距离计算方式，且仅返回距离≤0.68的文档。

示例文档 #

{
  "_id": {
    "$oid": "635ff70a16dfa4cd45f02c43"
  },
  "img_path": "1_11",
  "name": "11",
  "embedding": [
    0.04153144732117653,
    -0.0008036745712161064,
    ...
  ]
}

修改后的聚合管道代码 #

[
    {
        "$addFields": { 
            "target_embedding": target_embedding
        }
    },
    {"$unwind" : { "path" : "$embedding", "includeArrayIndex": "embedding_index"}},
    {"$unwind" : { "path" : "$target_embedding", "includeArrayIndex": "target_index" }},
    {
        "$project": {
            "img_path": 1,
            "embedding": 1,
            "name" : 1,
            "target_embedding": 1,
            "compare": { "$cmp": ['$embedding_index', '$target_index'] }
        }
    },
    {"$match": {"compare": 0}},
    {
      "$group": {
        "_id": "$img_path",
        "name" : { "$first": "$name" },
        # 计算两个向量的点积（余弦相似度分子）
        "dot_product": { "$sum": { "$multiply": ["$embedding", "$target_embedding"] } },
        # 计算当前embedding的模长平方
        "embedding_norm_sq": { "$sum": { "$pow": ["$embedding", 2] } },
        # 计算目标embedding的模长平方
        "target_norm_sq": { "$sum": { "$pow": ["$target_embedding", 2] } }
      }
    },
    { 
        "$project": {
            "name" : 1,
            "_id": 1,
            "dot_product": 1,
            # 计算两个向量的模长乘积（余弦相似度分母）
            "norm_product": { "$multiply": [{"$sqrt": "$embedding_norm_sq"}, {"$sqrt": "$target_norm_sq"}] }
        }
    },
    {
        "$project": {
            "name": 1,
            "_id": 1,
            # 余弦距离 = 1 - 余弦相似度，处理模长乘积为0的极端情况
            "cosine_distance": {
                "$cond": {
                    "if": { "$eq": ["$norm_product", 0] },
                    "then": 1.0,
                    "else": { "$subtract": [1, { "$divide": ["$dot_product", "$norm_product"] }] }
                }
            }
        }
    },
    # 过滤距离≤0.68的文档
    {"$match": {"cosine_distance": { "$lte": 0.68 }}},
    # 按距离升序排序
    { "$sort" : { "cosine_distance" : 1 } }
]

关键修改说明 #

1. 余弦距离计算逻辑：
   • 余弦相似度公式为 点积 / (向量A模长 * 向量B模长)，余弦距离定义为 1 - 余弦相似度，取值范围[0,2]，值越小表示向量越相似
2. 异常值处理：通过$cond判断模长乘积为0的情况，直接返回最大距离1.0，避免除以0报错
3. 过滤条件更新：直接在$match阶段筛选余弦距离≤0.68的文档，简化管道步骤
4. 结构优化：合并原欧氏计算中冗余的投影阶段，让管道逻辑更清晰

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