ES 如何实现向量搜索

在 ES 的使用过程中,通过设置分词器可以灵活地按照文本字面实现搜索和查询。但是在某些场景下,向量搜索非常有必要,比如 CV 方面的以图搜图和 NLP 领域的语义搜索。较新的 ES 版本支持稠密向量搜索,详情如下。相关片段设置重在强调特定的关键点,需要根据自己具体使用的工具或方式进行改动或设置。

设置 mappings

此处重点就是 “dense_vector” 的 type 设置,"dims"设置为具体的向量长度。

{
        "mappings": {
            "properties": {
                "entity_name": {
                    "type": "keyword"
                },
                "entity_name_embedding":{
                  	# 此处为 mappings 中关于稠密向量设置的重点
                    "type": "dense_vector",
                    "dims": 768
                }
            }
        }
}

新建 index

按照前面设置的 mappings 创建 index。

导入数据

以 BERT 预训练模型 [CLS] 位置的输出 768 维度的 embedding 为例,上传一维数组即可(python 中),但是长度必须严格和 mappings 中保持一致。

查询体 query body

{
  "query": {
    # 通过 script score 来自定义查询的分值计算
    "script_score": {
      "query": {
        # 此处可以结合其他约束条件
        "match_all": {}
      },
      # 稠密向量查询的重点在此处,通过 script 的方式指定具体的计算方式
      # 此处是计算余弦相似度,也可以替换其他的相似度
      "script": {
		# 余弦相似度计算需要两个参数:本次待查询的向量 queryVector 和索引中的字段 entity_name_embedding
        "source": "cosineSimilarity(params.queryVector, 'entity_name_embedding')",
        "lang": "painless",
		# 此处将本次待查询的向量传入,注意“params”中 queryVector 和 “source”中的书写要保持一致。
        "params": {
          # 768 维的待查询向量
          "queryVector": [0.99, 0.85, ..., 0.94, 0.96, 0.01]   
        }
      }
    }
  }
}

其他

  • index 中的数据和 query body 应该相互配合,保证涉及计算相似度的字段都有数据。比如 index 中一共 100 个 doc,entity_name_embedding 是稠密向量字段,那么每个 doc 都应该有 768 维的稠密向量,否则查询的时候会因为没有数据或参与相似度计算的数据不全而报错。
  • ES 实现字面分词的查询底层是通过倒排索引和 TF-IDF,实现稠密向量的查询是通过降低向量维度、提前构建近邻搜索图等方式来实现快速地召回。
  • 除了 ES 以外,Milvus 等工具也支持向量搜索。