如何在Elasticsearch中实现多向量kNN单请求查询？

问题描述 #

使用Elasticsearch的kNN功能检索最近邻时，无法像terms查询那样在单个请求中处理多向量。例如批量ID查询可以直接传入ID列表：

query_body = {
    "terms": {
        "transaction_id": id_list  # id_list是包含多个ID的列表
    }
}

search_results = client.search(index=index, query=query_body)

但kNN的script_score查询只能单次处理一个向量：

query_body = {
    "script_score": {
        "min_score": 0.80,
        "query": {"match_all": {}},
        "script": {
            "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
            "params": {"query_vector": query_vector},
        },
    }
}

search_results = client.search(index=index, query=query_body)

如果要处理10个向量，只能逐个调用上述查询，想知道：

• 是否支持向params.query_vector传入多向量？
• 能否通过Multi Search合并多查询为一次操作？

解决方案 #

1. 不支持直接传入多向量到script_score的query_vector #

目前Elasticsearch的script_score查询中的cosineSimilarity函数仅支持单个向量与文档向量计算相似度，无法直接传入多向量批量计算，因此不能通过这种方式实现单请求多向量kNN查询。

2. 使用Multi Search（msearch）合并多查询 #

可以通过Python客户端的msearch方法，将多个单向量kNN查询合并为一次请求发送，减少网络开销。示例代码如下：

# 假设你有10个查询向量组成的列表query_vectors
queries = []
for vec in query_vectors:
    # 每个查询的元数据（指定索引）
    queries.append({"index": index})
    # 每个查询的body，和单向量kNN查询一致
    queries.append({
        "script_score": {
            "min_score": 0.80,
            "query": {"match_all": {}},
            "script": {
                "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
                "params": {"query_vector": vec},
            },
        }
    })

# 执行批量查询
msearch_results = client.msearch(body=queries)

# 处理结果：每个查询的结果在msearch_results['responses']列表中，顺序和传入的query_vectors一致
for idx, result in enumerate(msearch_results['responses']):
    print(f"第{idx+1}个向量的查询结果：", result['hits']['hits'])

3. 原生批量kNN查询（Elasticsearch ≥8.7） #

如果你的Elasticsearch版本在8.7及以上，可以使用官方提供的批量kNN查询功能，相比msearch更高效。示例代码如下：

query_body = {
    "knn": {
        "field": "vector",
        "k": 10,  # 每个向量返回的最近邻数量
        "num_candidates": 100,  # 候选集大小，影响精度和性能
        "queries": [
            {"vector": query_vector_1, "id": "query_1"},
            {"vector": query_vector_2, "id": "query_2"},
            # 可添加更多查询向量
        ]
    }
}

search_results = client.search(index=index, knn=query_body['knn'])

# 处理结果：每个匹配文档会通过`matched_queries`字段关联对应的查询ID
for hit in search_results['hits']['hits']:
    print(f"文档ID：{hit['_id']}, 匹配的查询ID：{hit['matched_queries']}")

注意：该功能要求索引的向量字段为dense_vector类型，且配置了近似kNN索引（如hnsw）。

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