如何使用OpenSearch Java客户端执行KNN近似搜索？

核心实现代码 #

不管是高级客户端还是低级客户端，都可以直接对应控制台的KNN查询逻辑实现，以下是具体示例：

高级客户端（RestHighLevelClient）实现 #

import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.search.SearchResponse;
import org.opensearch.client.RequestOptions;
import org.opensearch.client.RestHighLevelClient;
import org.opensearch.index.query.KnnQueryBuilder;
import org.opensearch.index.query.QueryBuilders;
import org.opensearch.search.builder.SearchSourceBuilder;

import java.io.IOException;

public class OpenSearchKnnDemo {
    public static void main(String[] args) throws IOException {
        // 假设已完成RestHighLevelClient的初始化配置
        RestHighLevelClient client = new RestHighLevelClient(/* 你的客户端连接配置 */);

        // 构建KNN查询
        KnnQueryBuilder knnQuery = QueryBuilders.knnQuery(
            "embedding",       // 索引中存储向量的字段名
            new float[]{1f, 2f, 3f},  // 待搜索的向量数组
            2                  // 返回的近似邻居数量k
        );

        // 组装搜索请求
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
            .query(knnQuery)
            .size(2);  // 控制返回结果的总数量

        SearchRequest searchRequest = new SearchRequest("contentml-images-clip-batch1")
            .source(sourceBuilder);

        // 执行搜索并处理响应
        SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT);
        // 此处可根据业务需求解析response中的搜索结果
        client.close();
    }
}

低级客户端（RestClient）实现 #

如果使用低级客户端，可直接构造与控制台格式一致的JSON请求体：

import org.apache.http.HttpHost;
import org.apache.http.entity.ContentType;
import org.apache.http.nio.entity.NStringEntity;
import org.opensearch.client.Request;
import org.opensearch.client.Response;
import org.opensearch.client.RestClient;

import java.io.IOException;

public class OpenSearchLowLevelKnnDemo {
    public static void main(String[] args) throws IOException {
        RestClient client = RestClient.builder(new HttpHost("localhost", 9200, "http")).build();

        // 构造和控制台完全一致的请求体JSON
        String requestBody = "{\n" +
            "  \"size\": 2,\n" +
            "  \"query\": {\n" +
            "    \"knn\": {\n" +
            "      \"embedding\": {\n" +
            "        \"vector\": [1,2,3],\n" +
            "        \"k\": 2\n" +
            "      }\n" +
            "    }\n" +
            "  }\n" +
            "}";

        Request searchRequest = new Request("GET", "/contentml-images-clip-batch1/_search");
        searchRequest.setEntity(new NStringEntity(requestBody, ContentType.APPLICATION_JSON));

        Response response = client.performRequest(searchRequest);
        // 可通过response.getEntity().getContent()解析返回的JSON结果
        client.close();
    }
}

关键注意事项 #

• 确保Java客户端版本与OpenSearch集群版本完全一致，避免因版本差异导致的API不兼容问题
• 高级客户端的KnnQueryBuilder从OpenSearch 1.0版本开始提供，若使用旧版本客户端需检查API支持情况
• 代码中传入的向量数据类型要与索引定义的embedding字段类型匹配（比如索引字段是float类型向量，就用float[]传入）

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