如何优化ElasticSearch script_score查询 避免Python循环计算物品相似度

你可以通过Elasticsearch的批量搜索API一次性提交所有相似度查询请求，完全避免Python端循环发起HTTP请求的网络开销，核心优化思路如下：

优化方案1：使用_msearch批量执行相似度查询 #

_msearch是ES原生提供的批量查询接口，可以在单次请求中打包多个独立的查询条件，ES会并行处理所有查询后统一返回结果，和原逻辑的查询结果完全一致，但网络IO次数从N+1（N为输入ID数量）降到固定2次。

改造后的代码如下：

from typing import List, Dict
from elasticsearch import Elasticsearch

def query_es_for_similar_items(es: Elasticsearch, item_ids: List[int], index: str, n=100) -> Dict[int, List[Dict]]:
    # 第一步：批量查询所有输入ID对应的向量，和原逻辑一致
    id_query = {"query": {"ids": {"type": "_doc", "values": item_ids}}}
    id_resp = es.search(index=index, body=id_query)
    item_vectors = {}
    for hit in id_resp["hits"]["hits"]:
        item_id = int(hit["_id"])
        item_vectors[item_id] = hit["_source"]["embeddings"]
    
    # 第二步：构造msearch批量查询请求，避免循环调用search
    msearch_body = []
    for item_id, vector in item_vectors.items():
        # 每个查询前先指定查询的索引
        msearch_body.append({"index": index})
        # 每个相似度查询的具体逻辑和原逻辑完全一致
        msearch_body.append({
            "size": n,
            "query": {
                "script_score": {
                    "query": {"bool": {"must_not": [{"ids": {"values": [item_id]}}]}},
                    "script": {
                        "source": "cosineSimilarity(params.query_value, doc[params.field]) + 1",
                        "params": {
                            "field": "embeddings",
                            "query_value": vector,
                        },
                    },
                }
            }
        })
    
    # 执行批量查询
    msearch_resp = es.msearch(body=msearch_body)
    
    # 处理返回结果，格式和原返回值完全一致
    d_res = {}
    for idx, item_id in enumerate(item_vectors.keys()):
        resp = msearch_resp["responses"][idx]
        ranked_scores = [
            {
                "item_id": hit["_id"],
                "similarity": hit["_score"] / 2
            } for hit in resp["hits"]["hits"]
        ]
        d_res[item_id] = ranked_scores
    return d_res

优化方案2：升级使用原生KNN查询（可选，ES 7.10+版本支持） #

如果你使用的Elasticsearch版本≥7.10，推荐用原生的knn查询替代script_score计算余弦相似度，ES底层对向量检索做了专项优化，还可以开启HNSW向量索引大幅提升高维向量的检索速度，配合msearch使用效率提升会更明显，单条查询的结构可以调整为：

{
  "size": n,
  "knn": {
    "field": "embeddings",
    "query_vector": vector,
    "k": n,
    "num_candidates": n * 10, # 可根据召回精度需求调整
    "filter": {"bool": {"must_not": [{"ids": {"values": [item_id]}}]}}
  }
}

注意事项 #

• 单次msearch打包的查询数量建议控制在1000以内，避免请求体过大导致ES处理超时
• 如果输入ID的数量特别大，可以拆分多批msearch请求执行

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