PolarDB:OpenSearch プロトコルを使用したベクトル検索の実行

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.example</groupId>
    <artifactId>vector-search-demo</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>11</maven.compiler.source>
        <maven.compiler.target>11</maven.compiler.target>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.opensearch.client</groupId>
            <artifactId>opensearch-java</artifactId>
            <version>3.0.0</version>
        </dependency>
        <dependency>
            <groupId>org.apache.httpcomponents.client5</groupId>
            <artifactId>httpclient5</artifactId>
            <version>5.3</version>
        </dependency>
        <dependency>
            <groupId>org.apache.httpcomponents.core5</groupId>
            <artifactId>httpcore5</artifactId>
            <version>5.3</version>
        </dependency>
        <dependency>
            <groupId>org.apache.httpcomponents.core5</groupId>
            <artifactId>httpcore5-h2</artifactId>
            <version>5.3</version>
        </dependency>
        <!-- Jackson databind は opensearch-java に必要です -->
        <dependency>
            <groupId>com.fasterxml.jackson.core</groupId>
            <artifactId>jackson-databind</artifactId>
            <version>2.16.1</version>
        </dependency>
    </dependencies>
</project>

package com.example;

import org.apache.hc.client5.http.auth.AuthScope;
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManagerBuilder;
import org.apache.hc.core5.http.HttpHost;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch._types.mapping.KeywordProperty;
import org.opensearch.client.opensearch._types.mapping.KnnVectorMethod;
import org.opensearch.client.opensearch._types.mapping.KnnVectorProperty;
import org.opensearch.client.opensearch._types.mapping.Property;
import org.opensearch.client.opensearch._types.mapping.TextProperty;
import org.opensearch.client.opensearch._types.mapping.TypeMapping;
import org.opensearch.client.opensearch._types.query_dsl.KnnQuery;
import org.opensearch.client.opensearch._types.query_dsl.MatchQuery;
import org.opensearch.client.opensearch._types.query_dsl.Query;
import org.opensearch.client.opensearch._types.query_dsl.TermQuery;
import org.opensearch.client.opensearch.core.IndexRequest;
import org.opensearch.client.opensearch.core.SearchRequest;
import org.opensearch.client.opensearch.core.SearchResponse;
import org.opensearch.client.opensearch.core.search.Hit;
import org.opensearch.client.opensearch.indices.CreateIndexRequest;
import org.opensearch.client.opensearch.indices.DeleteIndexRequest;
import org.opensearch.client.transport.OpenSearchTransport;
import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBuilder;
import org.opensearch.client.json.JsonData;
import org.opensearch.client.json.jackson.JacksonJsonpMapper;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class VectorSearchDemo {

    private static final String INDEX_NAME = "test-java-demo-full-search";
    private static final String FIELD_NAME = "test-embedding";
    private static final int VECTOR_DIMENSION = 4;

    public static void main(String[] args) throws IOException {
        OpenSearchClient client = createClient("<polarsearch_host>", <polarsearch_port>, "<polarsearch_username>", "<polarsearch_password>");
        System.out.println("クライアントが初期化されました。");

        deleteIndexIfExists(client);
        createVectorIndex(client);
        System.out.println("インデックス '" + INDEX_NAME + "' が作成されました。");

        indexSampleData(client);
        System.out.println("サンプルデータがインデックス化されました。");
        
        try {
            Thread.sleep(2000);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        List<Float> queryVector = List.of(3.1f, 4.1f, 5.1f, 6.1f);

        // 3 種類の検索を順番に実行します。
        performBasicKnnSearch(client, queryVector);
        performHybridSearchWithText(client, queryVector);
        performFilteredSearchWithTerm(client, queryVector);
        
        client._transport().close();
        System.out.println("\nクライアントがクローズされました。");
    }

    // クライアントを初期化します。
    private static OpenSearchClient createClient(String hostName, int port, String username, String password) {
        final var host = new HttpHost("http", hostName, port);
        final var credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(new AuthScope(host), new UsernamePasswordCredentials(username, password.toCharArray()));
        final var connectionManager = PoolingAsyncClientConnectionManagerBuilder.create().build();
        OpenSearchTransport transport = ApacheHttpClient5TransportBuilder.builder(host)
            .setMapper(new JacksonJsonpMapper())
            .setHttpClientConfigCallback(httpClientBuilder ->
                httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider).setConnectionManager(connectionManager)
            ).build();
        return new OpenSearchClient(transport);
    }

    // クリーンなテストを保証するために、インデックスが存在する場合は削除します。
    private static void deleteIndexIfExists(OpenSearchClient client) throws IOException {
        if (client.indices().exists(r -> r.index(INDEX_NAME)).value()) {
            client.indices().delete(new DeleteIndexRequest.Builder().index(INDEX_NAME).build());
            System.out.println("インデックス '" + INDEX_NAME + "' が削除されました。");
        }
    }

    // ベクトルインデックスを作成します。
    private static void createVectorIndex(OpenSearchClient client) throws IOException {
        Property vectorProperty = Property.of(p -> p.knnVector(
            KnnVectorProperty.of(kvp -> kvp
                .dimension(VECTOR_DIMENSION)
                .dataType("float")
                .method(new KnnVectorMethod.Builder()
                    .name("hnsw")
                    .engine("faiss")
                    .spaceType("l2")
                    .parameters(Map.of(
                        "m", JsonData.of(16),
                        "ef_construction", JsonData.of(256)
                    ))
                    .build()
                )
            )
        ));
        
        TypeMapping mapping = TypeMapping.of(m -> m
            .properties(FIELD_NAME, vectorProperty)
            .properties("text", Property.of(p -> p.text(TextProperty.of(t -> t))))
            .properties("category", Property.of(p -> p.keyword(k -> k))) // カテゴリフィールドを追加します。
        );
        
        CreateIndexRequest request = new CreateIndexRequest.Builder()
            .index(INDEX_NAME)
            .settings(s -> s.knn(true))
            .mappings(mapping)
            .build();
            
        client.indices().create(request);
    }

    // ベクトルデータをインデックス化します。
    private static void indexSampleData(OpenSearchClient client) throws IOException {
        List<Map<String, Object>> documents = new ArrayList<>();
        documents.add(Map.of("text", "a book about data science", "category", "books", FIELD_NAME, List.of(1.0f, 2.0f, 3.0f, 4.0f)));
        documents.add(Map.of("text", "an intelligent smartphone with a great camera", "category", "electronics", FIELD_NAME, List.of(8.0f, 7.0f, 6.0f, 5.0f)));
        documents.add(Map.of("text", "a technical manual for a smart device", "category", "electronics", FIELD_NAME, List.of(3.0f, 4.0f, 5.0f, 6.0f)));

        for (int i = 0; i < documents.size(); i++) {
            IndexRequest<Map<String, Object>> request = new IndexRequest.Builder<Map<String, Object>>()
                .index(INDEX_NAME)
                .id("doc_" + i)
                .document(documents.get(i))
                .build();
            client.index(request);
        }
    }

    /**
     * 例 1: 基本的な k-NN 検索
     */
    private static void performBasicKnnSearch(OpenSearchClient client, List<Float> queryVector) throws IOException {
        System.out.println("\n--- 1. 基本的な k-NN 検索を実行 ---");
        System.out.println("最も類似したベクトルをクエリしています: " + queryVector);
        int k = 3;

        KnnQuery knnQuery = new KnnQuery.Builder()
            .field(FIELD_NAME)
            .vector(queryVector)
            .k(k)
            .build();

        SearchRequest searchRequest = new SearchRequest.Builder().index(INDEX_NAME).query(new Query.Builder().knn(knnQuery).build()).size(k).build();
        SearchResponse<Map> response = client.search(searchRequest, Map.class);

        printResults(response);
    }

    /**
     * 例 2: ハイブリッド検索 (k-NN + テキスト一致)
     */
    private static void performHybridSearchWithText(OpenSearchClient client, List<Float> queryVector) throws IOException {
        System.out.println("\n--- 2. ハイブリッド検索 (k-NN + テキスト一致) を実行 ---");
        String textQuery = "book";
        System.out.println("'" + textQuery + "' を含むドキュメントをフィルタリングし、最も類似したベクトルを検索します。");
        int k = 3;

        MatchQuery matchQuery = new MatchQuery.Builder().field("text").query(q -> q.stringValue(textQuery)).build();

        KnnQuery hybridKnnQuery = new KnnQuery.Builder()
            .field(FIELD_NAME)
            .vector(queryVector)
            .k(k)
            .filter(new Query.Builder().match(matchQuery).build())
            .build();
        
        SearchRequest searchRequest = new SearchRequest.Builder().index(INDEX_NAME).query(new Query.Builder().knn(hybridKnnQuery).build()).size(k).build();
        SearchResponse<Map> response = client.search(searchRequest, Map.class);
        
        printResults(response);
    }

    /**
     * 例 3: フィルター検索 (k-NN + term フィルター)
     */
    private static void performFilteredSearchWithTerm(OpenSearchClient client, List<Float> queryVector) throws IOException {
        System.out.println("\n--- 3. フィルター検索 (k-NN + Term フィルター) を実行 ---");
        String categoryFilter = "electronics";
        System.out.println("カテゴリ '" + categoryFilter + "' のドキュメントをフィルタリングし、最も類似したベクトルを検索します。");
        int k = 3;

        TermQuery termQuery = new TermQuery.Builder().field("category").value(v -> v.stringValue(categoryFilter)).build();
        
        KnnQuery filteredKnnQuery = new KnnQuery.Builder()
            .field(FIELD_NAME)
            .vector(queryVector)
            .k(k)
            .filter(new Query.Builder().term(termQuery).build())
            .build();
        
        SearchRequest searchRequest = new SearchRequest.Builder().index(INDEX_NAME).query(new Query.Builder().knn(filteredKnnQuery).build()).size(k).build();
        SearchResponse<Map> response = client.search(searchRequest, Map.class);

        printResults(response);
    }

    private static void printResults(SearchResponse<Map> response) {
        System.out.println("検索結果:");
        for (Hit<Map> hit : response.hits().hits()) {
            System.out.printf(" - ID: %s, スコア: %.4f, ソース: %s%n", hit.id(), hit.score(), hit.source());
        }
        if (response.hits().hits().isEmpty()) {
            System.out.println(" - 結果が見つかりませんでした。");
        }
    }
}