ElasticSearch dense_vector向量查询-Java实现

1. 简介

本文介绍如何使用ElasticSearch的Java High Level API执行向量查询，向量类型为dense_vector。

2. ElasticSearch 索引设计

PUT caster_vector1
{
  "settings": {
    "number_of_replicas": 0,
    "number_of_shards": 2
  },
  "mappings": {
    "properties": {
      "my_vector": {
        "type": "dense_vector",
        "dims": 2
      },
      "my_text": {
        "type": "text"
      }
    }
  }
}

3. 向量查询

package com.example.elasticsearchdemo;

import org.apache.http.HttpHost;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.ScriptScoreQueryBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class ESHighLevelRestQuery {
    public static void main(String[] args) throws IOException {
        // 创建一个RestHighLevelClient对象
        RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(new HttpHost("192.168.209.3", 9200, "http")));

        // 定义索引名称和类型名称
        String indexName = "caster_vector1";
        String typeName = "_doc";
        //定义最低得分
        float minScore = 1.0f;

        //通过分页限制返回的文档数量
        int page = 0;
        int size = 5;

        // 定义查询向量
        double[] queryVector = new double[]{0, -1};
        // 创建一个查询脚本
        Map<String, Object> scriptParams = new HashMap<>();
        scriptParams.put("query_vector", queryVector);

       //指定向量属性的名字，通过余弦相似度来衡量查询得分
        Script script = new Script(ScriptType.INLINE, "painless",
                "cosineSimilarity(params.query_vector, 'my_vector') + 1.0", scriptParams);


        // 创建一个ScriptScoreQueryBuilder对象
        ScriptScoreQueryBuilder queryBuilder = QueryBuilders.scriptScoreQuery(
                QueryBuilders.matchAllQuery(),
                script);

        // 创建一个SearchRequest对象并设置索引名称、类型名称和查询条件
        SearchRequest searchRequest = new SearchRequest(indexName);
        searchRequest.types(typeName);
        searchRequest.source(new SearchSourceBuilder().query(queryBuilder).from(page).size(size).minScore(minScore));

       // 执行搜索操作
        try {
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);

            // 处理搜索结果
            for (SearchHit hit : searchResponse.getHits().getHits()) {
                Map<String, Object> sourceAsMap = hit.getSourceAsMap();
                Float score = hit.getScore();
                System.out.println(sourceAsMap.toString() + ", score: " + score);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        // 关闭RestHighLevelClient对象
        client.close();

    }
}

在这个案例中，我们首先创建了一个RestHighLevelClient对象，并定义了索引名称和类型名称。
为了限制返回的文档数量， 我们设置了minScore，from，size变量。
然后，定义查询的向量，创建查询脚本，脚本中指定向量属性的名称，通过余弦相似度来衡量查询的得分。
接着，我们创建了一个ScriptScoreQueryBuilder对象和SearchRequest对象，执行搜索操作，我们将输出其错误信息，最后，关闭了RestHighLevelClient对象。