CLOVER🍀

That was when it all began.

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

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

前にApache Luceneでベクトル検索(kNN検索)を試してみました。

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

この後でANNに関係するHNSWのパッケージやクラスがあるのを見つけたので、今回はこちらを扱ってみたいなと思いまして。

kNN検索、ANNとHNSW

前回使ったのはkNN検索でした。これはQdrantの時にも扱いました。

Qdrantのチュートリアルから、「検索品質を測定する(Measure retrieval quality)」を試す - CLOVER🍀

あらためて内容を書いておきます。

k近傍法(ケイきんぼうほう、英: k-nearest neighbor algorithm, k-NN)は、入力との類似度が高い上位 k個の学習データで多数決/平均するアルゴリズムである。

k近傍法は以下の手順からなる:
1. 入力と全学習データとの類似度(距離)測定
2. 類似度上位 k 個の選出
3. 選出されたデータの多数決あるいは平均
すなわち「入力とよく似た k 個のデータで多数決/平均する」単純なアルゴリズムである

k近傍法 - Wikipedia

一方でANNは最近傍探索の派生で、近似最近傍探索と呼ばれるものです。

近似最近傍探索 (英: approximate nearest neighbor, ANN)

最近傍探索 - Wikipedia

kNN検索に比べると、近似で検索することになります。

ANNを使うことによる利点は、検索対象のベクトル空間が大きくなった時に計算量が指数関数的に増加する(次元の呪い)のですが
この対策になっています。

次元の呪い - Wikipedia

つまりkNN検索に比べると、ANNは速度を犠牲にして高速化を狙っていることになります。

ではどうやって近似するのかというと、それにはアルゴリズムがあり、そのひとつがHNSW(Hierarchical navigable small world)です。

Hierarchical navigable small world - Wikipedia

これは多くのベクトルデータベースで採用されているグラフベースの近似最近傍検索手法です。

HNSWは階層グラフになっていて、次の2つのパラメーターを取ります。

  • m(ノードごとのエッジの数=各ノードが接続する隣接ノードの数) … この値を大きくするほど検索の精度は向上するが、より多くのメモリやディスクを消費する
  • efConstruction(インデックス構築中に考慮する近傍の数) … この値を大きくするほど検索の精度は向上するが、インデックス作成速度は遅くなる

HNSWについては、Pineconeのドキュメントも参考になるでしょう。

Hierarchical Navigable Small Worlds (HNSW) | Pinecone

Apache LuceneのHNSW関連のパッケージ・クラス

Apache LuceneのHNSW関連のパッケージはこちらです。

org.apache.lucene.util.hnsw (Lucene 9.10.0 core API)

検索時には、このあたりを使うみたいですね。

ただ、これらのクラスを直接使うことはないと思います。

その理由は以下のクラスなのですが、こちらは抽象クラスです。

KnnVectorsFormat (Lucene 9.10.0 core API)

具象クラスとしてはこれらがあります。

このうちPerFieldKnnVectorsFormatは、フィールドごとに設定するためのクラスのようです。

これらのクラスでmやefConstructionを設定するのですが、その方法を知るにはApache LuceneのCodecについて知っておく必要があります。

Apache LuceneのCodec

Apache LuceneのCodecクラスは、転置インデックスセグメントをエンコード/デコードするためのクラスです。

Codec (Lucene 9.10.0 core API)

このクラスから、各種データ形式へのFormatを取得できます。

各Formatクラスからはインデックスファイル(セグメント)を読み書きするためのクラスが提供されるので、これを介してインデックス
ファイル(セグメント)にアクセスします。なのでEncoder/Decoderなわけですね。

Codecについては、elastic社のブログが参考になるでしょう。

What is an Apache Lucene Codec? | Elastic Blog

なお、Codecクラスは抽象クラスであり、具象クラスとしては現時点のApache LuceneではLucene99Codecクラスが使われます。

Lucene99Codec (Lucene 9.10.0 core API)

Apache LuceneのCodecとベクトル検索

Codecクラスから各種データ形式へのFormatが取得できると書きましたが、ベクトルで使われるKnnVectorsFormatもその一種です。

KnnVectorsFormat (Lucene 9.10.0 core API)

KnnVectorsFormatからは、次の2つのクラスのインスタンスが得られます。

KnnVectorsReaderにはsearchメソッドが定義されていて、この中で呼び出し方に応じてkNN検索またはANNが実行されることに
なります。

KnnVectorsReaderクラスは抽象クラスで、現時点の実装クラスはLucene99HnswVectorsReaderですね。

Lucene99HnswVectorsReader (Lucene 9.10.0 core API)

具体的には、KnnByteVectorQueryまたはKnnFloatVectorQueryで指定するkの値(検索するドキュメント数)がインデックス内の
全ベクトル数(≒全ドキュメントの数)より小さければHNSWを使った(HnswGraphSearcherによる)ANNを、そうでなければkNN検索を
実行するようです。

    if (knnCollector.k() < scorer.maxOrd()) {
      HnswGraphSearcher.search(scorer, collector, getGraph(fieldEntry), acceptedOrds);
    } else {
      // if k is larger than the number of vectors, we can just iterate over all vectors
      // and collect them
      for (int i = 0; i < scorer.maxOrd(); i++) {
        if (acceptedOrds == null || acceptedOrds.get(i)) {
          knnCollector.incVisitedCount(1);
          knnCollector.collect(scorer.ordToDoc(i), scorer.score(i));
        }
      }
    }

https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java#L237-L248

HnswGraphSearcherなどを直接使うことがないだろうと書いたのは、このあたりが理由です。

これらのクエリーの説明には、以下のようなことが書かれています。

  • If the filter cost is less than k, just execute an exact search
  • Otherwise run a kNN search subject to the filter
  • If the kNN search visits too many vectors without completing, stop and run an exact search

この中のフィルターのコストによっては動作が変わるという記述が、kNN検索が行われるのかANNが行われるのかの切り替えについて
触れている内容なのかなと思います。

なお、以下のクラスとこの文脈での関係はわかっていません…。

ところで、HNSWにはmとefConstructionという2つのパラメーターがあるという話を書きました。これらを設定するのは以下のクラスに
なります。

そしてデフォルト値はmが16、

https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsFormat.java#L104-L105

efConstruction(Apache Lucene上ではbeamWidthという名前になっています)が100になっています。

https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsFormat.java#L114-L117

ちなみにmの最大値は512、efConstructionの最大値は3200です。

Hibernate Searchのドキュメントでは、mは2〜100の範囲で指定することが推奨されていました。efConstructionについては特に推奨値は
ありませんでしたね。

Hibernate Search 7.1.1.Final: Reference Documentation

特にCodecの設定を行わない場合は、Lucene99Codecが作成したKnnVectorsFormatのインスタンスが使われます。

これはPerFieldKnnVectorsFormatとLucene99HnswVectorsFormatの組み合わせになっています。
※Lucene99Codec#getKnnVectorsFormatForFieldはdefaultKnnVectorsFormatを呼び出しているだけです。この時fieldの値は無視します

  private final KnnVectorsFormat defaultKnnVectorsFormat;
  private final KnnVectorsFormat knnVectorsFormat =
      new PerFieldKnnVectorsFormat() {
        @Override
        public KnnVectorsFormat getKnnVectorsFormatForField(String field) {
          return Lucene99Codec.this.getKnnVectorsFormatForField(field);
        }
      };

https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99Codec.java#L76-L83

    this.defaultKnnVectorsFormat = new Lucene99HnswVectorsFormat();

https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99Codec.java#L103

Lucene99HnswVectorsFormatをデフォルトコンストラクタで作成した場合、mもefConstructionもデフォルト値になるため、
これをカスタマイズする場合はApache LuceneのCodecの設定を行うことになります。

CodecはIndexWriterConfigで設定することになります。

IndexWriterConfig (Lucene 9.10.0 core API)

そして、設定したm(maxConn)やefConstruction(beamWidth)は書き込み時に使われます。

  @Override
  public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
    return new Lucene99HnswVectorsWriter(
        state,
        maxConn,
        beamWidth,
        flatVectorsFormat.fieldsWriter(state),
        numMergeWorkers,
        mergeExec);
  }

https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsFormat.java#L203-L212

このあたりが今回Apache LuceneのCodecの話を持ち出した事情ですね。

では、今回はCodecのカスタマイズをしたり、kNN検索、ANNをしたりしてみましょう。

準備

今回の環境はこちら。

$ 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-107-generic", arch: "amd64", family: "unix"

また、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にはないので、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.1
exceptiongroup           1.2.1
fastapi                  0.111.0
fastapi-cli              0.0.4
filelock                 3.14.0
fsspec                   2024.6.0
h11                      0.14.0
httpcore                 1.0.5
httptools                0.6.1
httpx                    0.27.0
huggingface-hub          0.23.2
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                    1.26.4
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.3
packaging                24.0
pillow                   10.3.0
pip                      22.0.2
pydantic                 2.7.3
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.0
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.0
tqdm                     4.66.4
transformers             4.41.2
triton                   2.3.0
typer                    0.12.3
typing_extensions        4.12.1
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.apache.lucene</groupId>
            <artifactId>lucene-core</artifactId>
            <version>9.10.0</version>
        </dependency>

        <dependency>
            <groupId>org.junit.jupiter</groupId>
            <artifactId>junit-jupiter</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>

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

        <dependency>
            <groupId>org.jboss.byteman</groupId>
            <artifactId>byteman-bmunit5</artifactId>
            <scope>test</scope>
            <version>4.0.23</version>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <artifactId>maven-surefire-plugin</artifactId>
                <version>3.2.5</version>
                <configuration>
                    <argLine>-Djdk.attach.allowAttachSelf -XX:+EnableDynamicAgentLoading</argLine>
                </configuration>
            </plugin>
        </plugins>
    </build>

Apache Luceneのベクトル検索を行うにはlucene-coreがあれば十分で、HNSWに関するパッケージもlucene-coreに含まれています。

確認はテストコードで行い、Jacksonは先ほど作成したAPIを呼び出す際に使います。

なぜかBytemanが入っていますが、これは検索に使ったのがkNN検索なのかANNなのか、
それからm(maxConn)やefConstruction(beamWidth)をカスタマイズした時に反映できているのか、といった確認に使います。

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

src/test/java/org/littlewings/lucene/ann/EmbeddingClient.java

package org.littlewings.lucene.ann;

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) {
    }
}

テストコードの雛形。

src/test/java/org/littlewings/lucene/ann/AnnSearchTest.java

package org.littlewings.lucene.ann;

import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.codecs.lucene99.Lucene99Codec;
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnFloatVectorQuery;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.ByteBuffersDirectory;
import org.apache.lucene.store.Directory;
import org.jboss.byteman.contrib.bmunit.BMScript;
import org.jboss.byteman.contrib.bmunit.BMScripts;
import org.jboss.byteman.contrib.bmunit.BMUnitConfig;
import org.jboss.byteman.contrib.bmunit.WithByteman;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Arrays;
import java.util.List;

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

@WithByteman
@BMUnitConfig(loadDirectory = "target/test-classes")
@BMScripts(scripts = {
        @BMScript("trace-ann-searcher.btm"),
        @BMScript("trace-hnsw-parameters.btm"),
}
)
class AnnSearchTest {
    EmbeddingClient embeddingClient;

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

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

    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;
    }

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

テキスト埋め込みに使うモデルは、intfloat/multilingual-e5-baseとします。

intfloat/multilingual-e5-base · Hugging Face

データの作成

データの作成は、以下のメソッドを使って行います。

    List<Document> createDocuments() {
        return List.of(
                createDocument(
                        "The Time Machine",
                        "A man travels through time and witnesses the evolution of humanity.",
                        "H.G. Wells",
                        1895
                ),
                createDocument(
                        "Ender's Game",
                        "A young boy is trained to become a military leader in a war against an alien race.",
                        "Orson Scott Card",
                        1985
                ),
                createDocument(
                        "Brave New World",
                        "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.",
                        "Aldous Huxley",
                        1932
                ),
                createDocument(
                        "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
                ),
                createDocument(
                        "Dune",
                        "A desert planet is the site of political intrigue and power struggles.",
                        "Frank Herbert",
                        1965
                ),
                createDocument(
                        "Foundation",
                        "A mathematician develops a science to predict the future of humanity and works to save civilization from collapse.",
                        "Isaac Asimov",
                        1951
                ),
                createDocument(
                        "Snow Crash",
                        "A futuristic world where the internet has evolved into a virtual reality metaverse.",
                        "Neal Stephenson",
                        1992
                ),
                createDocument(
                        "Neuromancer",
                        "A hacker is hired to pull off a near-impossible hack and gets pulled into a web of intrigue.",
                        "William Gibson",
                        1984
                ),
                createDocument(
                        "The War of the Worlds",
                        "A Martian invasion of Earth throws humanity into chaos.",
                        "H.G. Wells",
                        1898
                ),
                createDocument(
                        "The Hunger Games",
                        "A dystopian society where teenagers are forced to fight to the death in a televised spectacle.",
                        "Suzanne Collins",
                        2008
                ),
                createDocument(
                        "The Andromeda Strain",
                        "A deadly virus from outer space threatens to wipe out humanity.",
                        "Michael Crichton",
                        1969
                ),
                createDocument(
                        "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
                ),
                createDocument(
                        "The Three-Body Problem",
                        "Humans encounter an alien civilization that lives in a dying system.",
                        "Liu Cixin",
                        2008
                )
        );
    }

    Document createDocument(String name, String description, String author, int year) {
        Document document = new Document();
        document.add(new TextField("name", name, Field.Store.YES));
        document.add(new TextField("description", description, Field.Store.YES));
        document.add(new TextField("author", author, Field.Store.YES));
        document.add(new IntField("year", year, Field.Store.YES));

        float[] vector = textToVector("passage: " + description);
        document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));

        return document;
    }

    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;
    }

これはこちらのエントリーで使ったものと同じ、13個のデータで

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

Qdrantのチュートリアルのものでもあります。

Semantic Search 101 - Qdrant

テキスト埋め込みを行い、ベクトルを保存するフィールドはこちらです。ベクトル類似度関数はユークリッド距離とします。

        float[] vector = textToVector("passage: " + description);
        document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));

passage:はintfloat/multilingual-e5でドキュメントに指定する接頭辞です。

トレースの考え方

検索時にkNN検索とANNのどちらになっているのかを確認しようと、Explanationを使ってみたのですが。

Explanation (Lucene 9.10.0 core API)

得られた結果がこんな感じで、ほとんど情報がありませんでした。

0.75919974 = within top 13

これでは困るのでどう確認するか迷ったところ、Bytemanを使うことにしました。

BMUnit : Using Byteman with JUnit or TestNG fro...| JBoss.org Content Archive (Read Only)

スクリプトの指定はこちら。

@WithByteman
@BMUnitConfig(loadDirectory = "target/test-classes")
@BMScripts(scripts = {
        @BMScript("trace-ann-searcher.btm"),
        @BMScript("trace-hnsw-parameters.btm"),
}
)
class AnnSearchTest {

作成したBytemanのスクリプトはこちらです。

kNN検索かANNのどちらを行うかの条件であるKnnCollector#kとRandomVectorScorer#maxOrdの記録と、ANNを実行する際に使われる
HnswGraphSearcher#searchのトレース。HnswGraphSearcher#searchが呼び出されていない場合は、kNN検索が実行されていることに
なります。

src/test/resources/trace-ann-searcher.btm

RULE print KnnCollector k value
INTERFACE org.apache.lucene.search.KnnCollector
METHOD k
AT EXIT
IF TRUE
  DO traceln("KnnCollector#k = " + $!)
ENDRULE

RULE print RandomVectorScorer maxOrd value
INTERFACE org.apache.lucene.util.hnsw.RandomVectorScorer
METHOD maxOrd
AT EXIT
IF TRUE
  DO traceln("RandomVectorScorer#maxOrd = " + $!)
ENDRULE

RULE print HnswGraphSearcher called
CLASS org.apache.lucene.util.hnsw.HnswGraphSearcher
METHOD search(org.apache.lucene.util.hnsw.RandomVectorScorer, org.apache.lucene.search.KnnCollector, org.apache.lucene.util.hnsw.HnswGraph, org.apache.lucene.util.hnsw.HnswGraphSearcher, org.apache.lucene.util.Bits)
AT ENTRY
IF TRUE
  DO traceln("execute ANN(HNSW) HnswGraphSearcher#search")
ENDRULE

また、インデックスの書き込み時に使われるm(maxConn)、efConstruction(beamWidth)の値も記録するようにしましょう。

src/test/resources/trace-hnsw-parameters.btm

RULE print HNSW parameters
CLASS org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter
METHOD <init>
AT ENTRY
IF TRUE
  DO traceln("m = " + $2 + ", beamWith = " + $3)
ENDRULE

こちらを使って確認していきたいと思います。

kNN検索を行う

まず最初はkNN検索を行ってみましょう。

    @Test
    void knnSearch() throws IOException {
        try (Directory directory = new ByteBuffersDirectory()) {
            IndexWriterConfig config = new IndexWriterConfig();

            List<Document> documents = createDocuments();

            try (IndexWriter writer = new IndexWriter(directory, config)) {
                writer.addDocuments(documents);
            }

            try (DirectoryReader reader = DirectoryReader.open(directory)) {
                IndexSearcher searcher = new IndexSearcher(reader);

                int k = documents.size(); // all document

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

                TopDocs topDocs = searcher.search(query, 3);

                StoredFields storedFields = searcher.storedFields();

                List<Document> resultDocuments =
                        Arrays.stream(topDocs.scoreDocs)
                                .map(scoreDoc -> {
                                    try {
                                        return storedFields.document(scoreDoc.doc);
                                    } catch (IOException e) {
                                        throw new UncheckedIOException(e);
                                    }
                                })
                                .toList();

                assertThat(resultDocuments).hasSize(3);
                assertThat(resultDocuments.get(0).getField("name").stringValue())
                        .isEqualTo("The Hitchhiker's Guide to the Galaxy");
                assertThat(resultDocuments.get(0).getField("year").numericValue().intValue())
                        .isEqualTo(1979);
                assertThat(resultDocuments.get(1).getField("name").stringValue())
                        .isEqualTo("The Three-Body Problem");
                assertThat(resultDocuments.get(1).getField("year").numericValue().intValue())
                        .isEqualTo(2008);
                assertThat(resultDocuments.get(2).getField("name").stringValue())
                        .isEqualTo("The Andromeda Strain");
                assertThat(resultDocuments.get(2).getField("year").numericValue().intValue())
                        .isEqualTo(1969);
            }
        }
    }

KnnFloatVectorQueryに指定するkはドキュメントの数と同じにします。kはベクトル検索で検索対象とするドキュメントの数ですね。
厳密なことを言うと、比較対象はインデックス内の全ドキュメントの数というよりはその中のベクトルの数なのでえすが、今回は
すべてのベクトル用のフィールドに値を設定しているので全ベクトル数=全ドキュメントの数になっています。

                int k = documents.size(); // all document

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

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

Bytemanを使って差し込んだスクリプトで、実行時に得られたログはこちら。

KnnCollector#kとRandomVectorScorer#maxOrdの値。
※たくさん出力されます

KnnCollector#k = 13
RandomVectorScorer#maxOrd = 13

そして、HnswGraphSearcher#searchを呼び出したログはありませんでした。

m(maxConn)、efConstruction(beamWidth)の値はこちら。デフォルト値ですね。

m = 16, beamWith = 100

ANNを行う

次はANNを行ってみます。

    @Test
    void annSearch() throws IOException {
        try (Directory directory = new ByteBuffersDirectory()) {
            IndexWriterConfig config = new IndexWriterConfig();

            List<Document> documents = createDocuments();

            try (IndexWriter writer = new IndexWriter(directory, config)) {
                writer.addDocuments(documents);
            }

            try (DirectoryReader reader = DirectoryReader.open(directory)) {
                IndexSearcher searcher = new IndexSearcher(reader);

                int k = documents.size() - 1; // all documents - 1

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

                TopDocs topDocs = searcher.search(query, 3);

                StoredFields storedFields = searcher.storedFields();

                List<Document> resultDocuments =
                        Arrays.stream(topDocs.scoreDocs)
                                .map(scoreDoc -> {
                                    try {
                                        return storedFields.document(scoreDoc.doc);
                                    } catch (IOException e) {
                                        throw new UncheckedIOException(e);
                                    }
                                })
                                .toList();

                assertThat(resultDocuments).hasSize(3);
                assertThat(resultDocuments.get(0).getField("name").stringValue())
                        .isEqualTo("The Hitchhiker's Guide to the Galaxy");
                assertThat(resultDocuments.get(0).getField("year").numericValue().intValue())
                        .isEqualTo(1979);
                assertThat(resultDocuments.get(1).getField("name").stringValue())
                        .isEqualTo("The Three-Body Problem");
                assertThat(resultDocuments.get(1).getField("year").numericValue().intValue())
                        .isEqualTo(2008);
                assertThat(resultDocuments.get(2).getField("name").stringValue())
                        .isEqualTo("The Andromeda Strain");
                assertThat(resultDocuments.get(2).getField("year").numericValue().intValue())
                        .isEqualTo(1969);
            }
        }
    }

kNN検索との違いは、この1点です。

                int k = documents.size() - 1; // all documents - 1

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

KnnFloatVectorQueryに指定するkはドキュメントの数から1引いています。ここでもすべてのドキュメントでベクトルに値を保存して
いるので、インデックス内の全ベクトル数は全ドキュメントの数と同じになります。

Bytemanを使って差し込んだスクリプトで、実行時に得られたログはこちら。

KnnCollector#kとRandomVectorScorer#maxOrdの値。

KnnCollector#k = 12
RandomVectorScorer#maxOrd = 13

そして、HnswGraphSearcher#searchを呼び出したことを表すログが記録されました。

execute ANN(HNSW) HnswGraphSearcher#search

m(maxConn)、efConstruction(beamWidth)の値はデフォルト値で同じ値ですね。

m = 16, beamWith = 100

というわけで、検索対象とするドキュメント数を、インデックスに保存している全ベクトル数(≒全ドキュメントの数)よりも
小さい値を指定すると近似となりANNに切り替わることがわかりました。

デフォルトのCodecを確認する

次は、HNSWのパラメーターであるmとefConstructionの確認に移ります。

まずはデフォルトのCodecを確認してみましょう。

    @Test
    void defaultCodec() throws ReflectiveOperationException {
        Codec codec = Codec.getDefault();
        assertThat(codec).isExactlyInstanceOf(Lucene99Codec.class);
        assertThat(codec.knnVectorsFormat()).isInstanceOf(PerFieldKnnVectorsFormat.class);

        Lucene99Codec lucene99Codec = (Lucene99Codec) codec;

        java.lang.reflect.Field defaultKnnVectorsFormatField =
                Lucene99Codec.class.getDeclaredField("defaultKnnVectorsFormat");
        defaultKnnVectorsFormatField.setAccessible(true);
        KnnVectorsFormat knnVectorsFormat = (KnnVectorsFormat) defaultKnnVectorsFormatField.get(lucene99Codec);
        assertThat(knnVectorsFormat).isExactlyInstanceOf(Lucene99HnswVectorsFormat.class);
    }

このCodec#getDefaultがIndexWriterConfigで使われるCodecのデフォルトのインスタンスとなっています。

Apache Lucene 9.10.0ではLucene99Codecになります。

Lucene99Codecが使用するKnnVectorsFormatの実装は、Lucene99HnswVectorsFormatになっていることも確認できました。

独自のCodecを作成する

それでは、HNSWのパラメーターであるm(maxConn)、efConstruction(beamWidth)をカスタマイズするために独自の
Codecを作成します。

最初はKnnVectorsFormatのインスタンスだけ作ればいいのかと思っていたのですが、Codecは変更できないのでCodec自体を作成する
必要があります。

となると、Codecが要求するメソッドをすべて実装する必要があるのかなと思ったりするのですが

Codec (Lucene 9.10.0 core API)

そんなことはなく、FilterCodecを使うと既存のCodecの具象クラスを使って差分を実装することができます。

FilterCodec (Lucene 9.10.0 core API)

今回はLucene99Codecをベースにします。

Lucene99Codec (Lucene 9.10.0 core API)

そんなわけで、作成した独自のCodecはこちら。

src/test/java/org/littlewings/lucene/ann/MyCustomCodec.java

package org.littlewings.lucene.ann;

import org.apache.lucene.codecs.FilterCodec;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.codecs.lucene99.Lucene99Codec;
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;

public class MyCustomCodec extends FilterCodec {
    private KnnVectorsFormat knnVectorsFormat = new MyKnnVectorsFormat();

    public MyCustomCodec() {
        super("MyCustom", new Lucene99Codec());
    }

    @Override
    public KnnVectorsFormat knnVectorsFormat() {
        return knnVectorsFormat;
    }

    static class MyKnnVectorsFormat extends PerFieldKnnVectorsFormat {
        private KnnVectorsFormat defaultKnnVectorsFormat = new Lucene99HnswVectorsFormat();
        private KnnVectorsFormat forDescriptionVectorField = new Lucene99HnswVectorsFormat(32, 150);

        @Override
        public KnnVectorsFormat getKnnVectorsFormatForField(String field) {
            System.out.printf("KnnVectorsFormat#getKnnVectorsFormatForField = %s%n", field);

            if ("description_vector".equals(field)) {
                return forDescriptionVectorField;
            } else {
                return defaultKnnVectorsFormat;
            }
        }
    }
}

親コンストラクタには、Codecの名前と委譲先になるCodecの実装のインスタンスを渡します。オーバーライドしたメソッド以外は、
このインスタンスに処理が転送されることになります。

    public MyCustomCodec() {
        super("MyCustom", new Lucene99Codec());
    }

今回はdescription_vectorフィールドのみカスタマイズしたLucene99HnswVectorsFormatのインスタンスを割り当て、こちらの
m(maxConn)を32に、efConstruction(beamWidth)を150にします。

    private KnnVectorsFormat knnVectorsFormat = new MyKnnVectorsFormat();

    〜省略〜

    @Override
    public KnnVectorsFormat knnVectorsFormat() {
        return knnVectorsFormat;
    }

    static class MyKnnVectorsFormat extends PerFieldKnnVectorsFormat {
        private KnnVectorsFormat defaultKnnVectorsFormat = new Lucene99HnswVectorsFormat();
        private KnnVectorsFormat forDescriptionVectorField = new Lucene99HnswVectorsFormat(32, 150);

        @Override
        public KnnVectorsFormat getKnnVectorsFormatForField(String field) {
            System.out.printf("KnnVectorsFormat#getKnnVectorsFormatForField = %s%n", field);

            if ("description_vector".equals(field)) {
                return forDescriptionVectorField;
            } else {
                return defaultKnnVectorsFormat;
            }
        }
    }

呼び出された時にはフィールド名を出力するようにしています。

description_vector以外のフィールドにはデフォルトのLucene99HnswVectorsFormatのインスタンスを使います。

このCodecはServiceLoaderの仕組みで呼び出すので、以下のファイルを作成して実装クラスの名前を書いておきます。

src/test/resources/META-INF/services/org.apache.lucene.codecs.Codec

org.littlewings.lucene.ann.MyCustomCodec

こうすることで、Codec#forNameでこのインスタンスを取得できます。Codec#forNameで指定する値は、Codecを作成する時に
親コンストラクタに指定した名前です。

    @Test
    void customizeHsnwParamersCodec1() {
        Codec codec = Codec.forName("MyCustom");  // load from ServiceLoader
        assertThat(codec).isExactlyInstanceOf(MyCustomCodec.class);
    }

もっとも、ServiceLoaderの仕組みに従わずにふつうにnewしてもいいのですが、お作法に則ると前述のやり方ということになるでしょう。

    @Test
    void customizeHsnwParamersCodec2() {
        Codec codec = new MyCustomCodec();
    }

では、このCodecを使ってインデックスにドキュメントを登録、検索してみます。

その前に、利用するドキュメントを少し変えたいと思います。

    Document createDocument2(String name, String description, String author, int year) {
        Document document = new Document();
        document.add(new TextField("name", name, Field.Store.YES));
        document.add(new TextField("description", description, Field.Store.YES));
        document.add(new TextField("author", author, Field.Store.YES));
        document.add(new IntField("year", year, Field.Store.YES));

        float[] vector = textToVector("passage: " + description);
        document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));
        document.add(new KnnFloatVectorField("description_vector2", vector, VectorSimilarityFunction.EUCLIDEAN));

        return document;
    }

ベクトル用のフィールドを2つにしました。片方はデフォルトのm(maxConn)とefConstruction(beamWidth)を使います。

確認。

    @Test
    void customizeHnswParameters() throws IOException {
        try (Directory directory = new ByteBuffersDirectory()) {
            IndexWriterConfig config = new IndexWriterConfig();
            Codec codec = Codec.forName("MyCustom");
            config.setCodec(codec);

            List<Document> documents = List.of(
                    createDocument2(
                            "The Time Machine",
                            "A man travels through time and witnesses the evolution of humanity.",
                            "H.G. Wells",
                            1895
                    ),
                    createDocument2(
                            "Ender's Game",
                            "A young boy is trained to become a military leader in a war against an alien race.",
                            "Orson Scott Card",
                            1985
                    ),
                    createDocument2(
                            "Brave New World",
                            "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.",
                            "Aldous Huxley",
                            1932
                    )
            );

            try (IndexWriter writer = new IndexWriter(directory, config)) {
                writer.addDocuments(documents);
            }

            try (DirectoryReader reader = DirectoryReader.open(directory)) {
                IndexSearcher searcher = new IndexSearcher(reader);

                int k = 1; // ANN

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

                TopDocs topDocs = searcher.search(query, 1);

                StoredFields storedFields = searcher.storedFields();

                List<Document> resultDocuments =
                        Arrays.stream(topDocs.scoreDocs)
                                .map(scoreDoc -> {
                                    try {
                                        return storedFields.document(scoreDoc.doc);
                                    } catch (IOException e) {
                                        throw new UncheckedIOException(e);
                                    }
                                })
                                .toList();

                assertThat(resultDocuments).hasSize(1);
                assertThat(resultDocuments.get(0).getField("name").stringValue())
                        .isEqualTo("Ender's Game");
                assertThat(resultDocuments.get(0).getField("year").numericValue().intValue())
                        .isEqualTo(1985);
            }
        }
    }

Codecを設定しているのはこちら。

            IndexWriterConfig config = new IndexWriterConfig();
            Codec codec = Codec.forName("MyCustom");
            config.setCodec(codec);

検索はANNにしました。

                int k = 1; // ANN

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

ログを確認しましょう。

KnnVectorsFormat#getKnnVectorsFormatForField = description_vector
KnnVectorsFormat#getKnnVectorsFormatForField = description_vector
m = 32, beamWith = 150
KnnVectorsFormat#getKnnVectorsFormatForField = description_vector2
KnnVectorsFormat#getKnnVectorsFormatForField = description_vector2
m = 16, beamWith = 100

description_vectorの方は、カスタマイズしたm(maxConn)とefConstruction(beamWidth)の値が使われていることが確認できます。

全体のドキュメント数が3なのに対して、検索対象としたドキュメント数は1なのでANNになりました。

KnnCollector#k = 1
RandomVectorScorer#maxOrd = 3
KnnCollector#k = 1
KnnCollector#k = 1
execute ANN(HNSW) HnswGraphSearcher#search

これで、確認したいことはできたかなと。

オマケ

最後に、今回作成したテストコード全体を載せておきます。

src/test/java/org/littlewings/lucene/ann/AnnSearchTest.java

package org.littlewings.lucene.ann;

import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.codecs.lucene99.Lucene99Codec;
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnFloatVectorQuery;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.ByteBuffersDirectory;
import org.apache.lucene.store.Directory;
import org.jboss.byteman.contrib.bmunit.BMScript;
import org.jboss.byteman.contrib.bmunit.BMScripts;
import org.jboss.byteman.contrib.bmunit.BMUnitConfig;
import org.jboss.byteman.contrib.bmunit.WithByteman;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Arrays;
import java.util.List;

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

@WithByteman
@BMUnitConfig(loadDirectory = "target/test-classes")
@BMScripts(scripts = {
        @BMScript("trace-ann-searcher.btm"),
        @BMScript("trace-hnsw-parameters.btm"),
}
)
class AnnSearchTest {
    EmbeddingClient embeddingClient;

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

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

    List<Document> createDocuments() {
        return List.of(
                createDocument(
                        "The Time Machine",
                        "A man travels through time and witnesses the evolution of humanity.",
                        "H.G. Wells",
                        1895
                ),
                createDocument(
                        "Ender's Game",
                        "A young boy is trained to become a military leader in a war against an alien race.",
                        "Orson Scott Card",
                        1985
                ),
                createDocument(
                        "Brave New World",
                        "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.",
                        "Aldous Huxley",
                        1932
                ),
                createDocument(
                        "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
                ),
                createDocument(
                        "Dune",
                        "A desert planet is the site of political intrigue and power struggles.",
                        "Frank Herbert",
                        1965
                ),
                createDocument(
                        "Foundation",
                        "A mathematician develops a science to predict the future of humanity and works to save civilization from collapse.",
                        "Isaac Asimov",
                        1951
                ),
                createDocument(
                        "Snow Crash",
                        "A futuristic world where the internet has evolved into a virtual reality metaverse.",
                        "Neal Stephenson",
                        1992
                ),
                createDocument(
                        "Neuromancer",
                        "A hacker is hired to pull off a near-impossible hack and gets pulled into a web of intrigue.",
                        "William Gibson",
                        1984
                ),
                createDocument(
                        "The War of the Worlds",
                        "A Martian invasion of Earth throws humanity into chaos.",
                        "H.G. Wells",
                        1898
                ),
                createDocument(
                        "The Hunger Games",
                        "A dystopian society where teenagers are forced to fight to the death in a televised spectacle.",
                        "Suzanne Collins",
                        2008
                ),
                createDocument(
                        "The Andromeda Strain",
                        "A deadly virus from outer space threatens to wipe out humanity.",
                        "Michael Crichton",
                        1969
                ),
                createDocument(
                        "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
                ),
                createDocument(
                        "The Three-Body Problem",
                        "Humans encounter an alien civilization that lives in a dying system.",
                        "Liu Cixin",
                        2008
                )
        );
    }

    Document createDocument(String name, String description, String author, int year) {
        Document document = new Document();
        document.add(new TextField("name", name, Field.Store.YES));
        document.add(new TextField("description", description, Field.Store.YES));
        document.add(new TextField("author", author, Field.Store.YES));
        document.add(new IntField("year", year, Field.Store.YES));

        float[] vector = textToVector("passage: " + description);
        document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));

        return document;
    }

    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;
    }

    @Test
    void knnSearch() throws IOException {
        try (Directory directory = new ByteBuffersDirectory()) {
            IndexWriterConfig config = new IndexWriterConfig();

            List<Document> documents = createDocuments();

            try (IndexWriter writer = new IndexWriter(directory, config)) {
                writer.addDocuments(documents);
            }

            try (DirectoryReader reader = DirectoryReader.open(directory)) {
                IndexSearcher searcher = new IndexSearcher(reader);

                int k = documents.size(); // all document

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

                TopDocs topDocs = searcher.search(query, 3);

                StoredFields storedFields = searcher.storedFields();

                List<Document> resultDocuments =
                        Arrays.stream(topDocs.scoreDocs)
                                .map(scoreDoc -> {
                                    try {
                                        return storedFields.document(scoreDoc.doc);
                                    } catch (IOException e) {
                                        throw new UncheckedIOException(e);
                                    }
                                })
                                .toList();

                assertThat(resultDocuments).hasSize(3);
                assertThat(resultDocuments.get(0).getField("name").stringValue())
                        .isEqualTo("The Hitchhiker's Guide to the Galaxy");
                assertThat(resultDocuments.get(0).getField("year").numericValue().intValue())
                        .isEqualTo(1979);
                assertThat(resultDocuments.get(1).getField("name").stringValue())
                        .isEqualTo("The Three-Body Problem");
                assertThat(resultDocuments.get(1).getField("year").numericValue().intValue())
                        .isEqualTo(2008);
                assertThat(resultDocuments.get(2).getField("name").stringValue())
                        .isEqualTo("The Andromeda Strain");
                assertThat(resultDocuments.get(2).getField("year").numericValue().intValue())
                        .isEqualTo(1969);
            }
        }
    }

    @Test
    void annSearch() throws IOException {
        try (Directory directory = new ByteBuffersDirectory()) {
            IndexWriterConfig config = new IndexWriterConfig();

            List<Document> documents = createDocuments();

            try (IndexWriter writer = new IndexWriter(directory, config)) {
                writer.addDocuments(documents);
            }

            try (DirectoryReader reader = DirectoryReader.open(directory)) {
                IndexSearcher searcher = new IndexSearcher(reader);

                int k = documents.size() - 1; // all documents - 1

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

                TopDocs topDocs = searcher.search(query, 3);

                StoredFields storedFields = searcher.storedFields();

                List<Document> resultDocuments =
                        Arrays.stream(topDocs.scoreDocs)
                                .map(scoreDoc -> {
                                    try {
                                        return storedFields.document(scoreDoc.doc);
                                    } catch (IOException e) {
                                        throw new UncheckedIOException(e);
                                    }
                                })
                                .toList();

                assertThat(resultDocuments).hasSize(3);
                assertThat(resultDocuments.get(0).getField("name").stringValue())
                        .isEqualTo("The Hitchhiker's Guide to the Galaxy");
                assertThat(resultDocuments.get(0).getField("year").numericValue().intValue())
                        .isEqualTo(1979);
                assertThat(resultDocuments.get(1).getField("name").stringValue())
                        .isEqualTo("The Three-Body Problem");
                assertThat(resultDocuments.get(1).getField("year").numericValue().intValue())
                        .isEqualTo(2008);
                assertThat(resultDocuments.get(2).getField("name").stringValue())
                        .isEqualTo("The Andromeda Strain");
                assertThat(resultDocuments.get(2).getField("year").numericValue().intValue())
                        .isEqualTo(1969);
            }
        }
    }

    @Test
    void defaultCodec() throws ReflectiveOperationException {
        Codec codec = Codec.getDefault();
        assertThat(codec).isExactlyInstanceOf(Lucene99Codec.class);
        assertThat(codec.knnVectorsFormat()).isInstanceOf(PerFieldKnnVectorsFormat.class);

        Lucene99Codec lucene99Codec = (Lucene99Codec) codec;

        java.lang.reflect.Field defaultKnnVectorsFormatField =
                Lucene99Codec.class.getDeclaredField("defaultKnnVectorsFormat");
        defaultKnnVectorsFormatField.setAccessible(true);
        KnnVectorsFormat knnVectorsFormat = (KnnVectorsFormat) defaultKnnVectorsFormatField.get(lucene99Codec);
        assertThat(knnVectorsFormat).isExactlyInstanceOf(Lucene99HnswVectorsFormat.class);
    }

    @Test
    void customizeHsnwParamersCodec1() {
        Codec codec = Codec.forName("MyCustom");  // load from ServiceLoader
        assertThat(codec).isExactlyInstanceOf(MyCustomCodec.class);
    }

    @Test
    void customizeHsnwParamersCodec2() {
        Codec codec = new MyCustomCodec();
    }

    Document createDocument2(String name, String description, String author, int year) {
        Document document = new Document();
        document.add(new TextField("name", name, Field.Store.YES));
        document.add(new TextField("description", description, Field.Store.YES));
        document.add(new TextField("author", author, Field.Store.YES));
        document.add(new IntField("year", year, Field.Store.YES));

        float[] vector = textToVector("passage: " + description);
        document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));
        document.add(new KnnFloatVectorField("description_vector2", vector, VectorSimilarityFunction.EUCLIDEAN));

        return document;
    }

    @Test
    void customizeHnswParameters() throws IOException {
        try (Directory directory = new ByteBuffersDirectory()) {
            IndexWriterConfig config = new IndexWriterConfig();
            Codec codec = Codec.forName("MyCustom");
            config.setCodec(codec);

            List<Document> documents = List.of(
                    createDocument2(
                            "The Time Machine",
                            "A man travels through time and witnesses the evolution of humanity.",
                            "H.G. Wells",
                            1895
                    ),
                    createDocument2(
                            "Ender's Game",
                            "A young boy is trained to become a military leader in a war against an alien race.",
                            "Orson Scott Card",
                            1985
                    ),
                    createDocument2(
                            "Brave New World",
                            "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.",
                            "Aldous Huxley",
                            1932
                    )
            );

            try (IndexWriter writer = new IndexWriter(directory, config)) {
                writer.addDocuments(documents);
            }

            try (DirectoryReader reader = DirectoryReader.open(directory)) {
                IndexSearcher searcher = new IndexSearcher(reader);

                int k = 1; // ANN

                float[] vector = textToVector("query: alien invasion");
                KnnFloatVectorQuery query =
                        new KnnFloatVectorQuery("description_vector", vector, k);

                TopDocs topDocs = searcher.search(query, 1);

                StoredFields storedFields = searcher.storedFields();

                List<Document> resultDocuments =
                        Arrays.stream(topDocs.scoreDocs)
                                .map(scoreDoc -> {
                                    try {
                                        return storedFields.document(scoreDoc.doc);
                                    } catch (IOException e) {
                                        throw new UncheckedIOException(e);
                                    }
                                })
                                .toList();

                assertThat(resultDocuments).hasSize(1);
                assertThat(resultDocuments.get(0).getField("name").stringValue())
                        .isEqualTo("Ender's Game");
                assertThat(resultDocuments.get(0).getField("year").numericValue().intValue())
                        .isEqualTo(1985);
            }
        }
    }
}

おわりに

Apache LuceneでkNN検索とANNの切り替え方および、ANNで使われるHNSWのパラメーターの設定方法を見てみました。

けっこう大変だったのですが、使っているのがkNN検索なのかANNなのかぼやっとしていることが多かったので、こうやって切り替え条件が
わかるとだいぶクリアになりました。

また、mやefConstructionのカスタマイズ方法やどのように反映されるのかが確認できたのもよかったです。