CLOVER🍀

That was when it all began.

Infinispan 15.0で使えるようになったベクトル検索をInfinispan Server × Hot Rod Clientで試す

これは、なにをしたくて書いたもの?

Infinispan 15.0.0.Finalからベクトル検索ができるようになったので、試してみたいなということで。

Infinispan 15.0のベクトル検索

Infinispan 15.0のベクトル検索の情報は、ブログエントリーに出てきます。

まずはInfinispan 15.0.0.Finalのリリースブログエントリー。

Infinispan 15.0.0.Final

こちらはクエリーのみですが、インデックスにベクトルを保存するためのマッピングの書き方などはリリース前から少し出ています。

Infinispan kNN Vector Search

ドキュメントはQueryのものを見ることになります。

Querying Infinispan caches

エンティティにベクトル検索用のフィールドを定義するには、@Vectorアノテーションを使います。@Vectorアノテーションを付与する
フィールドはbyte(Byte)またはfloat(Float)の配列である必要があります。

Querying Infinispan caches / Indexing Infinispan caches / Infinispan native indexing annotations

設定可能な属性はなぜかベクトル検索用のクエリーの方に説明がありますが、次元数(dimension)、ベクトル類似度関数(similarity)、
HNSWのパラメーターm(maxConnections)、efConstruction(beamWidth)が指定できます。

Querying Infinispan caches / Creating Ickle queries / Vector search queries / Vector field attributes

このうち次元数は必須で、その他のデフォルト値はベクトル類似度関数がL2(ユークリッド距離)、mが16、efConstructionが512です。

Infinispanでインデックスを使った検索は、Hibernate Searchを使って実現しているものであって、さらにその背後にはApache Luceneが使われています。

アノテーションの属性の意味はHibernate Searchと同じなので、こちらのドキュメントを見るのもよいでしょう。

Hibernate Search 7.1.1.Final: Reference Documentation

次はクエリーです。

Querying Infinispan caches / Creating Ickle queries / Vector search queries

ドキュメントの例を使って構文を見てみます。

from play.Item i where i.myVector <-> [7,7,7]~3

[7,7,7]で指定しているのが検索条件となるベクトルです。3というのはkで、検索するドキュメント数を指定します。

ここでkをインデックス内の全ベクトル数(≒全ドキュメント数)と同じにするとkNN検索を行います。kが全ベクトル数を下回るとANNによる
近似検索を行います。

このあたりの動作は、こちらのエントリーで調べています。

Apache LuceneでkNN検索とANN+HNSWを使い分ける(Codecを使ったHNSWのパラメーター設定付き) - CLOVER🍀

ベクトルやkはプレースホルダーとして指定可能です。

Query<Item> query = cache.query("from play.Item i where i.floatVector <-> [:a,:b,:c]~:k");
query.setParameter("a", 1);
query.setParameter("b", 4.3);
query.setParameter("c", 3.3);
query.setParameter("k", 4);



Query<Item> query = cache.query("from play.Item i where i.floatVector <-> [:a]~:k");
query.setParameter("a", new float[]{7.1f, 7.0f, 3.1f});
query.setParameter("k", 3);

また、フィルターを使った絞り込みも可能です。

Query<Object[]> query = remoteCache.query(
   "select score(i), i from Item i where i.floatVector <-> [:a]~:k filtering (i.buggy : 'cat' or i.text : 'code')");
query.setParameter("a", new float[]{7, 7, 7});
query.setParameter("k", 3);

このあたり、Apache Luceneとほぼ変わらないので前に書いたエントリーをInfinispan ServerとHot Rod Clientで試してみることにします。

Apache Luceneでベクトル検索(kNN検索)を試す - CLOVER🍀

Apache LuceneでkNN検索とANN+HNSWを使い分ける(Codecを使ったHNSWのパラメーター設定付き) - CLOVER🍀

環境

今回の環境はこちら。

$ java --version
openjdk 21.0.3 2024-04-16
OpenJDK Runtime Environment (build 21.0.3+9-Ubuntu-1ubuntu122.04.1)
OpenJDK 64-Bit Server VM (build 21.0.3+9-Ubuntu-1ubuntu122.04.1, mixed mode, sharing)


$ mvn --version
Apache Maven 3.9.7 (8b094c9513efc1b9ce2d952b3b9c8eaedaf8cbf0)
Maven home: $HOME/.sdkman/candidates/maven/current
Java version: 21.0.3, vendor: Ubuntu, runtime: /usr/lib/jvm/java-21-openjdk-amd64
Default locale: ja_JP, platform encoding: UTF-8
OS name: "linux", version: "5.15.0-112-generic", arch: "amd64", family: "unix"

Infinispan Server。

$ java --version
openjdk 21.0.3 2024-04-16 LTS
OpenJDK Runtime Environment Temurin-21.0.3+9 (build 21.0.3+9-LTS)
OpenJDK 64-Bit Server VM Temurin-21.0.3+9 (build 21.0.3+9-LTS, mixed mode, sharing)


$ bin/server.sh --version

Infinispan Server 15.0.4.Final (I'm Still Standing)
Copyright (C) Red Hat Inc. and/or its affiliates and other contributors
License Apache License, v. 2.0. http://www.apache.org/licenses/LICENSE-2.0

Infinispan Serverは172.18.0.3〜5の3ノードのクラスター構成で動作しているものとして、以下のコマンドで起動させます。

$ bin/server.sh \
    -b 0.0.0.0 \
    -Djgroups.tcp.address=$(hostname -i)

Infinispan Serverには、以下のコマンドで管理用ユーザーとアプリケーション用ユーザーを作成しているものとします。

$ bin/cli.sh user create -g admin -p password ispn-admin
$ bin/cli.sh user create -g application -p password ispn-user

それから、Pythonも使います。

$ python3 --version
Python 3.10.12


$ pip3 --version
pip 22.0.2 from /usr/lib/python3/dist-packages/pip (python 3.10)

テキスト埋め込み

Apache Luceneの時もそうでしたが、Infinispan Serverにはテキスト埋め込みを行う機能はないのでSentence Transformersを使ったAPIを作って
こちらを利用することにします。

ライブラリーのインストール。

$ pip3 install sentence-transformers fastapi uvicorn[standard]

インストールされたライブラリー。

$ pip3 list
Package                  Version
------------------------ ----------
annotated-types          0.7.0
anyio                    4.4.0
certifi                  2024.6.2
charset-normalizer       3.3.2
click                    8.1.7
dnspython                2.6.1
email_validator          2.1.2
exceptiongroup           1.2.1
fastapi                  0.111.0
fastapi-cli              0.0.4
filelock                 3.15.1
fsspec                   2024.6.0
h11                      0.14.0
httpcore                 1.0.5
httptools                0.6.1
httpx                    0.27.0
huggingface-hub          0.23.4
idna                     3.7
Jinja2                   3.1.4
joblib                   1.4.2
markdown-it-py           3.0.0
MarkupSafe               2.1.5
mdurl                    0.1.2
mpmath                   1.3.0
networkx                 3.3
numpy                    2.0.0
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        8.9.2.26
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.20.5
nvidia-nvjitlink-cu12    12.5.40
nvidia-nvtx-cu12         12.1.105
orjson                   3.10.5
packaging                24.1
pillow                   10.3.0
pip                      22.0.2
pydantic                 2.7.4
pydantic_core            2.18.4
Pygments                 2.18.0
python-dotenv            1.0.1
python-multipart         0.0.9
PyYAML                   6.0.1
regex                    2024.5.15
requests                 2.32.3
rich                     13.7.1
safetensors              0.4.3
scikit-learn             1.5.0
scipy                    1.13.1
sentence-transformers    3.0.1
setuptools               59.6.0
shellingham              1.5.4
sniffio                  1.3.1
starlette                0.37.2
sympy                    1.12.1
threadpoolctl            3.5.0
tokenizers               0.19.1
torch                    2.3.1
tqdm                     4.66.4
transformers             4.41.2
triton                   2.3.1
typer                    0.12.3
typing_extensions        4.12.2
ujson                    5.10.0
urllib3                  2.2.1
uvicorn                  0.30.1
uvloop                   0.19.0
watchfiles               0.22.0
websockets               12.0

APIとなるスクリプトはこちら。

api.py

from fastapi import FastAPI
from pydantic import BaseModel
import os
from sentence_transformers import SentenceTransformer

app = FastAPI()

class EmbeddingRequest(BaseModel):
    model: str
    text: str
    normalize: bool = False

class EmbeddingResponse(BaseModel):
    model: str
    embedding: list[float]
    dimension: int

@app.post("/embeddings/encode")
def encode(request: EmbeddingRequest) -> EmbeddingResponse:
    sentence_transformer_model = SentenceTransformer(
        request.model,
        device=os.getenv("EMBEDDING_API_DEVICE", "cpu")
    )

    embeddings = sentence_transformer_model.encode(sentences=[request.text], normalize_embeddings=request.normalize)
    embedding = embeddings[0]

    # numpy array to float list
    embedding_as_float = embedding.tolist()

    return EmbeddingResponse(
        model=request.model,
        embedding=embedding_as_float,
        dimension=sentence_transformer_model.get_sentence_embedding_dimension()
    )

リクエストにはモデル、変換するテキストを指定します。

起動。

$ uvicorn api:app

こちらのAPIを使ってテキスト埋め込みを行います。

準備

Maven依存関係など。

    <properties>
        <maven.compiler.release>21</maven.compiler.release>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-api</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-client-hotrod</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-remote-query-client</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-query-dsl</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan.protostream</groupId>
            <artifactId>protostream-processor</artifactId>
            <version>5.0.4.Final</version>
            <scope>provided</scope>
        </dependency>

        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-core</artifactId>
            <version>15.0.4.Final</version>
            <scope>test</scope>
        </dependency>

        <dependency>
            <groupId>com.fasterxml.jackson.core</groupId>
            <artifactId>jackson-databind</artifactId>
            <version>2.17.1</version>
            <scope>test</scope>
        </dependency>

        <dependency>
            <groupId>org.junit.jupiter</groupId>
            <artifactId>junit-jupiter-engine</artifactId>
            <version>5.10.2</version>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>org.assertj</groupId>
            <artifactId>assertj-core</artifactId>
            <version>3.26.0</version>
            <scope>test</scope>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <artifactId>maven-compiler-plugin</artifactId>
                <version>3.13.0</version>
                <configuration>
                    <annotationProcessorPaths>
                        <annotationProcessorPath>
                            <groupId>org.infinispan.protostream</groupId>
                            <artifactId>protostream-processor</artifactId>
                            <version>5.0.4.Final</version>
                        </annotationProcessorPath>
                    </annotationProcessorPaths>
                </configuration>
            </plugin>
        </plugins>
    </build>

Hot Rod Clientでクエリーを使うのに必要な依存関係はここまでです。

        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-api</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-client-hotrod</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-remote-query-client</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan</groupId>
            <artifactId>infinispan-query-dsl</artifactId>
            <version>15.0.4.Final</version>
        </dependency>
        <dependency>
            <groupId>org.infinispan.protostream</groupId>
            <artifactId>protostream-processor</artifactId>
            <version>5.0.4.Final</version>
            <scope>provided</scope>
        </dependency>

infinispan-coreはテスト内でキャッシュを作成するために、Jacksonは先ほど作成したAPIを呼び出す際に使います。動作確認そのものはテストコードで
行います。

エンティティの作成

最初にキャッシュに保存するエンティティを作成しましょう。

src/main/java/org/littlewings/infinispan/remote/vectorsearch/Movie.java

package org.littlewings.infinispan.remote.vectorsearch;

import org.infinispan.api.annotations.indexing.Basic;
import org.infinispan.api.annotations.indexing.Indexed;
import org.infinispan.api.annotations.indexing.Text;
import org.infinispan.api.annotations.indexing.Vector;
import org.infinispan.api.annotations.indexing.option.VectorSimilarity;
import org.infinispan.protostream.annotations.Proto;

@Proto
@Indexed
public record Movie(
        @Text
        String name,
        @Text
        String description,
        @Vector(dimension = 768, similarity = VectorSimilarity.L2)
        float[] descriptionVector,
        @Text
        String author,
        @Basic(sortable = true)
        int year
) {
}

お題は映画にします。もっとも、使うデータもそうなのですが元ネタはQdrantのチュートリアルです。

Semantic Search 101 - Qdrant

ベクトル検索用のフィールドはこちらで、次元数は768、ベクトル類似度関数はユークリッド距離(デフォルト)にしてあります。

        @Vector(dimension = 768, similarity = VectorSimilarity.L2)
        float[] descriptionVector,

次元数が768なのは、この後でテキスト埋め込みに使うモデルがintfloat/multilingual-e5-baseだからで、こちらの次元数なわけですね。

intfloat/multilingual-e5-base · Hugging Face

次に、GeneratedSchemaインターフェースを拡張し、@ProtoSchemaアノテーションを付与したインターフェースを作成します。

src/main/java/org/littlewings/infinispan/remote/vectorsearch/EntitiesInitializer.java

package org.littlewings.infinispan.remote.vectorsearch;

import org.infinispan.protostream.GeneratedSchema;
import org.infinispan.protostream.annotations.ProtoSchema;
import org.infinispan.protostream.annotations.ProtoSyntax;

@ProtoSchema(
        includeClasses = {Movie.class},
        schemaFileName = "entities.proto",
        schemaFilePath = "proto/",
        schemaPackageName = "entity",
        syntax = ProtoSyntax.PROTO3
)
public interface EntitiesInitializer extends GeneratedSchema {
}

これでエンティティの準備は完了です。

こちらをコンパイルすると

$ mvn compile

Protocol Bufferes 3のIDLが生成されます。

target/classes/proto/entities.proto

// File name: entities.proto
// Generated from : entities.proto
syntax = "proto3";
package entity;



/**
 * @Indexed
 */
message Movie {

   /**
    * @Text
    */
   string name = 1;

   /**
    * @Text
    */
   string description = 2;

   /**
    * @Vector(dimension=768, similarity=L2)
    */
   repeated float descriptionVector = 3;

   /**
    * @Text
    */
   string author = 4;

   /**
    * @Basic(sortable=true)
    */
   int32 year = 5;
}

これでエンティティおよびProtocol BuffersのIDLが用意できました。

テストコードの雛形の作成

APIを呼び出すためのクライアントはこちら。

src/test/java/org/littlewings/infinispan/remote/vectorsearch/EmbeddingClient.java

package org.littlewings.infinispan.remote.vectorsearch;

import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.util.List;

public class EmbeddingClient implements AutoCloseable {
    private String url;
    private HttpClient httpClient;
    private ObjectMapper objectMapper;

    EmbeddingClient(String host, int port, HttpClient httpClient) {
        this.url = String.format("http://%s:%d/embeddings/encode", host, port);
        this.httpClient = httpClient;
        this.objectMapper = new ObjectMapper();
    }

    public static EmbeddingClient create(String host, int port) {
        return new EmbeddingClient(
                host,
                port,
                HttpClient
                        .newBuilder()
                        .version(HttpClient.Version.HTTP_1_1)
                        .followRedirects(HttpClient.Redirect.ALWAYS)
                        .build()
        );
    }

    public EmbeddingResponse execute(EmbeddingRequest request) {
        try {
            String json = objectMapper.writeValueAsString(request);

            HttpRequest httpRequest =
                    HttpRequest.newBuilder()
                            .uri(URI.create(url))
                            .header("Content-Type", "application/json")
                            .POST(HttpRequest.BodyPublishers.ofString(json))
                            .build();

            HttpResponse<String> httpResponse =
                    httpClient
                            .send(httpRequest, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8));

            return objectMapper.readValue(httpResponse.body(), EmbeddingResponse.class);
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void close() {
        httpClient.close();
    }

    public record EmbeddingRequest(String model, String text) {
    }

    public record EmbeddingResponse(String model, List<Float> embedding, int dimension) {
    }
}

このクラスを使って、先ほどPythonで作成したテキスト埋め込み用のAPIを呼び出します。

テストコードの雛形。

src/test/java/org/littlewings/infinispan/remote/vectorsearch/RemoteVectorSearchTest.java

package org.littlewings.infinispan.remote.vectorsearch;

import org.infinispan.client.hotrod.RemoteCache;
import org.infinispan.client.hotrod.RemoteCacheManager;
import org.infinispan.client.hotrod.RemoteCacheManagerAdmin;
import org.infinispan.commons.api.query.Query;
import org.infinispan.commons.api.query.QueryResult;
import org.infinispan.configuration.cache.CacheMode;
import org.infinispan.configuration.cache.IndexStartupMode;
import org.infinispan.configuration.cache.IndexStorage;
import org.infinispan.query.remote.client.ProtobufMetadataManagerConstants;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.function.Consumer;

import static org.assertj.core.api.Assertions.assertThat;

class RemoteVectorSearchTest {
    EmbeddingClient embeddingClient;

    String createUri(String userName, String password) {
        return String.format(
                "hotrod://%s:%s@172.18.0.2:11222,172.18.0.3:11222,172.18.0.4:11222"
                        + "?context-initializers=org.littlewings.infinispan.remote.vectorsearch.EntitiesInitializerImpl",
                userName,
                password
        );
    }

    @BeforeEach
    void setUp() {
        embeddingClient = EmbeddingClient.create("localhost", 8000);

        String uri = createUri("ispn-admin", "password");

        try (RemoteCacheManager manager = new RemoteCacheManager(uri)) {
            manager.getConfiguration().getContextInitializers().forEach(serializationContextInitializer -> {
                RemoteCache<String, String> protoCache = manager.getCache(ProtobufMetadataManagerConstants.PROTOBUF_METADATA_CACHE_NAME);
                protoCache.put("entities.proto", serializationContextInitializer.getProtoFile());
            });

            RemoteCacheManagerAdmin admin = manager.administration();

            // インデックスありのDistributed Cache
            org.infinispan.configuration.cache.Configuration indexedDistCacheConfiguration =
                    new org.infinispan.configuration.cache.ConfigurationBuilder()
                            .clustering()
                            .cacheMode(CacheMode.DIST_SYNC)
                            .encoding().key().mediaType("application/x-protostream")
                            .encoding().value().mediaType("application/x-protostream")
                            // indexing
                            .indexing()
                            .enable()
                            .addIndexedEntities("entity.Movie")
                            .storage(IndexStorage.FILESYSTEM)
                            .path("${infinispan.server.data.path}/index/movieCache")
                            .startupMode(IndexStartupMode.REINDEX)
                            .reader().refreshInterval(0L)  // default 0
                            .writer().commitInterval(1000)  // default null
                            .build();

            // キャッシュがない場合は作成、すでにある場合はデータを削除
            admin.getOrCreateCache("movieCache", indexedDistCacheConfiguration)
                    .clear();
        }
    }

    @AfterEach
    void tearDown() {
        embeddingClient.close();
    }

    <K, V> void withCache(String cacheName, Consumer<RemoteCache<K, V>> func) {
        String uri = createUri("ispn-user", "password");

        try (RemoteCacheManager manager = new RemoteCacheManager(uri)) {
            RemoteCache<K, V> cache = manager.getCache(cacheName);

            func.accept(cache);
        }
    }

    // ここに、データ作成とテストを書く
}

テストの開始時にHTTPクライアントのインスタンス作成およびキャッシュの作成、テストの終了時にHTTPクライアントのクローズを行います。

キャッシュは、インデックスを有効にしたDistributed Cacheとして作成します。

            // インデックスありのDistributed Cache
            org.infinispan.configuration.cache.Configuration indexedDistCacheConfiguration =
                    new org.infinispan.configuration.cache.ConfigurationBuilder()
                            .clustering()
                            .cacheMode(CacheMode.DIST_SYNC)
                            .encoding().key().mediaType("application/x-protostream")
                            .encoding().value().mediaType("application/x-protostream")
                            // indexing
                            .indexing()
                            .enable()
                            .addIndexedEntities("entity.Movie")
                            .storage(IndexStorage.FILESYSTEM)
                            .path("${infinispan.server.data.path}/index/movieCache")
                            .startupMode(IndexStartupMode.REINDEX)
                            .reader().refreshInterval(0L)  // default 0
                            .writer().commitInterval(1000)  // default null
                            .build();

これでテスト実行時に、Infinispan Serverにこんなキャッシュが定義されます。

server/data/caches.xml

<?xml version="1.0"?>
<infinispan xmlns="urn:infinispan:config:15.0">
    <cache-container>
        <caches>
            <distributed-cache name="movieCache" mode="SYNC" remote-timeout="17500" statistics="true">
                <encoding>
                    <key media-type="application/x-protostream"/>
                    <value media-type="application/x-protostream"/>
                </encoding>
                <locking concurrency-level="1000" acquire-timeout="15000" striping="false"/>
                <indexing enabled="true" storage="filesystem" startup-mode="REINDEX" path="/opt/infinispan-server/server/data/index/movieCache">
                    <index-writer commit-interval="1000"/>
                    <indexed-entities>
                        <indexed-entity>entity.Movie</indexed-entity>
                    </indexed-entities>
                </indexing>
                <state-transfer timeout="60000"/>
            </distributed-cache>
        </caches>
    </cache-container>
</infinispan>

テキスト埋め込みを行うメソッド。

    float[] textToVector(String text) {
        EmbeddingClient.EmbeddingRequest request =
                new EmbeddingClient.EmbeddingRequest(
                        "intfloat/multilingual-e5-base",
                        text
                );

        EmbeddingClient.EmbeddingResponse response = embeddingClient.execute(request);

        float[] vector = new float[response.embedding().size()];
        for (int i = 0; i < response.embedding().size(); i++) {
            vector[i] = response.embedding().get(i);
        }

        return vector;
    }

先ほど作成したHTTPクライアントを使って、テキスト埋め込み用のAPIを呼び出します。使うモデルはintfloat/multilingual-e5-baseですね。

intfloat/multilingual-e5-base · Hugging Face

データ作成用のメソッド。

    Movie createMovie(String name, String description, String author, int year) {
        return new Movie(
                name,
                description,
                textToVector("passage: " + description),
                author,
                year
        );
    }

    List<Movie> createMovies() {
        return List.of(
                createMovie(
                        "The Time Machine",
                        "A man travels through time and witnesses the evolution of humanity.",
                        "H.G. Wells",
                        1895
                ),
                createMovie(
                        "Ender's Game",
                        "A young boy is trained to become a military leader in a war against an alien race.",
                        "Orson Scott Card",
                        1985
                ),
                createMovie(
                        "Brave New World",
                        "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.",
                        "Aldous Huxley",
                        1932
                ),
                createMovie(
                        "The Hitchhiker's Guide to the Galaxy",
                        "A comedic science fiction series following the misadventures of an unwitting human and his alien friend.",
                        "Douglas Adams",
                        1979
                ),
                createMovie(
                        "Dune",
                        "A desert planet is the site of political intrigue and power struggles.",
                        "Frank Herbert",
                        1965
                ),
                createMovie(
                        "Foundation",
                        "A mathematician develops a science to predict the future of humanity and works to save civilization from collapse.",
                        "Isaac Asimov",
                        1951
                ),
                createMovie(
                        "Snow Crash",
                        "A futuristic world where the internet has evolved into a virtual reality metaverse.",
                        "Neal Stephenson",
                        1992
                ),
                createMovie(
                        "Neuromancer",
                        "A hacker is hired to pull off a near-impossible hack and gets pulled into a web of intrigue.",
                        "William Gibson",
                        1984
                ),
                createMovie(
                        "The War of the Worlds",
                        "A Martian invasion of Earth throws humanity into chaos.",
                        "H.G. Wells",
                        1898
                ),
                createMovie(
                        "The Hunger Games",
                        "A dystopian society where teenagers are forced to fight to the death in a televised spectacle.",
                        "Suzanne Collins",
                        2008
                ),
                createMovie(
                        "The Andromeda Strain",
                        "A deadly virus from outer space threatens to wipe out humanity.",
                        "Michael Crichton",
                        1969
                ),
                createMovie(
                        "The Left Hand of Darkness",
                        "A human ambassador is sent to a planet where the inhabitants are genderless and can change gender at will.",
                        "Ursula K. Le Guin",
                        1969
                ),
                createMovie(
                        "The Three-Body Problem",
                        "Humans encounter an alien civilization that lives in a dying system.",
                        "Liu Cixin",
                        2008
                )
        );
    }

先ほども書きましたが、元ネタはQdrantのチュートリアルです。全部で13個のデータになります。

Semantic Search 101 - Qdrant

またテキスト埋め込み時に付与しているpassage:はintfloat/multilingual-e5でドキュメントに指定する接頭辞です。

    Movie createMovie(String name, String description, String author, int year) {
        return new Movie(
                name,
                description,
                textToVector("passage: " + description),
                author,
                year
        );
    }

準備はできたので、あとはテストを書いてkNN検索、ANNを試していきます。

InfinispanでkNN検索およびANNを試す

まずはkNN検索から。

    @Test
    void knnSearch() {
        List<Movie> movies = createMovies();

        this.withCache("movieCache", remoteCache -> {
            movies.forEach(movie -> remoteCache.put(movie.name(), movie));

            int k = movies.size(); // all document
            float[] vector = textToVector("query: alien invasion");

            Query<Movie> query = remoteCache.query("from entity.Movie where descriptionVector <-> [:vector]~:k");
            query.maxResults(3); // result size
            query.setParameter("vector", vector);
            query.setParameter("k", k);

            QueryResult<Movie> result = query.execute();
            assertThat(result.count().value()).isEqualTo(movies.size()); // all document
            assertThat(result.count().isExact()).isTrue();

            List<Movie> resultMovies = result.list();
            assertThat(resultMovies).hasSize(3);

            assertThat(resultMovies.get(0).name())
                    .isEqualTo("The Hitchhiker's Guide to the Galaxy");
            assertThat(resultMovies.get(0).year())
                    .isEqualTo(1979);
            assertThat(resultMovies.get(1).name())
                    .isEqualTo("The Three-Body Problem");
            assertThat(resultMovies.get(1).year())
                    .isEqualTo(2008);
            assertThat(resultMovies.get(2).name())
                    .isEqualTo("The Andromeda Strain");
            assertThat(resultMovies.get(2).year())
                    .isEqualTo(1969);
        });
    }

データをキャッシュに登録して

            movies.forEach(movie -> remoteCache.put(movie.name(), movie));

kは全ドキュメント数と同じにして、クエリーに使うベクトルを作成します。ドキュメント数と同じにしたということは、kNN検索になります。

            int k = movies.size(); // all document
            float[] vector = textToVector("query: alien invasion");

query:というのはintfloat/multilingual-e5で検索時に使う接頭辞です。

あとはクエリーのパラメーターとしてバインドします。

            Query<Movie> query = remoteCache.query("from entity.Movie where descriptionVector <-> [:vector]~:k");
            query.maxResults(3); // result size
            query.setParameter("vector", vector);
            query.setParameter("k", k);

取得件数は3件にしました。これは、Apache Luceneを使ったエントリーと合わせてあります。

検索および結果取得。

            QueryResult<Movie> result = query.execute();
            assertThat(result.count().value()).isEqualTo(movies.size()); // all document
            assertThat(result.count().isExact()).isTrue();

            List<Movie> resultMovies = result.list();
            assertThat(resultMovies).hasSize(3);

QueryResult#count#vaueは、(ページングによる)取得件数ではなくてヒットした件数を返すみたいですね。

Querying Infinispan caches / Creating Ickle queries / Number of hits

検索結果は、Apache Luceneを使った時とまったく同じでした。

続いてANN。

    @Test
    void annSearch() {
        List<Movie> movies = createMovies();

        this.withCache("movieCache", remoteCache -> {
            movies.forEach(movie -> remoteCache.put(movie.name(), movie));

            int k = movies.size() - 1; // all documents - 1
            float[] vector = textToVector("query: alien invasion");

            Query<Movie> query = remoteCache.query("from entity.Movie where descriptionVector <-> [:vector]~:k");
            query.maxResults(3); // result size
            query.setParameter("vector", vector);
            query.setParameter("k", k);

            QueryResult<Movie> result = query.execute();
            assertThat(result.count().value()).isEqualTo(movies.size()); // all document
            assertThat(result.count().isExact()).isTrue();

            List<Movie> resultMovies = result.list();
            assertThat(resultMovies).hasSize(3);

            assertThat(resultMovies.get(0).name())
                    .isEqualTo("The Hitchhiker's Guide to the Galaxy");
            assertThat(resultMovies.get(0).year())
                    .isEqualTo(1979);
            assertThat(resultMovies.get(1).name())
                    .isEqualTo("The Three-Body Problem");
            assertThat(resultMovies.get(1).year())
                    .isEqualTo(2008);
            assertThat(resultMovies.get(2).name())
                    .isEqualTo("The Andromeda Strain");
            assertThat(resultMovies.get(2).year())
                    .isEqualTo(1969);
        });
    }

違いはkの値をドキュメント数より少なくしていることですね。これでkNN検索ではなく近似のANNになります。

            int k = movies.size() - 1; // all documents - 1
            float[] vector = textToVector("query: alien invasion");

目立った違いはこれくらいです。

今回のデータとクエリーの場合、kNN検索とANNは同じ結果になります。これもApache Luceneを使った時と同じですね。

フィルターを追加する

最後にkNN検索にフィルターを追加してみます。

    @Test
    void knnSearchWithFilter() {
        List<Movie> movies = createMovies();

        this.withCache("movieCache", remoteCache -> {
            movies.forEach(movie -> remoteCache.put(movie.name(), movie));

            int k = movies.size(); // all document
            float[] vector = textToVector("query: alien invasion");

            Query<Movie> query =
                    remoteCache.query("from entity.Movie where descriptionVector <-> [:vector]~:k filtering (year: [2000 to *])");
            query.maxResults(3); // result size
            query.setParameter("vector", vector);
            query.setParameter("k", k);
            // filteringの方はパラメーター化できなかったので

            QueryResult<Movie> result = query.execute();
            assertThat(result.count().value()).isEqualTo(2);
            assertThat(result.count().isExact()).isTrue();

            List<Movie> resultMovies = result.list();
            assertThat(resultMovies).hasSize(2);

            assertThat(resultMovies.get(0).name())
                    .isEqualTo("The Three-Body Problem");
            assertThat(resultMovies.get(0).year())
                    .isEqualTo(2008);
            assertThat(resultMovies.get(1).name())
                    .isEqualTo("The Hunger Games");
            assertThat(resultMovies.get(1).year())
                    .isEqualTo(2008);
        });
    }

filteringを追加して、さらに絞り込んでいます。

            Query<Movie> query =
                    remoteCache.query("from entity.Movie where descriptionVector <-> [:vector]~:k filtering (year: [2000 to *])");
            query.maxResults(3); // result size
            query.setParameter("vector", vector);
            query.setParameter("k", k);
            // filteringの方はパラメーター化できなかったので

今回は範囲クエリーを使ったのですが、これをパラメーター化するとこういう例外になったのでやめました…。

Caused by: java.lang.ClassCastException: Cannot cast org.infinispan.objectfilter.impl.syntax.ConstantValueExpr$ParamPlaceholder to java.lang.Integer

ベクトル検索のフィルターの問題というより、範囲クエリーの問題だと思います。

結果はちゃんと絞り込まれていますね。

            QueryResult<Movie> result = query.execute();
            assertThat(result.count().value()).isEqualTo(2);
            assertThat(result.count().isExact()).isTrue();

            List<Movie> resultMovies = result.list();
            assertThat(resultMovies).hasSize(2);

            assertThat(resultMovies.get(0).name())
                    .isEqualTo("The Three-Body Problem");
            assertThat(resultMovies.get(0).year())
                    .isEqualTo(2008);
            assertThat(resultMovies.get(1).name())
                    .isEqualTo("The Hunger Games");
            assertThat(resultMovies.get(1).year())
                    .isEqualTo(2008);

ANNの方は省略します。

おわりに

Infinispan 15.0で追加された、ベクトル検索を試してみました。

もっとも仕組み自体はApache Luceneのものなので、先にApache Luceneのベクトル検索をいろいろ調べておいたおかげでかなりすんなり使えました。
これでいきなりInfinispanでベクトル検索をしていたら、まず使えていなかったと思います…。あとQdrantを試していたことも大きいですが。

今回までで、Infinispan 15.0の目新しい機能はひととおり使えたかなと思います。

今回作成したソースコードは、こちらに置いています。

https://github.com/kazuhira-r/infinispan-getting-started/tree/master/remote-vector-search