これは、なにをしたくて書いたもの？
Infinispan 15.0.0.Finalからベクトル検索ができるようになったので、試してみたいなということで。
Infinispan 15.0のベクトル検索
Infinispan 15.0のベクトル検索の情報は、ブログエントリーに出てきます。
まずはInfinispan 15.0.0.Finalのリリースブログエントリー。
こちらはクエリーのみですが、インデックスにベクトルを保存するためのマッピングの書き方などはリリース前から少し出ています。
ドキュメントはQueryのものを見ることになります。
エンティティにベクトル検索用のフィールドを定義するには、@Vectorアノテーションを使います。@Vectorアノテーションを付与する
フィールドはbyte（Byte）またはfloat（Float）の配列である必要があります。
Querying Infinispan caches / Indexing Infinispan caches / Infinispan native indexing annotations
設定可能な属性はなぜかベクトル検索用のクエリーの方に説明がありますが、次元数（dimension）、ベクトル類似度関数（similarity）、
HNSWのパラメーターm（maxConnections）、efConstruction（beamWidth）が指定できます。
このうち次元数は必須で、その他のデフォルト値はベクトル類似度関数が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.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のチュートリアルです。
ベクトル検索用のフィールドはこちらで、次元数は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個のデータになります。
またテキスト埋め込み時に付与している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
