diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b897d49c50..656a87a7b9 100755 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -11,6 +11,10 @@ examples/cpp/ @NVIDIA/cuvs-cpp-codeowners java/ @NVIDIA/cuvs-java-codeowners examples/java/ @NVIDIA/cuvs-java-codeowners +#cuvs lucene code owners +# TODO: should we have this? +#java/cuvs-lucene @NVIDIA/cuvs-lucene-write + #python code owners python/ @NVIDIA/cuvs-python-codeowners diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 804c3cd877..78785a6a00 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -184,6 +184,39 @@ jobs: artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" sha: ${{ inputs.sha }} + lucene-build-matrix: + needs: cpp-build + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@main + with: + build_type: pull-request + matrix_name: conda-cpp-build + matrix_filter: map(select(.ARCH == "amd64")) + lucene-build: + # Consumes the cuvs-java artifact uploaded by java-build instead of rebuilding cuvs-java. + needs: [lucene-build-matrix, java-build] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@main + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.lucene-build-matrix.outputs.matrix) }} + with: + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + arch: "amd64" + date: ${{ inputs.date }} + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/build_lucene.sh" + artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" + file_to_upload: "java/cuvs-lucene/target/" + sha: ${{ inputs.sha }} python-build: needs: [cpp-build] permissions: diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2154fe74a7..ad0af5eb52 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -25,6 +25,8 @@ jobs: - rocky8-clib-tests - conda-java-build-and-tests-matrix - conda-java-build-and-tests + - conda-lucene-build-and-tests-matrix + - conda-lucene-build-and-tests - rust-build-matrix - rust-build - go-build-matrix @@ -548,6 +550,40 @@ jobs: script: "ci/test_java.sh" artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" + conda-lucene-build-and-tests-matrix: + needs: [conda-cpp-build, changed-files] + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@main + # cuvs-lucene lives under java/, so java/cuvs-lucene changes trigger the test_java group. + if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp + with: + build_type: pull-request + matrix_name: conda-cpp-build + matrix_filter: map(select(.ARCH == "amd64")) + conda-lucene-build-and-tests: + # Depends on conda-java-build-and-tests to consume its uploaded cuvs-java artifact + # (installed into the local Maven repo) instead of rebuilding cuvs-java. + needs: [conda-java-build-and-tests, conda-lucene-build-and-tests-matrix] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@main + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-lucene-build-and-tests-matrix.outputs.matrix) }} + with: + build_type: pull-request + node_type: "gpu-l4-latest-1" + arch: "amd64" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_lucene.sh" + artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" + file_to_upload: "java/cuvs-lucene/target/" rust-build-matrix: needs: [conda-cpp-build, changed-files] permissions: diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 26739e2deb..61e4aefac5 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -104,6 +104,40 @@ jobs: script: "ci/test_java.sh" artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" + conda-lucene-build-and-tests-matrix: + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@main + with: + build_type: ${{ inputs.build_type }} + matrix_name: conda-cpp-build + matrix_filter: map(select(.ARCH == "amd64")) + conda-lucene-build-and-tests: + # Consumes the cuvs-java artifact uploaded by conda-java-build-and-tests instead of + # rebuilding cuvs-java. + needs: [conda-java-build-and-tests, conda-lucene-build-and-tests-matrix] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@main + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-lucene-build-and-tests-matrix.outputs.matrix) }} + with: + build_type: ${{ inputs.build_type }} + branch: ${{ inputs.branch }} + date: ${{ inputs.date }} + sha: ${{ inputs.sha }} + node_type: "gpu-l4-latest-1" + arch: "amd64" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_lucene.sh" + artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" + file_to_upload: "java/cuvs-lucene/target/" wheel-tests-cuvs: permissions: actions: read diff --git a/build.sh b/build.sh index d5780b481a..7eeb877e44 100755 --- a/build.sh +++ b/build.sh @@ -19,7 +19,7 @@ ARGS=$* # scripts, and that this script resides in the repo dir! REPODIR=$(cd "$(dirname "$0")"; pwd) -VALIDARGS="clean libcuvs python rust go java docs tests bench-ann examples --uninstall -v -g -n --allgpuarch --no-mg --mnmg-tests --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h --run-java-tests" +VALIDARGS="clean libcuvs python rust go java lucene docs tests bench-ann examples --uninstall -v -g -n --allgpuarch --no-mg --mnmg-tests --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h --run-java-tests" HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool=] [--limit-tests=] [--limit-bench-ann=] [--build-metrics=] where is: clean - remove all existing build artifacts and configuration (start over) @@ -29,6 +29,7 @@ HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool= void _copy_matrix(cuvsResources_t res, DLManagedTensor* src_managed, DLManagedTensor* dst_managed) diff --git a/ci/build_lucene.sh b/ci/build_lucene.sh new file mode 100755 index 0000000000..4d1dbad619 --- /dev/null +++ b/ci/build_lucene.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# TODO: Remove this argument-handling when build and test workflows are separated, +# and test_lucene.sh no longer calls build_lucene.sh +# ref: https://github.com/rapidsai/cuvs/issues/868 +EXTRA_BUILD_ARGS=() +if [[ "${1:-}" == "--run-java-tests" ]]; then + EXTRA_BUILD_ARGS+=("--run-java-tests") +fi + +if [ -e "/opt/conda/etc/profile.d/conda.sh" ]; then + . /opt/conda/etc/profile.d/conda.sh +fi + +rapids-logger "Configuring conda strict channel priority" +conda config --set channel_priority strict + +rapids-logger "Downloading artifacts from previous jobs" +# libcuvs C++ conda package (provides libcuvs.so / libcuvs_c.so used at test time) +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") +# cuvs-java jar built by the java-build job (avoids rebuilding cuvs-java here) +CUVS_JAVA_ARTIFACT_DIR=$(rapids-download-from-github "cuvs-java-cuda${RAPIDS_CUDA_VERSION}") + +rapids-logger "Generate Java testing dependencies" + +ENV_YAML_DIR="$(mktemp -d)" + +rapids-dependency-file-generator \ + --output conda \ + --file-key java \ + --prepend-channel "${CPP_CHANNEL}" \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" | tee "${ENV_YAML_DIR}/env.yaml" + +rapids-mamba-retry env create --yes -f "${ENV_YAML_DIR}/env.yaml" -n java + +# Temporarily allow unbound variables for conda activation. +set +u +conda activate java +set -u + +rapids-print-env + +# cuvs-lucene depends on the plain cuvs-java jar (no bundled native libraries), so it loads +# libcuvs through the system loader. Make the conda-provided libcuvs discoverable at test time. +export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${LD_LIBRARY_PATH:-}" + +# Install the cuvs-java artifact built by the java-build job into the local Maven repository +# so that cuvs-lucene resolves com.nvidia.cuvs:cuvs-java without rebuilding it. The GAV is +# taken from the artifact's own pom.xml (bundled alongside the jar by java/build.sh). +VERSION=$(grep -oP 'VERSION="\K[^"]+' java/build.sh) +CUVS_JAVA_JAR=$(find "${CUVS_JAVA_ARTIFACT_DIR}" -name "cuvs-java-${VERSION}.jar" | head -1) +CUVS_JAVA_POM=$(find "${CUVS_JAVA_ARTIFACT_DIR}" -name "pom.xml" | head -1) +if [[ -z "${CUVS_JAVA_JAR}" || -z "${CUVS_JAVA_POM}" ]]; then + echo "Error: could not locate cuvs-java-${VERSION}.jar and/or pom.xml in ${CUVS_JAVA_ARTIFACT_DIR}" + exit 1 +fi +rapids-logger "Installing cuvs-java ${VERSION} from ${CUVS_JAVA_JAR}" +mvn install:install-file -Dfile="${CUVS_JAVA_JAR}" -DpomFile="${CUVS_JAVA_POM}" + +EXITCODE=0 +trap "EXITCODE=1" ERR +set +e + +rapids-logger "Run cuVS Lucene build" + +RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" +export RAPIDS_CUDA_MAJOR + +bash ./build.sh lucene "${EXTRA_BUILD_ARGS[@]}" + +rapids-logger "Build script exiting with value: $EXITCODE" +exit ${EXITCODE} diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index f36c14f771..81df4f7dc0 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 ######################## # CUVS Version Updater # @@ -175,3 +175,11 @@ done sed_runner "s| CuVS [[:digit:]]\{2\}\.[[:digit:]]\{2\} | CuVS ${NEXT_SHORT_TAG} |g" java/README.md sed_runner "s|-[[:digit:]]\{2\}\.[[:digit:]]\{2\}\.[[:digit:]]\{1,2\}\.jar|-${NEXT_FULL_JAVA_TAG}\.jar|g" java/examples/README.md sed_runner "s|/[[:digit:]]\{2\}\.[[:digit:]]\{2\}\.[[:digit:]]\{1,2\}/|/${NEXT_FULL_JAVA_TAG}/|g" java/examples/README.md + +# Update cuvs-lucene (integrated Java subproject; uses its own CUVS_LUCENE# markers) +sed_runner "s/VERSION=\".*\"/VERSION=\"${NEXT_FULL_JAVA_TAG}\"/g" java/cuvs-lucene/build.sh +for FILE in java/cuvs-lucene/pom.xml java/cuvs-lucene/examples/pom.xml java/cuvs-lucene/bench/pom.xml; do + sed_runner "/.*/s//${NEXT_FULL_JAVA_TAG}<\/version>/g" "${FILE}" +done +sed_runner "s|rapidsai/ci-conda:[\.0-9]*-cuda|rapidsai/ci-conda:${NEXT_SHORT_TAG}-cuda|g" java/cuvs-lucene/examples/README.md +sed_runner "s|target/examples-[\.0-9]*-jar|target/examples-${NEXT_FULL_JAVA_TAG}-jar|g" java/cuvs-lucene/examples/README.md diff --git a/ci/test_lucene.sh b/ci/test_lucene.sh new file mode 100755 index 0000000000..04e4c4fab5 --- /dev/null +++ b/ci/test_lucene.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +EXITCODE=0 +trap "EXITCODE=1" ERR +set +e + +rapids-logger "Check GPU usage" +nvidia-smi + +rapids-logger "Run cuVS Lucene build and tests" + +RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" +export RAPIDS_CUDA_MAJOR + +# TODO: switch to installing pre-built artifacts instead of rebuilding in test jobs +# ref: https://github.com/rapidsai/cuvs/issues/868 +ci/build_lucene.sh --run-java-tests + +rapids-logger "Test script exiting with value: $EXITCODE" +exit ${EXITCODE} diff --git a/java/.gitignore b/java/.gitignore index 2e98aff7ba..a4da3c4db7 100644 --- a/java/.gitignore +++ b/java/.gitignore @@ -7,3 +7,5 @@ openjdk-22-jextract* /cuvs-java/bin/ /cuvs-java/src/main/java22/com/nvidia/cuvs/internal/panama/ /cuvs-java/*.cag +# cuvs-lucene +/cuvs-lucene/bin/ diff --git a/java/README.md b/java/README.md index ff62ff1225..a14995b761 100644 --- a/java/README.md +++ b/java/README.md @@ -59,3 +59,13 @@ times, by passing `-Dtests.iters=10` through the command line. ## Examples A few starter examples of CAGRA, HNSW, and Bruteforce indexing and searching are provided in the `examples` directory. + +## cuVS Lucene + +The [`cuvs-lucene`](cuvs-lucene) directory contains the cuVS integration for +[Apache Lucene](https://github.com/apache/lucene) — a pluggable `KnnVectorsFormat` that enables +GPU-accelerated vector search. It depends on the `cuvs-java` bindings described above. + +`lucene` is a separate top-level build target: build it with `./build.sh lucene` (or +`./build.sh java lucene` to rebuild `cuvs-java` first, or `./build.sh libcuvs java lucene` to build +everything from scratch). See [`cuvs-lucene/README.md`](cuvs-lucene/README.md) for details. diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 60a2928aec..89a5acf880 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -218,6 +218,18 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams merg /** Disables pooled memory on the current device, reverting back to the default setting. */ void resetRMMPooledMemory(); + /** + * Bogus method used to verify end-to-end build propagation from the cuVS C layer, through the + * cuvs-java bindings, to cuvs-lucene. Returns the sum of its two arguments. + * + * @param a first addend + * @param b second addend + * @return the sum {@code a + b} + */ + default int bogusAddV1(int a, int b) { + throw new UnsupportedOperationException("bogusAddV1 is not supported by this provider"); + } + /** Retrieves the system-wide provider. */ static CuVSProvider provider() { return CuVSServiceProvider.Holder.INSTANCE; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index c9cff1e272..29d73d89b3 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -473,6 +473,11 @@ public void resetRMMPooledMemory() { checkCuVSError(cuvsRMMMemoryResourceResetInvoker.apply(), "cuvsRMMMemoryResourceReset"); } + @Override + public int bogusAddV1(int a, int b) { + return cuvsBogusAddV1(a, b); + } + @Override public CuVSMatrix.Builder newHostMatrixBuilder( long size, long columns, CuVSMatrix.DataType dataType) { diff --git a/java/cuvs-lucene/.mvn/maven.config b/java/cuvs-lucene/.mvn/maven.config new file mode 100644 index 0000000000..c83cb17ff4 --- /dev/null +++ b/java/cuvs-lucene/.mvn/maven.config @@ -0,0 +1,6 @@ +-e +-B +-Daether.connector.basic.downstreamThreads=1 +-Daether.transport.http.retryHandler.count=5 +-Daether.transport.http.retryHandler.interval=10000 +-Dmaven.wagon.http.retryHandler.count=5 diff --git a/java/cuvs-lucene/README.md b/java/cuvs-lucene/README.md new file mode 100644 index 0000000000..b26bf33ef8 --- /dev/null +++ b/java/cuvs-lucene/README.md @@ -0,0 +1,67 @@ +# Lucene cuVS + +This is a project for using [cuVS](https://github.com/rapidsai/cuvs), NVIDIA's GPU accelerated vector search library, with [Apache Lucene](https://github.com/apache/lucene). + +## Overview + +This library provides a new [KnnVectorFormat](https://lucene.apache.org/core/10_3_1/core/org/apache/lucene/codecs/KnnVectorsFormat.html) which can be plugged into a Lucene codec. + +cuvs-lucene is part of the [cuVS](https://github.com/rapidsai/cuvs) repository and depends on +the `com.nvidia.cuvs:cuvs-java` artifact built by [`../cuvs-java`](../cuvs-java). See +[`../README.md`](../README.md) for an overview of the Java projects. + +## Building + +### Prerequisites + +- [CUDA 12.0+](https://developer.nvidia.com/cuda-toolkit-archive) +- [Maven 3.9.6+](https://maven.apache.org/download.cgi) +- [JDK 22](https://jdk.java.net/archive/) +- The `libcuvs` C/C++ libraries and the `cuvs-java` bindings (both built from this repo, see below) + +### From the repository root + +`lucene` is a build target of the top-level `build.sh`, alongside `libcuvs` and `java`: + +```sh +# Build everything cuvs-lucene needs, from scratch: +./build.sh libcuvs java lucene + +# If libcuvs and cuvs-java are already built and installed, just (re)build cuvs-lucene: +./build.sh lucene +``` + +Like `java` (which builds against an already-built `libcuvs`), `lucene` builds against the +already-installed `cuvs-java` artifact, so it can be rebuilt on its own without waiting for +`cuvs-java`. Add `java` (i.e. `./build.sh java lucene`) when `cuvs-java` has changed and needs +rebuilding first. The `cuvs-lucene` jar is produced under `target/`. + +### From this directory + +Once `cuvs-java` has been installed into your local Maven repository (via `../build.sh java`), +you can build cuvs-lucene directly: + +```sh +./build.sh +``` + +### Running Tests + +Append `--run-java-tests` to run the test suite. `build.sh` sets `LD_LIBRARY_PATH` to the in-tree +`../../cpp/build` so the tests can locate `libcuvs`: + +```sh +# from the repository root +./build.sh lucene --run-java-tests + +# or from this directory +./build.sh --run-java-tests +``` + +If your `libcuvs` libraries live elsewhere (e.g. a conda environment), set `LD_LIBRARY_PATH` to +the directory containing the matching `libcuvs_c.so` before running the tests. + +## Contributing + +> [!NOTE] +> The code style format is automatically enforced (including the missing license header, if any) using the [Spotless maven plugin](https://github.com/diffplug/spotless/tree/main/plugin-maven). This currently happens in the maven's `validate` stage. diff --git a/java/cuvs-lucene/bench/README.md b/java/cuvs-lucene/bench/README.md new file mode 100644 index 0000000000..c11db0b630 --- /dev/null +++ b/java/cuvs-lucene/bench/README.md @@ -0,0 +1,9 @@ +# Microbenchmarks + +This maven project provides microbenchmarks for `cuvs-lucene` using JMH. + +To run do (if on current directory): + +```sh +cd .. && mvn clean install -DskipTests && cd benchmarks && mvn clean install && java -jar target/benchmarks.jar +``` diff --git a/java/cuvs-lucene/bench/pom.xml b/java/cuvs-lucene/bench/pom.xml new file mode 100644 index 0000000000..29557d18a9 --- /dev/null +++ b/java/cuvs-lucene/bench/pom.xml @@ -0,0 +1,234 @@ + + + + 4.0.0 + + com.nvidia.cuvs.lucene.benchmarks + cuvs-lucene-benchmarks + 26.10.0 + jar + + cuvs-lucene-benchmarks + + + + + gcs-maven-central-mirror + GCS Maven Central mirror + https://maven-central.storage-download.googleapis.com/maven2/ + + true + + + false + + + + + central + Maven Central + https://repo.maven.apache.org/maven2 + + true + + + false + + + + + + + + gcs-maven-central-mirror + GCS Maven Central mirror + https://maven-central.storage-download.googleapis.com/maven2/ + + true + + + false + + + + + central + Maven Plugin Repository + https://repo.maven.apache.org/maven2 + + true + + + false + + + + + + + org.openjdk.jmh + jmh-core + ${jmh.version} + + + org.openjdk.jmh + jmh-generator-annprocess + ${jmh.version} + provided + + + com.nvidia.cuvs.lucene + cuvs-lucene + 26.10.0 + + + commons-io + commons-io + 2.18.0 + + + + + UTF-8 + 1.37 + 22 + 22 + 22 + benchmarks + + + + + + org.apache.maven.plugins + maven-jar-plugin + 3.4.2 + + + + true + + + true + true + + + + + + com.diffplug.spotless + spotless-maven-plugin + 2.44.5 + + + + apply + + validate + + + + + + src/**/*.java + + + 1.27.0 + + true + false + + + + + + org.apache.maven.plugins + maven-compiler-plugin + 3.8.0 + + ${javac.target} + ${javac.target} + ${javac.target} + + + + org.apache.maven.plugins + maven-shade-plugin + 3.2.1 + + + package + + shade + + + ${uberjar.name} + + + org.openjdk.jmh.Main + + + + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + + + + + + maven-clean-plugin + 2.5 + + + maven-deploy-plugin + 2.8.1 + + + maven-install-plugin + 2.5.1 + + + maven-jar-plugin + 3.4.2 + + + maven-javadoc-plugin + 2.9.1 + + + maven-resources-plugin + 2.6 + + + maven-site-plugin + 3.3 + + + maven-source-plugin + 2.2.1 + + + maven-surefire-plugin + 2.17 + + + + + diff --git a/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/AcceleratedHnswIndexingBenchmarks.java b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/AcceleratedHnswIndexingBenchmarks.java new file mode 100644 index 0000000000..08a41f6a5b --- /dev/null +++ b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/AcceleratedHnswIndexingBenchmarks.java @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.benchmarks; + +import static com.nvidia.cuvs.lucene.benchmarks.Utils.cleanup; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.generateDataset; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.index; + +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.IndexWriterConfig; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@BenchmarkMode(Mode.All) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@State(Scope.Benchmark) +@Fork(value = 1) +@Warmup(iterations = 0, time = 1, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS) +public class AcceleratedHnswIndexingBenchmarks { + + @SuppressWarnings("unused") + private static Logger log = Logger.getLogger(AcceleratedHnswIndexingBenchmarks.class.getName()); + + private static Random random; + private static Path indexDirPath; + private static Codec codec; + + private final int COMMIT_FREQ = 100; + private final String ID_FIELD = "id"; + private final String VECTOR_FIELD = "vector_field"; + + private int numDocs; + private int dimension; + private float[][] dataset; + + @Setup(Level.Trial) + public void setup() throws Exception { + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + numDocs = 1000; + dimension = 128; + dataset = generateDataset(random, numDocs, dimension); + } + + @Benchmark + public void benchmarkAcceleratedHnswIndexing(Blackhole blackhole) throws Exception { + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + index(indexDirPath, config, dataset, numDocs, VECTOR_FIELD, ID_FIELD, COMMIT_FREQ); + } + + @TearDown(Level.Trial) + public void tearDown() throws IOException { + cleanup(indexDirPath); + } +} diff --git a/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/AcceleratedHnswSearchBenchmarks.java b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/AcceleratedHnswSearchBenchmarks.java new file mode 100644 index 0000000000..32c8d6f970 --- /dev/null +++ b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/AcceleratedHnswSearchBenchmarks.java @@ -0,0 +1,86 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.benchmarks; + +import static com.nvidia.cuvs.lucene.benchmarks.Utils.cleanup; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.generateDataset; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.index; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.search; + +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@BenchmarkMode(Mode.All) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@State(Scope.Benchmark) +@Fork(value = 1) +@Warmup(iterations = 0, time = 1, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS) +public class AcceleratedHnswSearchBenchmarks { + + @SuppressWarnings("unused") + private static Logger log = Logger.getLogger(AcceleratedHnswSearchBenchmarks.class.getName()); + + private static Random random; + private static Path indexDirPath; + private static Codec codec; + + private final int COMMIT_FREQ = 100; + private final String ID_FIELD = "id"; + private final String VECTOR_FIELD = "vector_field"; + + private int numDocs; + private int dimension; + private int topK; + private float[][] dataset; + private float[] queryVector; + + @Setup(Level.Trial) + public void setup() throws Exception { + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + numDocs = 1000; + dimension = 128; + topK = 5; + dataset = generateDataset(random, numDocs, dimension); + queryVector = generateDataset(random, 1, dimension)[0]; + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + index(indexDirPath, config, dataset, numDocs, VECTOR_FIELD, ID_FIELD, COMMIT_FREQ); + } + + @Benchmark + public void benchmarkAcceleratedHnswSearch(Blackhole blackhole) throws Exception { + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + search(indexDirPath, VECTOR_FIELD, queryVector, topK, query); + } + + @TearDown(Level.Trial) + public void tearDown() throws IOException { + cleanup(indexDirPath); + } +} diff --git a/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.java b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.java new file mode 100644 index 0000000000..644e852794 --- /dev/null +++ b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.java @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.benchmarks; + +import static com.nvidia.cuvs.lucene.benchmarks.Utils.cleanup; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.generateDataset; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.index; + +import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.IndexWriterConfig; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@BenchmarkMode(Mode.All) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@State(Scope.Benchmark) +@Fork(value = 1) +@Warmup(iterations = 0, time = 1, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS) +public class CagraIndexingBenchmarks { + + @SuppressWarnings("unused") + private static Logger log = Logger.getLogger(CagraIndexingBenchmarks.class.getName()); + + private static Random random; + private static Path indexDirPath; + private static Codec codec; + + private final int COMMIT_FREQ = 100; + private final String ID_FIELD = "id"; + private final String VECTOR_FIELD = "vector_field"; + + private int numDocs; + private int dimension; + private float[][] dataset; + + @Setup(Level.Trial) + public void setup() throws Exception { + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + codec = new CuVS2510GPUSearchCodec(); + numDocs = 1000; + dimension = 128; + dataset = generateDataset(random, numDocs, dimension); + } + + @Benchmark + public void benchmarkAcceleratedHnswIndexing(Blackhole blackhole) throws Exception { + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + index(indexDirPath, config, dataset, numDocs, VECTOR_FIELD, ID_FIELD, COMMIT_FREQ); + } + + @TearDown(Level.Trial) + public void tearDown() throws IOException { + cleanup(indexDirPath); + } +} diff --git a/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.java b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.java new file mode 100644 index 0000000000..9032701095 --- /dev/null +++ b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.java @@ -0,0 +1,88 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.benchmarks; + +import static com.nvidia.cuvs.lucene.benchmarks.Utils.cleanup; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.generateDataset; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.index; +import static com.nvidia.cuvs.lucene.benchmarks.Utils.search; + +import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; +import com.nvidia.cuvs.lucene.GPUKnnFloatVectorQuery; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@BenchmarkMode(Mode.All) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@State(Scope.Benchmark) +@Fork(value = 1) +@Warmup(iterations = 0, time = 1, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS) +public class CagraSearchBenchmarks { + + @SuppressWarnings("unused") + private static Logger log = Logger.getLogger(CagraSearchBenchmarks.class.getName()); + + private static Random random; + private static Path indexDirPath; + private static Codec codec; + + private final int COMMIT_FREQ = 100; + private final String ID_FIELD = "id"; + private final String VECTOR_FIELD = "vector_field"; + + private int numDocs; + private int dimension; + private int topK; + private float[][] dataset; + private float[] queryVector; + + @Setup(Level.Trial) + public void setup() throws Exception { + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + codec = new CuVS2510GPUSearchCodec(); + numDocs = 1000; + dimension = 128; + topK = 5; + dataset = generateDataset(random, numDocs, dimension); + queryVector = generateDataset(random, 1, dimension)[0]; + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + index(indexDirPath, config, dataset, numDocs, VECTOR_FIELD, ID_FIELD, COMMIT_FREQ); + } + + @Benchmark + public void benchmarkAcceleratedHnswSearch(Blackhole blackhole) throws Exception { + KnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK, null, topK, 1); + search(indexDirPath, VECTOR_FIELD, queryVector, topK, query); + } + + @TearDown(Level.Trial) + public void tearDown() throws IOException { + cleanup(indexDirPath); + } +} diff --git a/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/Utils.java b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/Utils.java new file mode 100644 index 0000000000..ceda520be4 --- /dev/null +++ b/java/cuvs-lucene/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/Utils.java @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.benchmarks; + +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.util.Random; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +public class Utils { + + public static void index( + Path indexDirPath, + IndexWriterConfig config, + float[][] dataset, + int numDocs, + String idField, + String vectorField, + int commitFreq) + throws Exception { + int count = commitFreq; + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(idField, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(vectorField, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = commitFreq; + } + } + } + } + + public static void search( + Path indexDirPath, + String vectorField, + float[] queryVector, + int topK, + KnnFloatVectorQuery query) + throws IOException { + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + IndexSearcher searcher = new IndexSearcher(reader); + searcher.search(query, topK); + } + } + + public static void cleanup(Path indexDirPath) throws IOException { + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + public static float[][] generateDataset(Random random, int size, int dimensions) { + float[][] dataset = new float[size][dimensions]; + for (int i = 0; i < size; i++) { + for (int j = 0; j < dimensions; j++) { + dataset[i][j] = random.nextFloat() * 100; + } + } + return dataset; + } + + public static float[] generateRandomVector(int dimensions, Random random) { + float[] vector = new float[dimensions]; + for (int i = 0; i < dimensions; i++) { + vector[i] = random.nextFloat() * 100; + } + return vector; + } + + public static float[][] generateQueries(Random random, int dimensions, int numQueries) { + float[][] queries = new float[numQueries][dimensions]; + for (int i = 0; i < numQueries; i++) { + for (int j = 0; j < dimensions; j++) { + queries[i][j] = random.nextFloat() * 100; + } + } + return queries; + } +} diff --git a/java/cuvs-lucene/build.sh b/java/cuvs-lucene/build.sh new file mode 100755 index 0000000000..103fe83b38 --- /dev/null +++ b/java/cuvs-lucene/build.sh @@ -0,0 +1,37 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -e -u -o pipefail + +ARGS="$*" +NUMARGS=$# + +VERSION="26.10.0" # Note: The version is updated automatically when ci/release/update-version.sh is invoked +GROUP_ID="com.nvidia.cuvs.lucene" + +function hasArg { + (( NUMARGS != 0 )) && (echo " ${ARGS} " | grep -q " $1 ") +} + +if [ -z "${CMAKE_PREFIX_PATH:=}" ]; then + CMAKE_PREFIX_PATH="$(pwd)/../../cpp/build" + export CMAKE_PREFIX_PATH +fi + +if [ -z ${LD_LIBRARY_PATH+x} ]; then + export LD_LIBRARY_PATH=$CMAKE_PREFIX_PATH +else + export LD_LIBRARY_PATH=$CMAKE_PREFIX_PATH:${LD_LIBRARY_PATH} +fi + +MAVEN_VERIFY_ARGS=() +if ! hasArg --run-java-tests; then + MAVEN_VERIFY_ARGS=("-DskipTests") +fi + +mvn clean verify "${MAVEN_VERIFY_ARGS[@]}" \ + && mvn jacoco:report \ + && mvn install:install-file -Dfile=./target/cuvs-lucene-$VERSION.jar -DgroupId=$GROUP_ID -DartifactId=cuvs-lucene -Dversion=$VERSION -Dpackaging=jar \ + && cp pom.xml ./target/ diff --git a/java/cuvs-lucene/examples/README.md b/java/cuvs-lucene/examples/README.md new file mode 100644 index 0000000000..f4a18caab4 --- /dev/null +++ b/java/cuvs-lucene/examples/README.md @@ -0,0 +1,78 @@ +# Examples + +This maven project contains basic examples that showcase how `cuvs-lucene` can be used. + +You can build the prerequisites either inside a RAPIDS conda container (the quickest way to try +things out, since it provides a prebuilt `libcuvs` without compiling the C/C++ libraries) or +entirely from source. + +## Option A — RAPIDS container (no local libcuvs build) + +### Prerequisites + +- [Docker](https://www.docker.com/) +- [Nvidia Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) +- A machine with an Nvidia GPU + +### Steps + +Launch the container **from the repository root** (the directory containing `build.sh` and +`dependencies.yaml`) — not from this `examples/` directory — so the whole repo is mounted and +becomes the container's working directory: + +```sh +docker run --rm --gpus all --pull=always --volume $PWD:$PWD --workdir $PWD -it rapidsai/ci-conda:26.10-cuda13.3.0-ubuntu24.04-py3.13 +``` + +Inside the container you are now at the repository root. Create a conda environment with `libcuvs` +and the Java toolchain (this reads the repo's `dependencies.yaml`, so it must be run from the +repository root), then build `cuvs-java` and `cuvs-lucene` against it: + +```sh +rapids-dependency-file-generator --output conda --file-key java \ + --matrix "cuda=13.3;arch=$(arch)" | tee /tmp/java_env.yaml +rapids-mamba-retry env create --yes -f /tmp/java_env.yaml -n java +conda activate java +export CMAKE_PREFIX_PATH=$CONDA_PREFIX +export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH +./build.sh java lucene +``` + +## Option B — build everything from source + +### Prerequisites + +- A local [CUDA toolkit](https://developer.nvidia.com/cuda-toolkit-archive) and a machine with an Nvidia GPU +- [Maven 3.9.6+](https://maven.apache.org/download.cgi) and [JDK 22](https://jdk.java.net/archive/) + +### Steps + +From the repository root, build `libcuvs`, the `cuvs-java` bindings, and `cuvs-lucene` (see the +[top-level Java README](../../README.md)), then point `LD_LIBRARY_PATH` at the in-tree libraries: + +```sh +./build.sh libcuvs java lucene +export LD_LIBRARY_PATH=$PWD/cpp/build:$LD_LIBRARY_PATH +``` + +## Build and run the examples + +With either option done (and its `LD_LIBRARY_PATH`/conda env active), build the examples from this +directory: + +```sh +cd java/cuvs-lucene/examples +mvn clean install +``` + +To run the Accelerated HNSW example do: + +```sh +java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.AcceleratedHnswExample +``` + +To run the Index and Search on GPU example do: + +```sh +java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample +``` diff --git a/java/cuvs-lucene/examples/pom.xml b/java/cuvs-lucene/examples/pom.xml new file mode 100644 index 0000000000..82cda4981e --- /dev/null +++ b/java/cuvs-lucene/examples/pom.xml @@ -0,0 +1,205 @@ + + + + 4.0.0 + + com.nvidia.cuvs.lucene.examples + examples + 26.10.0 + examples + + + 22 + 22 + UTF-8 + UTF-8 + + + + + + gcs-maven-central-mirror + GCS Maven Central mirror + https://maven-central.storage-download.googleapis.com/maven2/ + + true + + + false + + + + + central + Maven Central + https://repo.maven.apache.org/maven2 + + true + + + false + + + + + + + + gcs-maven-central-mirror + GCS Maven Central mirror + https://maven-central.storage-download.googleapis.com/maven2/ + + true + + + false + + + + + central + Maven Plugin Repository + https://repo.maven.apache.org/maven2 + + true + + + false + + + + + + + org.apache.lucene + lucene-core + 10.2.0 + + + org.apache.lucene + lucene-codecs + 10.2.0 + test + + + org.apache.lucene + lucene-backward-codecs + 10.2.0 + + + commons-io + commons-io + 2.18.0 + + + com.nvidia.cuvs.lucene + cuvs-lucene + 26.10.0 + + + + + + + com.diffplug.spotless + spotless-maven-plugin + 2.44.5 + + + + apply + + validate + + + + + + src/**/*.java + + + 1.27.0 + + true + false + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.6.0 + + + src/main/assembly/jar-with-merged-services.xml + + + + true + + + true + true + + + + + + make-assembly + package + + single + + + + + + + + + maven-clean-plugin + 3.4.0 + + + maven-resources-plugin + 3.3.1 + + + maven-compiler-plugin + 3.13.0 + + + maven-surefire-plugin + 3.3.0 + + + maven-jar-plugin + 3.4.2 + + + maven-install-plugin + 3.1.2 + + + maven-deploy-plugin + 3.1.2 + + + maven-site-plugin + 3.12.1 + + + maven-project-info-reports-plugin + 3.6.1 + + + + + diff --git a/java/cuvs-lucene/examples/src/main/assembly/jar-with-merged-services.xml b/java/cuvs-lucene/examples/src/main/assembly/jar-with-merged-services.xml new file mode 100644 index 0000000000..dc6706188d --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/assembly/jar-with-merged-services.xml @@ -0,0 +1,22 @@ + + jar-with-merged-services + + jar + + false + + + / + true + true + runtime + + + + + metaInf-services + + + diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java new file mode 100644 index 0000000000..d9fbffe9f0 --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java @@ -0,0 +1,132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static com.nvidia.cuvs.lucene.examples.Utils.generateDataset; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +public class AcceleratedHnswExample { + + private static Logger log = Logger.getLogger(AcceleratedHnswExample.class.getName()); + + private static Random random; + private static Path indexDirPath; + + public static void main(String[] args) throws Exception { + + AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.log(Level.FINE, "Successfully opened index"); + + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + log.log( + Level.FINE, + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + } + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + } + } + + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java new file mode 100644 index 0000000000..9d32032cc7 --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java @@ -0,0 +1,134 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static com.nvidia.cuvs.lucene.examples.Utils.generateDataset; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; +import com.nvidia.cuvs.lucene.GPUKnnFloatVectorQuery; +import com.nvidia.cuvs.lucene.GPUSearchParams; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +public class IndexAndSearchonGPUExample { + + private static Logger log = Logger.getLogger(IndexAndSearchonGPUExample.class.getName()); + + private static Random random; + private static Path indexDirPath; + + public static void main(String[] args) throws Exception { + + GPUSearchParams params = new GPUSearchParams.Builder().build(); + Codec codec = new CuVS2510GPUSearchCodec(params); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.log(Level.FINE, "Successfully opened index"); + + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + log.log( + Level.FINE, + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + } + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK, null, topK, 1); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + } + } + + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java new file mode 100644 index 0000000000..5e27130f6a --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import java.util.Random; + +public class Utils { + + public static float[][] generateDataset(Random random, int size, int dimensions) { + float[][] dataset = new float[size][dimensions]; + for (int i = 0; i < size; i++) { + for (int j = 0; j < dimensions; j++) { + dataset[i][j] = random.nextFloat() * 100; + } + } + return dataset; + } + + public static float[] generateRandomVector(int dimensions, Random random) { + float[] vector = new float[dimensions]; + for (int i = 0; i < dimensions; i++) { + vector[i] = random.nextFloat() * 100; + } + return vector; + } + + public static float[][] generateQueries(Random random, int dimensions, int numQueries) { + // Generate random query vectors + float[][] queries = new float[numQueries][dimensions]; + for (int i = 0; i < numQueries; i++) { + for (int j = 0; j < dimensions; j++) { + queries[i][j] = random.nextFloat() * 100; + } + } + return queries; + } +} diff --git a/java/cuvs-lucene/examples/src/main/resources/logging.properties b/java/cuvs-lucene/examples/src/main/resources/logging.properties new file mode 100644 index 0000000000..53f2c42727 --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/resources/logging.properties @@ -0,0 +1,9 @@ +# Set root level +.level=FINEST + +# Define handlers +handlers=java.util.logging.ConsoleHandler + +java.util.logging.ConsoleHandler.level=FINEST +java.util.logging.ConsoleHandler.formatter=java.util.logging.SimpleFormatter +java.util.logging.SimpleFormatter.format=[%1$tF %1$tT] %4$s %2$s %5$s%6$s%n diff --git a/java/cuvs-lucene/pom.xml b/java/cuvs-lucene/pom.xml new file mode 100644 index 0000000000..f76b6e6111 --- /dev/null +++ b/java/cuvs-lucene/pom.xml @@ -0,0 +1,249 @@ + + + + 4.0.0 + com.nvidia.cuvs.lucene + cuvs-lucene + + 26.10.0 + cuvs-lucene + jar + + + This project provides pluggable codecs for Apache Lucene to enable GPU-based accelerated vector search. + + https://rapids.ai + + + + Apache License, Version 2.0 + https://www.apache.org/licenses/LICENSE-2.0.txt + repo + A business-friendly OSS license + + + + scm:git:https://github.com/rapidsai/cuvs.git + scm:git:git@github.com:rapidsai/cuvs.git + HEAD + https://github.com/rapidsai/cuvs + + + + 22 + 22 + UTF-8 + UTF-8 + + + + + + gcs-maven-central-mirror + GCS Maven Central mirror + https://maven-central.storage-download.googleapis.com/maven2/ + + true + + + false + + + + + central + Maven Central + https://repo.maven.apache.org/maven2 + + true + + + false + + + + + + + + gcs-maven-central-mirror + GCS Maven Central mirror + https://maven-central.storage-download.googleapis.com/maven2/ + + true + + + false + + + + + central + Maven Plugin Repository + https://repo.maven.apache.org/maven2 + + true + + + false + + + + + + + org.apache.lucene + lucene-core + 10.2.0 + + + org.apache.lucene + lucene-codecs + 10.2.0 + test + + + org.apache.lucene + lucene-backward-codecs + 10.2.0 + + + org.apache.lucene + lucene-test-framework + 10.2.0 + test + + + commons-io + commons-io + 2.18.0 + test + + + com.nvidia.cuvs + cuvs-java + 26.10.0 + + + + + + + com.diffplug.spotless + spotless-maven-plugin + 2.44.5 + + + + apply + + validate + + + + + + src/**/*.java + + + 1.27.0 + + true + false + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.6.0 + + + jar-with-dependencies + + + + + make-assembly + package + + single + + + + + + org.apache.maven.plugins + maven-surefire-plugin + 3.5.4 + + plain + false + false + false + 0 + + ${project.basedir}/src/main/resources/logging.properties + + + + + org.apache.maven.plugins + maven-source-plugin + 3.3.1 + + + attach-sources + + jar-no-fork + + + + + + org.apache.maven.plugins + maven-javadoc-plugin + 3.12.0 + + + attach-javadocs + + jar + + + all,-missing + ${project.build.directory}/javadocs + + + + + + org.jacoco + jacoco-maven-plugin + 0.8.14 + + + + prepare-agent + + + + report + prepare-package + + report + + + + + + + diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java new file mode 100644 index 0000000000..aac314153d --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -0,0 +1,589 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSIvfPqParams; +import java.util.Objects; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.function.Supplier; + +public class AcceleratedHNSWParams { + + public static enum Strategy { + /* + * This strategy allows for automatic selection of the underlying CAGRA build algorithm. + * With this strategy we use NN_DESCENT for dataset less than 5M vectors, else we use IVF_PQ. + * Indexing parameters, especially for IVF_PQ, are heuristically identified automatically. + * + * This is the default and the recommended strategy. + */ + HEURISTIC, + /* + * This is an option when the end-user would want to use custom parameter values. + * This strategy should only be used under expert guidance. + */ + CUSTOM + } + + /* + * TODO: Update boundaries for all parameters when a consensus is reached. + * Issue: https://github.com/rapidsai/cuvs-lucene/issues/99 + */ + public static final int MIN_WRITER_THREADS = 1; + public static final int MAX_WRITER_THREADS = 512; + public static final int MIN_INT_GRAPH_DEG = 2; + public static final int MAX_INT_GRAPH_DEG = 512; + public static final int MIN_GRAPH_DEG = 1; + public static final int MAX_GRAPH_DEG = 512; + public static final int MIN_HNSW_LAYERS = 1; + public static final int MAX_HNSW_LAYERS = 99; + public static final int MIN_MAX_CONN = 1; + public static final int MAX_MAX_CONN = 512; + public static final int MIN_BEAM_WIDTH = 1; + public static final int MAX_BEAM_WIDTH = 512; + public static final int MIN_NUM_MERGE_WORKERS = 1; + public static final int MAX_NUM_MERGE_WORKERS = 512; + public static final int MIN_NN_DESCENT_NUM_ITERATIONS = 1; + public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; + + public static final int DEFAULT_WRITER_THREADS = 1; + public static final int DEFAULT_INT_GRAPH_DEGREE = 128; + public static final int DEFAULT_GRAPH_DEGREE = 64; + public static final int DEFAULT_HNSW_LAYERS = 1; + public static final int DEFAULT_MAX_CONN = 32; + public static final int DEFAULT_BEAM_WIDTH = 32; + public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = + CagraGraphBuildAlgo.NN_DESCENT; + public static final int DEFAULT_NUM_MERGE_WORKERS = 1; + public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; + public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; + public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; + + public static final Supplier DEFAULT_IVF_PQ_PARAMS = + () -> { + return new CuVSIvfPqParams.Builder().build(); + }; + + public static final Supplier DEFAULT_MERGE_EXE_SRVC = + () -> { + return Executors.newFixedThreadPool(DEFAULT_NUM_MERGE_WORKERS); + }; + + private final int writerThreads; + private final int intermediateGraphDegree; + private final int graphdegree; + private final int hnswLayers; + private final int maxConn; + private final int beamWidth; + private final CagraGraphBuildAlgo cagraGraphBuildAlgo; + private final CuVSIvfPqParams cuVSIvfPqParams; + private final int numMergeWorkers; + private final ExecutorService mergeExec; + private final Strategy strategy; + private final CuvsDistanceType cuvsDistanceType; + private final int nnDescentNumIterations; + + /** + * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. + * + * @param writerThreads Number of cuVS writer threads to use. + * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. + * @param graphdegree The graph degree to use while building the CAGRA index. + * @param indexType The type of index to build - CAGRA, BRUTEFORCE, or both. + * @param hnswLayers The number of HNSW layers to build in the HNSW index. + * @param maxConn The max connection parameter used when building HNSW index with the fallback mechanism. + * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism. + * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. + * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. + * @param numMergeWorkers The number of merge workers to use with the fallback mechanism. + * @param mergeExec The instance of {@link ExecutorService} to use with the fallback mechanism. + * @param strategy either HEURISTIC [Default] that automatically chooses build algorithm and its parameters based on data set size or CUSTOM that uses the parameters passed though this class. + * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. + * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. + */ + private AcceleratedHNSWParams( + int writerThreads, + int intermediateGraphDegree, + int graphdegree, + int hnswLayers, + int maxConn, + int beamWidth, + CagraGraphBuildAlgo cagraGraphBuildAlgo, + CuVSIvfPqParams cuVSIvfPqParams, + int numMergeWorkers, + ExecutorService mergeExec, + Strategy strategy, + CuvsDistanceType cuvsDistanceType, + int nnDescentNumIterations) { + super(); + this.writerThreads = writerThreads; + this.intermediateGraphDegree = intermediateGraphDegree; + this.graphdegree = graphdegree; + this.hnswLayers = hnswLayers; + this.maxConn = maxConn; + this.beamWidth = beamWidth; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + this.cuVSIvfPqParams = cuVSIvfPqParams; + this.numMergeWorkers = numMergeWorkers; + this.mergeExec = mergeExec; + this.strategy = strategy; + this.cuvsDistanceType = cuvsDistanceType; + this.nnDescentNumIterations = nnDescentNumIterations; + } + + /** + * Get the cuVS writer threads parameter + * + * @return cuVS writer threads parameter + */ + public int getWriterThreads() { + return writerThreads; + } + + /** + * Get the intermediate graph degree + * + * @return the graph degree parameter + */ + public int getIntermediateGraphDegree() { + return intermediateGraphDegree; + } + + /** + * Get the graph degree + * + * @return the graph degree parameter + */ + public int getGraphdegree() { + return graphdegree; + } + + /** + * Get the number of HNSW layers + * + * @return the number of HNSW layers + */ + public int getHnswLayers() { + return hnswLayers; + } + + /** + * Get the max connection parameter + * + * @return the max connection parameter + */ + public int getMaxConn() { + return maxConn; + } + + /** + * Get the beam width parameter + * + * @return the beam width parameter + */ + public int getBeamWidth() { + return beamWidth; + } + + /** + * Get the CAGRA graph build algorithm + * + * @return the CAGRA graph build algorithm + */ + public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { + return cagraGraphBuildAlgo; + } + + /** + * Get the instance of {@link CuVSIvfPqParams} + * + * @return the instance of {@link CuVSIvfPqParams} + */ + public CuVSIvfPqParams getCuVSIvfPqParams() { + return cuVSIvfPqParams; + } + + /** + * Get the number of merge workers set to be used in the fallback mechanism + * + * @return the number of merge workers + */ + public int getNumMergeWorkers() { + return numMergeWorkers; + } + + /** + * Get the instance of the {@link ExecutorService} to be used in the fallback mechanism + * + * @return the instance of the {@link ExecutorService} + */ + public ExecutorService getMergeExec() { + return mergeExec; + } + + /** + * Get the chosen strategy: + * + * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set + * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally + * + * @return get the chosen {@link Strategy} + */ + public Strategy getStrategy() { + return strategy; + } + + /** + * Get the cuvs distance type + * + * @return the distance type + */ + public CuvsDistanceType getCuvsDistanceType() { + return cuvsDistanceType; + } + + /** + * get the number of Iterations to run if building with NN_DESCENT + * + * @return the number of iterations for NN_DESCENT + */ + public int getNNDescentNumIterations() { + return nnDescentNumIterations; + } + + @Override + public String toString() { + return "AcceleratedHNSWParams [writerThreads=" + + writerThreads + + ", intermediateGraphDegree=" + + intermediateGraphDegree + + ", graphdegree=" + + graphdegree + + ", hnswLayers=" + + hnswLayers + + ", maxConn=" + + maxConn + + ", beamWidth=" + + beamWidth + + ", cagraGraphBuildAlgo=" + + cagraGraphBuildAlgo + + ", cuVSIvfPqParams=" + + cuVSIvfPqParams + + ", numMergeWorkers=" + + numMergeWorkers + + ", mergeExec=" + + mergeExec + + ", strategy=" + + strategy + + ", cuvsDistanceType=" + + cuvsDistanceType + + ", nnDescentNumIterations=" + + nnDescentNumIterations + + "]"; + } + + /** + * Builder class for creating an instance of {@link GPUSearchParams} + */ + public static class Builder { + + private int writerThreads = DEFAULT_WRITER_THREADS; + private int intermediateGraphDegree = DEFAULT_INT_GRAPH_DEGREE; + private int graphdegree = DEFAULT_GRAPH_DEGREE; + private int hnswLayers = DEFAULT_HNSW_LAYERS; + private int maxConn = DEFAULT_MAX_CONN; + private int beamWidth = DEFAULT_BEAM_WIDTH; + private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO; + private int numMergeWorkers = DEFAULT_NUM_MERGE_WORKERS; + private CuVSIvfPqParams cuVSIvfPqParams = null; + private ExecutorService mergeExec = null; + private Strategy strategy = DEFAULT_STRATEGY; + private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE; + private int nnDescentNumIterations = DEFAULT_NN_DESCENT_NUM_ITERATIONS; + + /** + * Set the number of cuVS writer threads while building the index + * Valid range - Minimum: {@value MIN_WRITER_THREADS}, Maximum: {@value MAX_WRITER_THREADS} + * Default value - {@value DEFAULT_WRITER_THREADS} + * + * @param writerThreads + * @return instance of {@link Builder} + */ + public Builder withWriterThreads(int writerThreads) { + this.writerThreads = writerThreads; + return this; + } + + /** + * Set the intermediate graph degree to use while building CAGRA index + * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} + * Default value - {@value DEFAULT_INT_GRAPH_DEGREE} + * + * @param intermediateGraphDegree + * @return instance of {@link Builder} + */ + public Builder withIntermediateGraphDegree(int intermediateGraphDegree) { + this.intermediateGraphDegree = intermediateGraphDegree; + return this; + } + + /** + * Set the graph degree to use while building CAGRA index + * Valid range - Minimum: {@value MIN_GRAPH_DEG}, Maximum: {@value MAX_GRAPH_DEG} + * Default value - {@value DEFAULT_GRAPH_DEGREE} + * + * @param graphDegree + * @return instance of {@link Builder} + */ + public Builder withGraphDegree(int graphDegree) { + this.graphdegree = graphDegree; + return this; + } + + /** + * Set the number of HNSW layers to construct while building the HNSW index + * Valid range - Minimum: {@value MIN_HNSW_LAYERS}, Maximum: {@value MAX_HNSW_LAYERS} + * Default value - {@value DEFAULT_HNSW_LAYERS} + * + * @param hnswLayers the number of HNSW layers + * @return instance of {@link Builder} + */ + public Builder withHNSWLayer(int hnswLayers) { + this.hnswLayers = hnswLayers; + return this; + } + + /** + * Set the max connections parameter while building HNSW index with fallback mechanism + * Valid range - Minimum: {@value MIN_MAX_CONN}, Maximum: {@value MAX_MAX_CONN} + * Default value - {@value DEFAULT_MAX_CONN} + * + * @param maxConn the max connections parameter + * @return instance of {@link Builder} + */ + public Builder withMaxConn(int maxConn) { + this.maxConn = maxConn; + return this; + } + + /** + * Set the beam width parameter while building HNSW index with fallback mechanism + * Valid range - Minimum: {@value MIN_BEAM_WIDTH}, Maximum: {@value MAX_BEAM_WIDTH} + * Default value - {@value DEFAULT_BEAM_WIDTH} + * + * @param beamWidth the beam width parameter + * @return instance of {@link Builder} + */ + public Builder withBeamWidth(int beamWidth) { + this.beamWidth = beamWidth; + return this; + } + + /** + * Set the CAGRA graph build algorithm to use + * Default value - NN_DESCENT + * + * @param cagraGraphBuildAlgo + * @return instance of {@link Builder} + */ + public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo) { + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + return this; + } + + /** + * Set the instance of {@link CuVSIvfPqParams} + * + * @param cuVSIvfPqParams + * @return instance of {@link Builder} + */ + public Builder withCuVSIvfPqParams(CuVSIvfPqParams cuVSIvfPqParams) { + this.cuVSIvfPqParams = cuVSIvfPqParams; + return this; + } + + /** + * Set the number of merge workers to be used with the fallback mechanism + * Default value - {@value DEFAULT_NUM_MERGE_WORKERS} + * + * @param numMergeWorkers number of merge workers to set + * @return instance of {@link Builder} + */ + public Builder withNumMergeWorkers(int numMergeWorkers) { + this.numMergeWorkers = numMergeWorkers; + return this; + } + + /** + * Set the merge executor service to be used in the fallback mechanism + * Default value an instance with one thread + * + * @param mergeExec an instance of {@link ExecutorService} + * @return instance of {@link Builder} + */ + public Builder withMergeExecutorService(ExecutorService mergeExec) { + this.mergeExec = mergeExec; + return this; + } + + /** + * Set the chosen strategy: + * + * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set + * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally + * + * Valid options - HEURISTIC, CUSTOM + * Default value - HEURISTIC + * + * @param strategy, the strategy to choose + * @return instance of {@link Builder} + */ + public Builder withStrategy(Strategy strategy) { + this.strategy = strategy; + return this; + } + + /** + * Set the CuvsDistanceType + * + * @param cuvsDistanceType the CuvsDistanceType to set + * @return instance of {@link Builder} + */ + public Builder withCuvsDistanceType(CuvsDistanceType cuvsDistanceType) { + this.cuvsDistanceType = cuvsDistanceType; + return this; + } + + /** + * Set the number of Iterations to run if building with NN_DESCENT + * + * Valid range - Minimum: {@value MIN_NN_DESCENT_NUM_ITERATIONS}, Maximum: {@value MAX_NN_DESCENT_NUM_ITERATIONS} + * Default value - {@value DEFAULT_NN_DESCENT_NUM_ITERATIONS} + * + * @param nnDescentNumIterations number of merge workers to set + * @return instance of {@link Builder} + */ + public Builder withNNDescentNumIterations(int nnDescentNumIterations) { + this.nnDescentNumIterations = nnDescentNumIterations; + return this; + } + + /** + * Validates the input parameters. + * + * @throws IllegalArgumentException + */ + private void validate() throws IllegalArgumentException { + if (writerThreads < MIN_WRITER_THREADS || writerThreads > MAX_WRITER_THREADS) { + throw new IllegalArgumentException( + "writerThreads not in valid range. Valid range: [" + + MIN_WRITER_THREADS + + ", " + + MAX_WRITER_THREADS + + "]"); + } + if (intermediateGraphDegree < MIN_INT_GRAPH_DEG + || intermediateGraphDegree > MAX_INT_GRAPH_DEG) { + throw new IllegalArgumentException( + "intermediateGraphDegree not in valid range. Valid range: [" + + MIN_INT_GRAPH_DEG + + ", " + + MAX_INT_GRAPH_DEG + + "]"); + } + if (graphdegree < MIN_GRAPH_DEG || graphdegree > MAX_GRAPH_DEG) { + throw new IllegalArgumentException( + "graphdegree not in valid range. Valid range: [" + + MIN_GRAPH_DEG + + ", " + + MAX_GRAPH_DEG + + "]"); + } + if (hnswLayers < MIN_HNSW_LAYERS || hnswLayers > MAX_HNSW_LAYERS) { + throw new IllegalArgumentException( + "hnswLayers not in valid range. Valid range: [" + + MIN_HNSW_LAYERS + + ", " + + MAX_HNSW_LAYERS + + "]"); + } + if (maxConn < MIN_MAX_CONN || maxConn > MAX_MAX_CONN) { + throw new IllegalArgumentException( + "maxConn not in valid range. Valid range: [" + + MIN_MAX_CONN + + ", " + + MAX_MAX_CONN + + "]"); + } + if (beamWidth < MIN_BEAM_WIDTH || beamWidth > MAX_BEAM_WIDTH) { + throw new IllegalArgumentException( + "beamWidth not in valid range. Valid range: [" + + MIN_BEAM_WIDTH + + ", " + + MAX_BEAM_WIDTH + + "]"); + } + if (Objects.isNull(cagraGraphBuildAlgo)) { + throw new IllegalArgumentException("cagraGraphBuildAlgo cannot be null."); + } + if (numMergeWorkers < MIN_NUM_MERGE_WORKERS || numMergeWorkers > MAX_NUM_MERGE_WORKERS) { + throw new IllegalArgumentException( + "numMergeWorkers not in valid range. Valid range: [" + + MIN_NUM_MERGE_WORKERS + + ", " + + MAX_NUM_MERGE_WORKERS + + "]"); + } + if (Objects.isNull(strategy)) { + throw new IllegalArgumentException("strategy cannot be null."); + } + if (Objects.isNull(cuvsDistanceType)) { + throw new IllegalArgumentException("cuvsDistanceType cannot be null."); + } + if (nnDescentNumIterations < MIN_NN_DESCENT_NUM_ITERATIONS + || nnDescentNumIterations > MAX_NN_DESCENT_NUM_ITERATIONS) { + throw new IllegalArgumentException( + "nnDescentNumIterations not in valid range. Valid range: [" + + MIN_NN_DESCENT_NUM_ITERATIONS + + ", " + + MAX_NN_DESCENT_NUM_ITERATIONS + + "]"); + } + } + + /** + * Create an instance of {@link AcceleratedHNSWParams} + * + * @return instance of {@link AcceleratedHNSWParams} + */ + public AcceleratedHNSWParams build() { + if (Objects.isNull(cuVSIvfPqParams)) { + cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS.get(); + } + if (Objects.isNull(mergeExec)) { + mergeExec = DEFAULT_MERGE_EXE_SRVC.get(); + } + validate(); + return new AcceleratedHNSWParams( + writerThreads, + intermediateGraphDegree, + graphdegree, + hnswLayers, + maxConn, + beamWidth, + cagraGraphBuildAlgo, + cuVSIvfPqParams, + numMergeWorkers, + mergeExec, + strategy, + cuvsDistanceType, + nnDescentNumIterations); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java new file mode 100644 index 0000000000..18d2645335 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -0,0 +1,489 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.createByteMatrixFromArray; + +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.RowView; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.SortedSet; +import java.util.TreeSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.apache.lucene.util.packed.DirectMonotonicWriter; + +public class AcceleratedHNSWUtils { + + public enum QuantizationType { + BINARY, + SCALAR, + NONE + } + + private static final LuceneProvider LUCENE_PROVIDER; + private static final List VECTOR_SIMILARITY_FUNCTIONS; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Creates a dummy HNSW graph for a single vector. + * The graph will have 1 level with 1 node and no neighbors. + */ + public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimensions) + throws Throwable { + // Create adjacency list for single node with no neighbors + int[][] singleNodeAdjacency = new int[][] {{-1}}; // -1 indicates no neighbors + + // Create CuVSMatrix from the adjacency list + CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency); + + // Create layer data for single-level graph + List layerNodes = new ArrayList<>(); + List layerAdjacencies = new ArrayList<>(); + + // Layer 0: contains all nodes (just the single node) + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list + layerAdjacencies.add(adjacencyMatrix); + + // Create the single-layer graph + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } + + /** + * Creates a multi-layer HNSW graph with dynamic number of layers. + * M = cagraGraphDegree/2 + * Each layer contains 1/M nodes from the previous layer + * Creates layers until the highest layer has ≤ M nodes + */ + public static GPUBuiltHnswGraph createMultiLayerHnswGraph( + FieldInfo fieldInfo, + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + int graphDegree, + CagraIndexParams params, + QuantizationType quantization) + throws Throwable { + + // Calculate M as cagraGraphDegree/2 + int M = graphDegree / 2; + + // Store all layers data + List layerNodes = new ArrayList<>(); + List layerAdjacencies = new ArrayList<>(); + + // Layer 0: Use full CAGRA adjacency list + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list + layerAdjacencies.add(adjacencyListMatrix); + + int currentLayerSize = size; + int layerIndex = 1; + Random random = new Random(); + + while (layerIndex < hnswLayers && currentLayerSize > 1) { + // Calculate size for next layer (1/M of current layer) + int nextLayerSize = Math.max(2, currentLayerSize / M); + // Select nodes for this layer + SortedSet selectedNodesSet = new TreeSet<>(); + + if (layerIndex == 1) { + // Select from all nodes (Layer 0) + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(random.nextInt(size)); + } + } else { + // Select from previous layer nodes + int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); + while (selectedNodesSet.size() < nextLayerSize) { + int idx = random.nextInt(prevLayerNodes.length); + selectedNodesSet.add(prevLayerNodes[idx]); + } + } + + // Convert to sorted array + int[] selectedNodes = + selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); + + layerNodes.add(selectedNodes); + + if (quantization == QuantizationType.NONE) { + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); + } + + // Build CAGRA graph for this layer + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + + } else { + + // Extract vectors for selected nodes + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + } + + // Build CAGRA graph for this layer + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); + } + + // Update for next iteration + currentLayerSize = nextLayerSize; + layerIndex++; + + // Use different seed for each layer + random = new Random(new Random().nextLong()); + } + + // Create the multi-layer graph with all layers + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } + + /** + * Builds a CAGRA graph for a subset of binary quantized vectors + */ + private static CuVSMatrix buildCagraGraphForSubset( + Object vectors, + int[] selectedNodes, + int bytesPerVector, + CagraIndexParams params, + int dimensions, + QuantizationType quantization) + throws Throwable { + + CuVSMatrix subsetDataset; + + if (quantization == QuantizationType.BINARY) { + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); + } else if (quantization == QuantizationType.SCALAR) { + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); + } else { + subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); + } + + // Build CAGRA index for the subset + CagraIndex subsetIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(subsetDataset) + .withIndexParams(params) + .build(); + + // Get adjacency list from subset CAGRA index + CuVSMatrix cagraGraph = subsetIndex.getGraph(); + + long numNodes = cagraGraph.size(); + long degree = cagraGraph.columns(); + + // Create a re-mapped adjacency list + int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; + + for (int i = 0; i < numNodes; i++) { + RowView rv = cagraGraph.getRow(i); + for (int j = 0; j < degree && j < rv.size(); j++) { + int subsetIndex1 = rv.getAsInt(j); + // Map subset index to original node ID + if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { + remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; + } else { + // Invalid index, use self-reference + remappedAdjacency[i][j] = selectedNodes[i]; + } + } + } + + subsetIndex.close(); + return CuVSMatrix.ofArray(remappedAdjacency); + } + + /** + * Returns a 2D array of offsets (information written while writing the meta info) + * + * @param graph instance of GPUBuiltHnswGraph + * @param vectorIndex instance of IndexOutput + * @return a 2D array of offsets + * @throws IOException I/O Exceptions + */ + public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) + throws IOException { + // write vectors' neighbors on each level into the vectorIndex file + int countOnLevel0 = graph.size(); + int[][] offsets = new int[graph.numLevels()][]; + int[] scratch = new int[graph.maxConn() * 2]; + for (int level = 0; level < graph.numLevels(); level++) { + int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); + offsets[level] = new int[sortedNodes.length]; + int nodeOffsetId = 0; + + for (int node : sortedNodes) { + // Get node neighbors + NeighborArray neighbors = graph.getNeighbors(level, node); + // Get the size of the neighbor array + int size = neighbors.size(); + // Write size in VInt as the neighbors list is typically small + long offsetStart = vectorIndex.getFilePointer(); + // Get neighbors + int[] nnodes = neighbors.nodes(); + // Sort them + Arrays.sort(nnodes, 0, size); + // Now that we have sorted, do delta encoding to minimize the required bits to store the + // information + int actualSize = 0; + if (size > 0) { + scratch[0] = nnodes[0]; + actualSize = 1; + } + // De-duplication + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + // Sorting step helps here + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + } + // Write the size after duplicates are removed + vectorIndex.writeVInt(actualSize); + // Write de-duplicated neighbors + for (int i = 0; i < actualSize; i++) { + vectorIndex.writeVInt(scratch[i]); + } + offsets[level][nodeOffsetId++] = + Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); + } + } + // Return offsets (information written while writing the meta info) + return offsets; + } + + /** + * Writes the meta information for the index. + * + * @param vectorIndex instance of IndexOutput + * @param meta instance of IndexOutput + * @param field instance of FieldInfo + * @param vectorIndexOffset vector index offset + * @param vectorIndexLength vector index length + * @param count the count of vectors + * @param graph instance of HnswGraph + * @param graphLevelNodeOffsets graph level node offsets + * @throws IOException I/O Exceptions + */ + public static void writeMeta( + IndexOutput vectorIndex, + IndexOutput meta, + FieldInfo field, + long vectorIndexOffset, + long vectorIndexLength, + int count, + HnswGraph graph, + int[][] graphLevelNodeOffsets, + int graphDegree) + throws IOException { + + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); + meta.writeVLong(vectorIndexOffset); + meta.writeVLong(vectorIndexLength); + meta.writeVInt(field.getVectorDimension()); + meta.writeInt(count); + meta.writeVInt(graphDegree / 2); // M = cagraGraphDegree/2 + + // write graph nodes on each level + if (graph == null) { + meta.writeVInt(0); + } else { + meta.writeVInt(graph.numLevels()); + long valueCount = 0; + for (int level = 0; level < graph.numLevels(); level++) { + NodesIterator nodesOnLevel = graph.getNodesOnLevel(level); + valueCount += nodesOnLevel.size(); + if (level > 0) { + int[] nol = new int[nodesOnLevel.size()]; + int numberConsumed = nodesOnLevel.consume(nol); + Arrays.sort(nol); + assert numberConsumed == nodesOnLevel.size(); + meta.writeVInt(nol.length); // number of nodes on a level + for (int i = nodesOnLevel.size() - 1; i > 0; --i) { + nol[i] -= nol[i - 1]; + } + for (int n : nol) { + meta.writeVInt(n); + } + } else { + assert nodesOnLevel.size() == count : "Level 0 expects to have all nodes"; + } + } + + long start = vectorIndex.getFilePointer(); + meta.writeLong(start); + meta.writeVInt(16); // DIRECT_MONOTONIC_BLOCK_SHIFT); + + final DirectMonotonicWriter memoryOffsetsWriter = + DirectMonotonicWriter.getInstance(meta, vectorIndex, valueCount, 16); + long cumulativeOffsetSum = 0; + for (int[] levelOffsets : graphLevelNodeOffsets) { + for (int v : levelOffsets) { + memoryOffsetsWriter.add(cumulativeOffsetSum); + cumulativeOffsetSum += v; + } + } + + memoryOffsetsWriter.finish(); + meta.writeLong(vectorIndex.getFilePointer() - start); + } + } + + public static int distFuncToOrd(VectorSimilarityFunction func) { + for (int i = 0; i < VECTOR_SIMILARITY_FUNCTIONS.size(); i++) { + if (VECTOR_SIMILARITY_FUNCTIONS.get(i).equals(func)) { + return (byte) i; + } + } + throw new IllegalArgumentException("invalid distance function: " + func); + } + + /** + * A utility method to print info/debugging messages using InfoStream. + * + * @param msg the debugging message to print + */ + public static void printInfoStream(InfoStream infoStream, String component, String msg) { + if (infoStream.isEnabled(component)) { + infoStream.message(component, msg); + } + } + + /** + * Writes an empty meta information for the field. + * + * @param fieldInfo instance of FieldInfo + * @throws IOException I/O Exceptions + */ + public static void writeEmpty(FieldInfo fieldInfo, IndexOutput op) throws IOException { + writeMeta(null, op, fieldInfo, 0, 0, 0, null, null, 0); + } + + /** + * Quantizes FLOAT32 vectors to binary (1 bit per dimension, packed into bytes). + * Binary quantization: each dimension is compared to a centroid (mean of all values for that dimension). + * If value > centroid, bit = 1, else bit = 0. + * Bits are packed: 8 dimensions per byte. + * + * @param floatVectors A list of float vectors + * @return A list of byte binary representation for the input vectors + */ + public static List quantizeFloatVectorsToBinary(List floatVectors) { + if (floatVectors.isEmpty()) { + return new ArrayList<>(); + } + + int dimensions = floatVectors.get(0).length; + int numVectors = floatVectors.size(); + int bytesPerVector = (dimensions + 7) / 8; + + float[] centroids = new float[dimensions]; + for (float[] vector : floatVectors) { + for (int d = 0; d < dimensions; d++) { + centroids[d] += vector[d]; + } + } + for (int d = 0; d < dimensions; d++) { + centroids[d] /= numVectors; + } + + List quantizedVectors = new ArrayList<>(numVectors); + for (float[] vector : floatVectors) { + byte[] quantized = new byte[bytesPerVector]; + for (int d = 0; d < dimensions; d++) { + boolean bit = vector[d] > centroids[d]; + int byteIndex = d / 8; + int bitIndex = d % 8; + if (bit) { + quantized[byteIndex] |= (1 << bitIndex); + } + } + quantizedVectors.add(quantized); + } + + return quantizedVectors; + } + + /** + * Scalar quantization. + * + * @param floatVectors A list of float vectors + * @return A list of byte scalar representation for the input vectors + */ + public static List quantizeFloatVectorsToScalar(List floatVectors) { + if (floatVectors.isEmpty()) { + return new ArrayList<>(); + } + + int dimensions = floatVectors.get(0).length; + int numVectors = floatVectors.size(); + + float[] minPerDim = new float[dimensions]; + float[] maxPerDim = new float[dimensions]; + Arrays.fill(minPerDim, Float.MAX_VALUE); + Arrays.fill(maxPerDim, Float.MIN_VALUE); + + for (float[] vector : floatVectors) { + for (int d = 0; d < dimensions; d++) { + minPerDim[d] = Math.min(minPerDim[d], vector[d]); + maxPerDim[d] = Math.max(maxPerDim[d], vector[d]); + } + } + + List quantizedVectors = new ArrayList<>(numVectors); + for (float[] vector : floatVectors) { + byte[] quantized = new byte[dimensions]; + for (int d = 0; d < dimensions; d++) { + float range = maxPerDim[d] - minPerDim[d]; + if (range > 0) { + float normalized = (vector[d] - minPerDim[d]) / range; + int quantizedValue = Math.round(normalized * 127.0f) - 64; + quantized[d] = (byte) Math.max(-64, Math.min(63, quantizedValue)); + } else { + quantized[d] = 0; + } + } + quantizedVectors.add(quantized); + } + + return quantizedVectors; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BogusPropagation.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BogusPropagation.java new file mode 100644 index 0000000000..0a7c1d39ae --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BogusPropagation.java @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.spi.CuVSProvider; + +/** + * Bogus helper used to verify end-to-end build propagation from the cuVS C layer, through the + * cuvs-java bindings, to cuvs-lucene. It delegates to {@link CuVSProvider#bogusAddV1(int, int)}, which + * in turn calls the {@code cuvsBogusAddV1} C function. If any layer of the build fails to propagate, + * code that references this method will fail to compile or fail at runtime. + */ +public final class BogusPropagation { + + private BogusPropagation() {} + + /** + * Returns the sum of {@code a} and {@code b} by delegating to the bogus native function exposed + * through cuvs-java. + * + * @param a first addend + * @param b second addend + * @return the sum {@code a + b} + */ + public static int bogusAddV1(int a, int b) { + return CuVSProvider.provider().bogusAddV1(a, b); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java new file mode 100644 index 0000000000..20fed04bc9 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -0,0 +1,225 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CodebookGen; +import com.nvidia.cuvs.CagraIndexParams.CudaDataType; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSIvfPqIndexParams; +import com.nvidia.cuvs.CuVSIvfPqParams; +import com.nvidia.cuvs.CuVSIvfPqSearchParams; + +/** + * A centralized approach to producing instances of {@link CagraIndexParams} based on the chosen strategy + */ +public class CagraIndexParamsFactory { + + private static final int ALGO_SWITCH_THRESHOLD = 5_000_000; + + /** + * Translation of the internal logic found here: + * https://github.com/rapidsai/cuvs/blob/main/cpp/include/cuvs/neighbors/ivf_pq.hpp#L3385-L3428 + * + * Ideally we should hook into the internal API but this is currently replicated to avoid complications + * in other parts of code base. + */ + private static CuVSIvfPqParams getCuVSIvfPqParams(long rows, long dimension) { + + int pqDim; + int pqBits; + + if (dimension <= 32) { + pqDim = 16; + pqBits = 8; + } else { + pqBits = 4; + if (dimension <= 64) { + pqDim = 32; + } else if (dimension <= 128) { + pqDim = 64; + } else if (dimension <= 192) { + pqDim = 96; + } else { + pqDim = (int) roundUpSafe(dimension / 2, 128); + } + } + + int nLists = (int) Math.max(1, rows / 2000); + final int kmeansNIters = 10; + final double kMinPointsPerCluster = 32; + double minKmeansTrainsetPoints = kMinPointsPerCluster * nLists; + final double maxKmeansTrainsetFraction = 1.0; + double minKmeansTrainsetFraction = + Math.min(maxKmeansTrainsetFraction, minKmeansTrainsetPoints / rows); + double kmeansTrainsetFraction = + Math.clamp( + 1.0 / Math.sqrt(rows * 1e-5), minKmeansTrainsetFraction, maxKmeansTrainsetFraction); + final CodebookGen codebookKind = CodebookGen.PER_SUBSPACE; + int nProbes = (int) Math.round(Math.sqrt(nLists) / 20 + 4); + final int refinementRate = 1; + + CuVSIvfPqIndexParams cuVSIvfPqIndexParams = + new CuVSIvfPqIndexParams.Builder() + .withCodebookKind(codebookKind) + .withKmeansNIters(kmeansNIters) + .withKmeansTrainsetFraction(kmeansTrainsetFraction) + .withNLists(nLists) + .withPqBits(pqBits) + .withPqDim(pqDim) + .withAddDataOnBuild(true) + .withConservativeMemoryAllocation(true) + .build(); + + CuVSIvfPqSearchParams cuVSIvfPqSearchParams = + new CuVSIvfPqSearchParams.Builder() + .withLutDtype(CudaDataType.CUDA_R_16F) + .withInternalDistanceDtype(CudaDataType.CUDA_R_16F) + .withNProbes(nProbes) + .build(); + + CuVSIvfPqParams cuVSIvfPqParams = + new CuVSIvfPqParams.Builder() + .withCuVSIvfPqIndexParams(cuVSIvfPqIndexParams) + .withCuVSIvfPqSearchParams(cuVSIvfPqSearchParams) + .withRefinementRate(refinementRate) + .build(); + + return cuVSIvfPqParams; + } + + /* + * Rough translation from raft's internal utility found here: + * https://github.com/rapidsai/raft/blob/main/cpp/include/raft/util/integer_utils.hpp#L47-L56 + */ + private static long roundUpSafe(long numberToRound, long modulus) { + long remainder = numberToRound % modulus; + if (remainder == 0) { + return numberToRound; + } + long roundedUp = numberToRound - remainder + modulus; + return roundedUp; + } + + private static CagraIndexParams getNNDescentParams( + int graphDegree, + int intGraphDegree, + int writerThreads, + long nnDescentNumIterations, + CuvsDistanceType cuvsDistanceType) { + return new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(graphDegree) + .withIntermediateGraphDegree(intGraphDegree) + .withNNDescentNumIterations(nnDescentNumIterations) + .withNumWriterThreads(writerThreads) + .withMetric(cuvsDistanceType) + .build(); + } + + private static CagraIndexParams getIVFPQParams( + int graphDegree, + int intGraphDegree, + int writerThreads, + long rows, + long dimension, + CuvsDistanceType cuvsDistanceType) { + return new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.IVF_PQ) + .withCuVSIvfPqParams(getCuVSIvfPqParams(rows, dimension)) + .withNumWriterThreads(writerThreads) + .withIntermediateGraphDegree(intGraphDegree) + .withGraphDegree(graphDegree) + .withMetric(cuvsDistanceType) + .build(); + } + + /** + * Creates an instance of {@link CagraIndexParams} based on the chosen strategy in the {@link GPUSearchParams}. + * + * @param gPUSearchParams an instance of {@link GPUSearchParams} containing input params incoming via the build and search on the GPU API. + * @param rows number of vectors in the data set + * @param dimension the dimension of the vectors in the data set + * @return an instance of {@link CagraIndexParams} + */ + public static CagraIndexParams create( + GPUSearchParams gPUSearchParams, long rows, long dimension) { + if (gPUSearchParams.getStrategy().equals(GPUSearchParams.Strategy.HEURISTIC)) { + if (rows < ALGO_SWITCH_THRESHOLD) { + return getNNDescentParams( + gPUSearchParams.getGraphdegree(), + gPUSearchParams.getIntermediateGraphDegree(), + gPUSearchParams.getWriterThreads(), + gPUSearchParams.getnNDescentNumIterations(), + gPUSearchParams.getCuvsDistanceType()); + } else { + return getIVFPQParams( + gPUSearchParams.getGraphdegree(), + gPUSearchParams.getIntermediateGraphDegree(), + gPUSearchParams.getWriterThreads(), + rows, + dimension, + gPUSearchParams.getCuvsDistanceType()); + } + } else { + return new CagraIndexParams.Builder() + .withNumWriterThreads(gPUSearchParams.getWriterThreads()) + .withIntermediateGraphDegree(gPUSearchParams.getIntermediateGraphDegree()) + .withGraphDegree(gPUSearchParams.getGraphdegree()) + .withCagraGraphBuildAlgo(gPUSearchParams.getCagraGraphBuildAlgo()) + .withCuVSIvfPqParams(gPUSearchParams.getCuVSIvfPqParams()) + .withNNDescentNumIterations(gPUSearchParams.getnNDescentNumIterations()) + .build(); + } + } + + /* + * Ideally there should be just one create method instead of two. + * We should do that when both the input parameter classes can be unified in the future. + */ + + /** + * Creates an instance of {@link CagraIndexParams} based on the chosen strategy in the {@link AcceleratedHNSWParams}. + * + * @param acceleratedHNSWParams an instance of {@link AcceleratedHNSWParams} containing input params incoming via the build and search on the GPU API. + * @param rows number of vectors in the data set + * @param dimension the dimension of the vectors in the data set + * @return an instance of {@link CagraIndexParams} + */ + public static CagraIndexParams create( + AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { + if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { + if (rows + < ALGO_SWITCH_THRESHOLD) { // TODO: maybe consider making this threshold configurable from + // outside later. + return getNNDescentParams( + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getNNDescentNumIterations(), + acceleratedHNSWParams.getCuvsDistanceType()); + } else { + return getIVFPQParams( + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getWriterThreads(), + rows, + dimension, + acceleratedHNSWParams.getCuvsDistanceType()); + } + } else { + return new CagraIndexParams.Builder() + .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) + .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) + .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) + .withCagraGraphBuildAlgo(acceleratedHNSWParams.getCagraGraphBuildAlgo()) + .withCuVSIvfPqParams(acceleratedHNSWParams.getCuVSIvfPqParams()) + .withNNDescentNumIterations(acceleratedHNSWParams.getNNDescentNumIterations()) + .build(); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java new file mode 100644 index 0000000000..9eb5317b64 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.LibraryException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; + +/** + * cuVS based codec for GPU based vector search that enables both - indexing and search on the GPU. + * cuVS serialization formats are in experimental phase and hence backward compatibility cannot be guaranteed. + * + * @since 25.10 + */ +public class CuVS2510GPUSearchCodec extends FilterCodec { + + private static final Logger log = Logger.getLogger(CuVS2510GPUSearchCodec.class.getName()); + private static final String NAME = "CuVS2510GPUSearchCodec"; + private KnnVectorsFormat format; + + /** + * Default constructor for {@link CuVS2510GPUSearchCodec}. + * + * @throws Exception + */ + public CuVS2510GPUSearchCodec() throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + initializeFormat(new GPUSearchParams.Builder().build()); + } + + /** + * Initialize {@link CuVS2510GPUSearchCodec} with an instance of {@link GPUSearchParams} + * having default parameter values. + * + * @param name the name of the codec + * @param delegate the delegate codec + */ + public CuVS2510GPUSearchCodec(String name, Codec delegate) { + super(name, delegate); + initializeFormat(new GPUSearchParams.Builder().build()); + } + + /** + * Initialize the codec with an instance of {@link GPUSearchParams} having either default + * or overridden parameter values. + * + * @param params An instance of {@link GPUSearchParams} + * @throws Exception Exception raised when initializing the codec + */ + public CuVS2510GPUSearchCodec(GPUSearchParams params) throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + initializeFormat(params); + } + + /** + * Initialize the {@link CuVS2510GPUVectorsFormat} instance using {@link GPUSearchParams}. + * + * @param params an instance of {@link GPUSearchParams} + */ + private void initializeFormat(GPUSearchParams params) { + try { + format = new CuVS2510GPUVectorsFormat(params); + setKnnFormat(format); + } catch (LibraryException ex) { + log.log( + Level.SEVERE, + "Couldn't load native library, possible classloader issue. " + ex.getMessage()); + } + } + + /** + * Get the configured {@link KnnVectorsFormat}. + * + * @return the instance of the {@link KnnVectorsFormat} + */ + @Override + public KnnVectorsFormat knnVectorsFormat() { + return format; + } + + /** + * Set the {@link KnnVectorsFormat}. + * + * @param format the {@link KnnVectorsFormat} to set + */ + public void setKnnFormat(KnnVectorsFormat format) { + this.format = format; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java new file mode 100644 index 0000000000..a38a28289f --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.assertIsSupported; + +import com.nvidia.cuvs.LibraryException; +import java.io.IOException; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** + * Extends upon the KnnVectorsFormat - Encodes/decodes per-document vector and any associated indexing structures required to support + * GPU-based accelerated nearest-neighbor search. + * + * @since 25.10 + */ +public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { + + private static final int MAX_DIMENSIONS = 4096; + private static final LuceneProvider LUCENE_PROVIDER; + private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; + + public static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; + public static final String CUVS_META_CODEC_EXT = "vemc"; + public static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex"; + public static final String CUVS_INDEX_EXT = "vcag"; + public static final int VERSION_START = 0; + public static final int VERSION_CURRENT = VERSION_START; + + private GPUSearchParams gpuSearchParams; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + FLAT_VECTORS_FORMAT = + LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes the {@link CuVS2510GPUVectorsFormat} with default parameter values. + * + * @throws LibraryException if the native library fails to load + */ + public CuVS2510GPUVectorsFormat() { + this(new GPUSearchParams.Builder().build()); + } + + /** + * Initializes the {@link CuVS2510GPUVectorsFormat} with an instance of {@link GPUSearchParams}. + * + * @param gpuSearchParams An instance of {@link GPUSearchParams} + * @throws LibraryException if the native library fails to load + */ + public CuVS2510GPUVectorsFormat(GPUSearchParams gpuSearchParams) { + super("CuVS2510GPUVectorsFormat"); + this.gpuSearchParams = gpuSearchParams; + } + + /** + * Returns a KnnVectorsReader instance to write the vectors to the index. + */ + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + assertIsSupported(); + var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + return new CuVS2510GPUVectorsWriter(state, gpuSearchParams, flatWriter); + } + + /** + * Returns a KnnVectorsReader instance to read the vectors from the index. + */ + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + assertIsSupported(); + return new CuVS2510GPUVectorsReader(state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + } + + /** + * Returns the maximum number of vector dimensions supported by this codec for the given field name. + */ + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMENSIONS; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java new file mode 100644 index 0000000000..4480d7d78f --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -0,0 +1,629 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_START; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; + +import com.nvidia.cuvs.BruteForceIndex; +import com.nvidia.cuvs.BruteForceQuery; +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraQuery; +import com.nvidia.cuvs.CagraSearchParams; +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.BitSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.stream.Stream; +import java.util.stream.StreamSupport; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntObjectHashMap; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.store.ChecksumIndexInput; +import org.apache.lucene.store.DataInput; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IOContext.Context; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.ReadAdvice; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.hnsw.IntToIntFunction; + +/** + * KnnVectorsReader instance associated with cuVS format for reading vectors from an index. + * + * @since 25.10 + */ +public class CuVS2510GPUVectorsReader extends KnnVectorsReader { + + private static final LuceneProvider LUCENE_PROVIDER; + private static final List VECTOR_SIMILARITY_FUNCTIONS; + + private final FlatVectorsReader flatVectorsReader; + private final FieldInfos fieldInfos; + private final IntObjectHashMap fields; + private final IntObjectHashMap cuvsIndices; + private final IndexInput cuvsIndexInput; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes the {@link CuVS2510GPUVectorsReader}, checks and loads the index. + * + * @param state instance of the SegmentReadState + * @param flatReader instance of the FlatVectorsReader + * + * @throws IOException I/O exception + */ + public CuVS2510GPUVectorsReader(SegmentReadState state, FlatVectorsReader flatReader) + throws IOException { + this.flatVectorsReader = flatReader; + this.fieldInfos = state.fieldInfos; + this.fields = new IntObjectHashMap<>(); + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, CUVS_META_CODEC_EXT); + boolean success = false; + int versionMeta = -1; + try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName)) { + Throwable priorException = null; + try { + versionMeta = + CodecUtil.checkIndexHeader( + meta, + CUVS_META_CODEC_NAME, + VERSION_START, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + readFields(meta); + } catch (Throwable exception) { + priorException = exception; + } finally { + CodecUtil.checkFooter(meta, priorException); + } + var ioContext = state.context.withReadAdvice(ReadAdvice.SEQUENTIAL); + cuvsIndexInput = openCuVSInput(state, versionMeta, ioContext); + /* + * Only load indexes on the GPU when this reader is opening for searches. + * Do not load indexes on the GPU when this reader is opening during merge calls. + * With this approach we reduce device memory usage by approximately 50% during merges. + */ + if (state.context.context().equals(Context.MERGE)) { + cuvsIndices = null; + } else { + cuvsIndices = loadCuVSIndices(); + } + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Opens and returns the IndexInput for the segment file. + * + * @param state instance of the SegmentReadState + * @param versionMeta the version number + * @param context instance of the IOContext + * @return an instance of the IndexInput + * @throws IOException + */ + private static IndexInput openCuVSInput( + SegmentReadState state, int versionMeta, IOContext context) throws IOException { + String fileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, CUVS_INDEX_EXT); + IndexInput in = state.directory.openInput(fileName, context); + boolean success = false; + try { + int versionVectorData = + CodecUtil.checkIndexHeader( + in, + CUVS_INDEX_CODEC_NAME, + VERSION_START, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + checkVersion(versionMeta, versionVectorData, in); + CodecUtil.retrieveChecksum(in); + success = true; + return in; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(in); + } + } + } + + /** + * Confirms that the vector dimensions are as expected. + * + * @param info instance of the FieldInfo that describes document fields + * @param fieldEntry instance of the FieldEntry that holds the meta information for the field + */ + private void validateFieldEntry(FieldInfo info, FieldEntry fieldEntry) { + int dimension = info.getVectorDimension(); + if (dimension != fieldEntry.dims()) { + throw new IllegalStateException( + "Inconsistent vector dimension for field=\"" + + info.name + + "\"; " + + dimension + + " != " + + fieldEntry.dims()); + } + } + + /** + * Reads the fieldInfo for each index field and loads FieldEntry in a map. + * + * @param meta instance of the ChecksumIndexInput + * @throws IOException + */ + private void readFields(ChecksumIndexInput meta) throws IOException { + for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) { + FieldInfo info = fieldInfos.fieldInfo(fieldNumber); + if (info == null) { + throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta); + } + FieldEntry fieldEntry = readField(meta, info); + validateFieldEntry(info, fieldEntry); + fields.put(info.number, fieldEntry); + } + } + + /** + * Checks the distance function validity and returns it. + * + * @param input instance of DataInput + * @return an instance of VectorSimilarityFunction + * @throws IOException + */ + static VectorSimilarityFunction readSimilarityFunction(DataInput input) throws IOException { + int i = input.readInt(); + if (i < 0 || i >= VECTOR_SIMILARITY_FUNCTIONS.size()) { + throw new IllegalArgumentException("invalid distance function: " + i); + } + return VECTOR_SIMILARITY_FUNCTIONS.get(i); + } + + /** + * Reads the vector encoding (The numeric data type of the vector values) from the DataInput. + * + * @param input instance of DataInput + * @return the vector encoding + * @throws IOException + */ + static VectorEncoding readVectorEncoding(DataInput input) throws IOException { + int encodingId = input.readInt(); + if (encodingId < 0 || encodingId >= VectorEncoding.values().length) { + throw new CorruptIndexException("Invalid vector encoding id: " + encodingId, input); + } + return VectorEncoding.values()[encodingId]; + } + + /** + * Reads the field from IndexInput using FieldInfo. + * + * @param input instance of IndexInput + * @param info instance of FieldInfo + * @return the field entry + * @throws IOException + */ + private FieldEntry readField(IndexInput input, FieldInfo info) throws IOException { + VectorEncoding vectorEncoding = readVectorEncoding(input); + VectorSimilarityFunction similarityFunction = readSimilarityFunction(input); + if (similarityFunction != info.getVectorSimilarityFunction()) { + throw new IllegalStateException( + "Inconsistent vector similarity function for field=\"" + + info.name + + "\"; " + + similarityFunction + + " != " + + info.getVectorSimilarityFunction()); + } + return FieldEntry.readEntry(input, vectorEncoding, info.getVectorSimilarityFunction()); + } + + /** + * Gets the FieldEntry from the map using the field name. Check the encoding as well. + * + * @param field name of the field + * @param expectedEncoding expected encoding + * @return an instance of FieldEntry that Holds the meta information for the field + */ + private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) { + final FieldInfo info = fieldInfos.fieldInfo(field); + final FieldEntry fieldEntry; + if (info == null || (fieldEntry = fields.get(info.number)) == null) { + throw new IllegalArgumentException("field=\"" + field + "\" not found"); + } + if (fieldEntry.vectorEncoding != expectedEncoding) { + throw new IllegalArgumentException( + "field=\"" + + field + + "\" is encoded as: " + + fieldEntry.vectorEncoding + + " expected: " + + expectedEncoding); + } + return fieldEntry; + } + + /** + * Invokes loadCuVSIndex for each field and returns the map of {@link GPUIndex}. + * + * @return the map containing {@link GPUIndex} objects + * @throws IOException + */ + private IntObjectHashMap loadCuVSIndices() throws IOException { + var indices = new IntObjectHashMap(); + for (var field : fields) { + var fieldEntry = field.value; + int fieldNumber = field.key; + var cuvsIndex = loadCuVSIndex(fieldEntry); + indices.put(fieldNumber, cuvsIndex); + } + return indices; + } + + /** + * Loads the CAGRA and brute force index (if exists) onto the GPU. + * + * @param fieldEntry instance of {@link FieldEntry} + * @return return the instance of {@link GPUIndex} + * @throws IOException + */ + private GPUIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { + CagraIndex cagraIndex = null; + BruteForceIndex bruteForceIndex = null; + try { + long len = fieldEntry.cagraIndexLength(); + if (len > 0) { + long off = fieldEntry.cagraIndexOffset(); + try (var slice = cuvsIndexInput.slice("cagra index", off, len); + var in = new IndexInputInputStream(slice)) { + cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()).from(in).build(); + } + } + len = fieldEntry.bruteForceIndexLength(); + if (len > 0) { + long off = fieldEntry.bruteForceIndexOffset(); + try (var slice = cuvsIndexInput.slice("bf index", off, len); + var in = new IndexInputInputStream(slice)) { + bruteForceIndex = BruteForceIndex.newBuilder(getCuVSResourcesInstance()).from(in).build(); + } + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + return new GPUIndex(cagraIndex, bruteForceIndex); + } + + /** + * Closes the resources. + */ + @Override + public void close() throws IOException { + var closeableStream = Stream.of(flatVectorsReader, cuvsIndexInput); + IOUtils.close(closeableStream::iterator); + if (cuvsIndices != null) { + var indexClosableStream = stream(cuvsIndices.values().iterator()).map(cursor -> cursor.value); + IOUtils.close(indexClosableStream::iterator); + } + closeCuVSResourcesInstance(); + } + + static Stream stream(Iterator iterator) { + return StreamSupport.stream(((Iterable) () -> iterator).spliterator(), false); + } + + /** + * Checks consistency of this reader. + */ + @Override + public void checkIntegrity() throws IOException { + flatVectorsReader.checkIntegrity(); + CodecUtil.checksumEntireFile(cuvsIndexInput); + } + + /** + * Returns the FloatVectorValues for the given field. + */ + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return flatVectorsReader.getFloatVectorValues(field); + } + + /** + * Returns the ByteVectorValues for the given field. + * + * This is not supported. + */ + @Override + public ByteVectorValues getByteVectorValues(String field) { + throw new UnsupportedOperationException("byte vectors are not currently supported"); + } + + /** Native float to float function */ + private interface FloatToFloatFunction { + float apply(float v); + } + + /** + * Get the score normalization function. + * + * @param sim instance of VectorSimilarityFunction + * @return an instance of the FloatToFloatFunction + */ + private static FloatToFloatFunction getScoreNormalizationFunc(VectorSimilarityFunction sim) { + // TODO: check for different similarities + return score -> (1f / (1f + score)); + } + + /** + * Returns the k nearest neighbor documents using cuVS's CAGRA or brute force algorithm for this field, to the given vector. + */ + @Override + public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) + throws IOException { + var fieldEntry = getFieldEntry(field, VectorEncoding.FLOAT32); + if (fieldEntry.count() == 0 || knnCollector.k() == 0) { + return; + } + + var fieldNumber = fieldInfos.fieldInfo(field).number; + GPUIndex cuvsIndex = cuvsIndices != null ? cuvsIndices.get(fieldNumber) : null; + if (cuvsIndex == null) { + throw new IllegalStateException("Index not found for field:" + field); + } + + final FloatVectorValues rawValues = flatVectorsReader.getFloatVectorValues(field); + final Bits acceptedOrds = rawValues.getAcceptOrds(acceptDocs); + BitSet[] mask = null; + int maskLength = 0; + int topK = knnCollector.k(); + + if (acceptDocs != null) { + mask = new BitSet[1]; // As there is only one query "target" + mask[0] = new BitSet(acceptedOrds.length()); + /* + * Need to find if there is a better alternative for below loop + * in subsequent code improvement iterations. This is needed as + * there is a difference between Lucene's Bits "acceptDocs" and + * what our API accepts. + */ + for (int i = 0; i < acceptedOrds.length(); i++) { + if (acceptedOrds.get(i)) { + mask[0].set(i); + } + } + topK = Math.min(knnCollector.k() + 10, mask[0].cardinality()); + maskLength = mask[0].length(); + } + + try { + List> searchResult = null; + if (knnCollector.k() <= 1024 && cuvsIndex.getCagraIndex() != null) { + CagraSearchParams searchParams; + if (knnCollector instanceof GPUPerLeafCuVSKnnCollector) { + GPUPerLeafCuVSKnnCollector collector = (GPUPerLeafCuVSKnnCollector) knnCollector; + searchParams = + new CagraSearchParams.Builder() + .withItopkSize(Math.max(collector.getiTopK(), topK)) + .withSearchWidth(collector.getSearchWidth()) + .build(); + } else { + // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK + searchParams = new CagraSearchParams.Builder().withItopkSize(topK).build(); + } + CagraIndex cagraIndex = cuvsIndex.getCagraIndex(); + assert cagraIndex != null; + CagraQuery query = null; + + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + getCuVSResourcesInstance(), 1, target.length, CuVSMatrix.DataType.FLOAT); + builder.addVector(target); + CuVSMatrix queryVector = builder.build(); + + if (acceptDocs != null) { + query = + new CagraQuery.Builder(getCuVSResourcesInstance()) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queryVector) + .withPrefilter(mask[0], maskLength) + .build(); + } else { + query = + new CagraQuery.Builder(getCuVSResourcesInstance()) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queryVector) + .build(); + } + searchResult = cagraIndex.search(query).getResults(); + } else { + BruteForceIndex bruteforceIndex = cuvsIndex.getBruteforceIndex(); + assert bruteforceIndex != null; + BruteForceQuery query = null; + float[][] queryVector = new float[][] {target}; + if (acceptDocs != null) { + query = + new BruteForceQuery.Builder(getCuVSResourcesInstance()) + .withQueryVectors(queryVector) + .withPrefilters(mask, maskLength) + .withTopK(topK) + .build(); + } else { + query = + new BruteForceQuery.Builder(getCuVSResourcesInstance()) + .withQueryVectors(queryVector) + .withTopK(topK) + .build(); + } + searchResult = bruteforceIndex.search(query).getResults(); + } + + // List expected to have only one entry because of single query "target". + assert searchResult.size() == 1; + final IntToIntFunction ordToDocFunction = (IntToIntFunction) rawValues::ordToDoc; + final FloatToFloatFunction scoreCorrectionFunction = + getScoreNormalizationFunc(fieldEntry.similarityFunction); + + for (Entry entry : searchResult.getFirst().entrySet()) { + int ord = entry.getKey(); + float score = entry.getValue(); + if (knnCollector.earlyTerminated()) { + break; + } + if (ord < 0) { + continue; + } + float correctedScore = scoreCorrectionFunction.apply(score); + int doc = ordToDocFunction.apply(ord); + knnCollector.incVisitedCount(1); + knnCollector.collect(doc, correctedScore); + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Return the k nearest neighbor documents as determined by comparison of their vector values for this field, to the given vector. + * + * This is not supported. + */ + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) + throws IOException { + throw new UnsupportedOperationException("Byte vectors are not currently supported"); + } + + /** + * Holds the meta information for the field. + */ + record FieldEntry( + VectorEncoding vectorEncoding, + VectorSimilarityFunction similarityFunction, + int dims, + int count, + long cagraIndexOffset, + long cagraIndexLength, + long bruteForceIndexOffset, + long bruteForceIndexLength) { + + /** + * Returns an instance of FieldEntry. + * + * @param input instance of IndexInput + * @param vectorEncoding The numeric data type of the vector values + * @param similarityFunction Vector similarity function; used in search to return top K most similar vectors to a target vector + * @return an instance of FieldEntry + * @throws IOException I/O Exceptions + */ + static FieldEntry readEntry( + IndexInput input, + VectorEncoding vectorEncoding, + VectorSimilarityFunction similarityFunction) + throws IOException { + var dims = input.readInt(); + var count = input.readInt(); + var cagraIndexOffset = input.readVLong(); + var cagraIndexLength = input.readVLong(); + var bruteForceIndexOffset = input.readVLong(); + var bruteForceIndexLength = input.readVLong(); + return new FieldEntry( + vectorEncoding, + similarityFunction, + dims, + count, + cagraIndexOffset, + cagraIndexLength, + bruteForceIndexOffset, + bruteForceIndexLength); + } + } + + /** + * Checks the version and throws CorruptIndexException on mismatch. + * + * @param versionMeta + * @param versionVectorData + * @param in + * @throws CorruptIndexException + */ + private static void checkVersion(int versionMeta, int versionVectorData, IndexInput in) + throws CorruptIndexException { + if (versionMeta != versionVectorData) { + throw new CorruptIndexException( + "Format versions mismatch: meta=" + + versionMeta + + ", " + + CUVS_META_CODEC_NAME + + "=" + + versionVectorData, + in); + } + } + + /** + * Gets the instance of FieldInfos. + * + * @return the instance of FieldInfos + */ + public FieldInfos getFieldInfos() { + return fieldInfos; + } + + /** + * Gets the map of {@link GPUIndex} objects. + * + * @return the map of GPU index objects + */ + public IntObjectHashMap getCuvsIndexes() { + return cuvsIndices; + } + + /** + * Gets the map of FieldEntry objects that hold the meta information for the field. + * + * @return the map of FieldEntry objects + */ + public IntObjectHashMap getFieldEntries() { + return fields; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java new file mode 100644 index 0000000000..757ca0573e --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -0,0 +1,539 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.info; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +import com.nvidia.cuvs.BruteForceIndex; +import com.nvidia.cuvs.BruteForceIndexParams; +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntObjectHashMap; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; + +/** + * extends upon KnnVectorsWriter and has implementation for critical methods like flush, merge etc. + * + * @since 25.10 + */ +public class CuVS2510GPUVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(CuVS2510GPUVectorsWriter.class); + private static final String COMPONENT = "CuVS2510GPUVectorsWriter"; + private static final LuceneProvider LUCENE_PROVIDER; + private static final List VECTOR_SIMILARITY_FUNCTIONS; + private static final int MIN_CAGRA_INDEX_SIZE = 2; + + private final GPUSearchParams gpuSearchParams; + private final FlatVectorsWriter flatVectorsWriter; + private final List fields = new ArrayList<>(); + private final InfoStream infoStream; + private IndexOutput meta = null, cuvsIndex = null; + private boolean finished; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * The cuVS index Types. + */ + public enum IndexType { + + /** Builds a CAGRA index. */ + CAGRA(true, false), + + /** Builds a Brute Force index. */ + BRUTE_FORCE(false, true), + + /** Builds both - CAGRA and Brute Force indexes. */ + CAGRA_AND_BRUTE_FORCE(true, true); + private final boolean cagra, bruteForce; + + IndexType(boolean cagra, boolean bruteForce) { + this.cagra = cagra; + this.bruteForce = bruteForce; + } + + public boolean isCagra() { + return cagra; + } + + public boolean isBruteForce() { + return bruteForce; + } + } + + /** + * Initializes {@link CuVS2510GPUVectorsWriter}. + * + * @param state instance of the SegmentWriteState + * @param gpuSearchParams An instance of {@link GPUSearchParams} + * @param flatVectorsWriter instance of FlatVectorsWriter + * + * @throws IOException I/O exceptions + */ + public CuVS2510GPUVectorsWriter( + SegmentWriteState state, GPUSearchParams gpuSearchParams, FlatVectorsWriter flatVectorsWriter) + throws IOException { + super(); + this.gpuSearchParams = gpuSearchParams; + this.flatVectorsWriter = flatVectorsWriter; + this.infoStream = state.infoStream; + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, CUVS_META_CODEC_EXT); + String cagraFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, CUVS_INDEX_EXT); + boolean success = false; + try { + meta = state.directory.createOutput(metaFileName, state.context); + cuvsIndex = state.directory.createOutput(cagraFileName, state.context); + CodecUtil.writeIndexHeader( + meta, + CUVS_META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + cuvsIndex, + CUVS_INDEX_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Add new field for indexing. + */ + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + var encoding = fieldInfo.getVectorEncoding(); + if (encoding != FLOAT32) { + throw new IllegalArgumentException("Expected float32, got:" + encoding); + } + var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); + @SuppressWarnings("unchecked") + var flatWriter = (FlatFieldVectorsWriter) writer; + var cuvsFieldWriter = new GPUFieldWriter(fieldInfo, flatWriter); + fields.add(cuvsFieldWriter); + return writer; + } + + /** + * Creates CAGRA and/or brute force indexes and writes them. + * + * @param fieldInfo Instance of the FieldInFo to use + * @param vectors list of float vectors to index + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors != null && vectors.size() == 0) { + writeEmpty(fieldInfo); + return; + } + long cagraIndexOffset, cagraIndexLength = 0L; + long bruteForceIndexOffset, bruteForceIndexLength = 0L; + + /* + * CAGRA has an issue when asked to build an index with just one vector. + * Hence, we currently fallback to brute force in such a case. + */ + IndexType indexType = + gpuSearchParams.getIndexType().isCagra() && vectors.size() < MIN_CAGRA_INDEX_SIZE + ? IndexType.BRUTE_FORCE + : gpuSearchParams.getIndexType(); + + try { + cagraIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.isCagra()) { + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + try { + CuVSMatrix cagraDataset = + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + writeCagraIndex(cagraIndexOutputStream, cagraDataset); + } catch (Throwable t) { + // Fallback to brute force in a few cases, for now. + Utils.handleThrowableWithIgnore(t, t.getMessage()); + indexType = IndexType.BRUTE_FORCE; + } + cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + } + bruteForceIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.isBruteForce()) { + var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix bruteforceDataset = + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + + writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); + bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; + } + writeMeta( + fieldInfo, + vectors.size(), + cagraIndexOffset, + cagraIndexLength, + bruteForceIndexOffset, + bruteForceIndexLength); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Builds and writes the CAGRA index. + * + * @param os Instance of the OutputStream + * @param dataset The instance of CuVSMatrix holding the dataset + * @throws Throwable + */ + private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + CagraIndexParams params = + CagraIndexParamsFactory.create(gpuSearchParams, dataset.size(), dataset.columns()); + CagraIndex index = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + index.serialize(os); + index.close(); + } + + /** + * Builds and writes the brute force index. + * + * @param os Instance of OutputStream to write the index bytes to + * @param dataset Instance of CuVSMatrix that holds the data set + * @throws Throwable + */ + private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + BruteForceIndexParams params = + new BruteForceIndexParams.Builder() + .withNumWriterThreads(gpuSearchParams.getWriterThreads()) + .build(); + var index = + BruteForceIndex.newBuilder(getCuVSResourcesInstance()) + .withIndexParams(params) + .withDataset(dataset) + .build(); + index.serialize(os); + index.close(); + } + + /** + * Creates the CAGRA and/or brute force indexes and writes them to the disk. + */ + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + flatVectorsWriter.flush(maxDoc, sortMap); + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } + } + + /** + * Calls the method that builds indexes and writes them to the disk. + * + * @param fieldData reference to the {@link GPUFieldWriter} + * @throws IOException + */ + private void writeField(GPUFieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); + } + + /** + * Builds indexes and writes them to the disk. + * + * @param fieldData reference to the {@link GPUFieldWriter} + * @param sortMap reference to DocMap + * @throws IOException I/O Exceptions + */ + private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) + throws IOException { + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; + mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); + List sortedVectors = new ArrayList(); + for (int i = 0; i < fieldData.getVectors().size(); i++) { + sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); + } + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); + } + + /** + * Writes empty meta information for the field. + * + * @param fieldInfo instance of the FieldInfo + * @throws IOException I/O Exceptions + */ + private void writeEmpty(FieldInfo fieldInfo) throws IOException { + writeMeta(fieldInfo, 0, 0L, 0L, 0L, 0L); + } + + /** + * Writes the meta information for the index. + * + * @param field instance of FieldInfo + * @param count number of vectors + * @param cagraIndexOffset CAGRA index offset + * @param cagraIndexLength CAGRA index length + * @param bruteForceIndexOffset Brute force index offset + * @param bruteForceIndexLength Brute force index length + * @throws IOException I/O Exceptions + */ + private void writeMeta( + FieldInfo field, + int count, + long cagraIndexOffset, + long cagraIndexLength, + long bruteForceIndexOffset, + long bruteForceIndexLength) + throws IOException { + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); + meta.writeInt(field.getVectorDimension()); + meta.writeInt(count); + meta.writeVLong(cagraIndexOffset); + meta.writeVLong(cagraIndexLength); + meta.writeVLong(bruteForceIndexOffset); + meta.writeVLong(bruteForceIndexLength); + } + + static int distFuncToOrd(VectorSimilarityFunction func) { + for (int i = 0; i < VECTOR_SIMILARITY_FUNCTIONS.size(); i++) { + if (VECTOR_SIMILARITY_FUNCTIONS.get(i).equals(func)) { + return (byte) i; + } + } + throw new IllegalArgumentException("Invalid distance function: " + func); + } + + /** + * Uses the cuVS API to merge CAGRA indexes. + * + * This is currently (and intentionally) marked as unused and will be plugged in later. + * + * @param fieldInfo instance of the FieldInfo + * @param mergeState instance of the MergeState + * @throws IOException I/O Exceptions + */ + @SuppressWarnings("unused") + private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + List cagraIndexes = new ArrayList<>(); + // We need this count so that the merged segment's meta information has the vector count. + int totalVectorCount = 0; + for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { + KnnVectorsReader knnReader = mergeState.knnVectorsReaders[i]; + // Access the CAGRA index for this field from the reader + if (knnReader != null) { + if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { + if (cvr != null) { + totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); + CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); + if (cagraIndex != null) { + cagraIndexes.add(cagraIndex); + } + } + } else { + // This should never happen + throw new RuntimeException( + "Reader is not of CuVSVectorsReader type. Instead it is: " + knnReader.getClass()); + } + } + } + assert cagraIndexes.size() > 1; + CagraIndex mergedIndex = + CagraIndex.merge(cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()])); + writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + info( + infoStream, + COMPONENT, + "Successfully merged " + cagraIndexes.size() + " CAGRA indexes using native merge API"); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. + */ + private CagraIndex getCagraIndexFromReader(CuVS2510GPUVectorsReader reader, String fieldName) { + try { + IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); + FieldInfos fieldInfos = reader.getFieldInfos(); + FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); + if (fieldInfo != null) { + GPUIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); + if (cuvsIndex != null) { + return cuvsIndex.getCagraIndex(); + } + } + } catch (Exception e) { + info( + infoStream, + COMPONENT, + "Failed to extract CAGRA index for field " + fieldName + ": " + e.getMessage()); + throw e; + } + return null; + } + + /** + * Writes a pre-built merged CAGRA index to the output. + */ + private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, int vectorCount) + throws IOException { + try { + long cagraIndexOffset = cuvsIndex.getFilePointer(); + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + Path tmpFile = + Files.createTempFile(getCuVSResourcesInstance().tempDirectory(), "mergedindex", "cag"); + mergedIndex.serialize(cagraIndexOutputStream, tmpFile); + long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L); + mergedIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Creates List from merged vectors. + */ + private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List res = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + int ordinal = iter.index(); + float[] vector = mergedVectorValues.vectorValue(ordinal); + res.add(vector.clone()); + } + return res; + } + + /** + * Fallback method that rebuilds indexes from merged vectors. + * Used when native CAGRA merge() is not possible. Also used + * when non-CAGRA index types are used (for e.g. Brute Force index). + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Write field for merging. + */ + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + flatVectorsWriter.mergeOneField(fieldInfo, mergeState); + vectorBasedMerge(fieldInfo, mergeState); + } + + /** + * Returns the memory usage of this object in bytes. + */ + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (var field : fields) { + total += field.ramBytesUsed(); + } + return total; + } + + /** + * Called once at the end before close. + */ + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + flatVectorsWriter.finish(); + if (meta != null) { + // write end of fields marker + meta.writeInt(-1); + CodecUtil.writeFooter(meta); + } + if (cuvsIndex != null) { + CodecUtil.writeFooter(cuvsIndex); + } + } + + /** + * Close the applicable resources. + */ + @Override + public void close() throws IOException { + IOUtils.close(meta, cuvsIndex, flatVectorsWriter); + closeCuVSResourcesInstance(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java new file mode 100644 index 0000000000..8e2d9a70e2 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; + +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; +import java.io.IOException; +import java.util.List; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.util.RamUsageEstimator; + +public class FieldWriter extends KnnFieldVectorsWriter { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); + + private final FieldInfo fieldInfo; + private final FlatFieldVectorsWriter flatFieldVectorsWriter; + private int lastDocID = -1; + private QuantizationType quantizationType; + + @SuppressWarnings("unchecked") + public FieldWriter( + QuantizationType quantizationType, + FieldInfo fieldInfo, + FlatFieldVectorsWriter flatFieldVectorsWriter) { + this.quantizationType = quantizationType; + this.fieldInfo = fieldInfo; + this.flatFieldVectorsWriter = (FlatFieldVectorsWriter) flatFieldVectorsWriter; + } + + @Override + public void addValue(int docID, Object vectorValue) throws IOException { + if (docID == lastDocID) { + throw new IllegalArgumentException( + "VectorValuesField \"" + + fieldInfo.name + + "\" appears more than once in this document (only one value is allowed per" + + " field)"); + } + flatFieldVectorsWriter.addValue(docID, (float[]) vectorValue); + } + + List getByteVectors() { + if (quantizationType == QuantizationType.BINARY) { + return quantizeFloatVectorsToBinary(flatFieldVectorsWriter.getVectors()); + } else if (quantizationType == QuantizationType.SCALAR) { + return quantizeFloatVectorsToScalar(flatFieldVectorsWriter.getVectors()); + } else { + throw new UnsupportedOperationException("Not applicable for QuantizationType.NONE"); + } + } + + List getFloatVectors() { + return flatFieldVectorsWriter.getVectors(); + } + + FieldInfo fieldInfo() { + return fieldInfo; + } + + DocsWithFieldSet getDocsWithFieldSet() { + return flatFieldVectorsWriter.getDocsWithFieldSet(); + } + + @Override + public Object copyValue(Object vectorValue) { + throw new UnsupportedOperationException(); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java new file mode 100644 index 0000000000..56c59b6b0c --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -0,0 +1,168 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.BruteForceIndex; +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSMatrix.Builder; +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.GPUInfoProvider; +import com.nvidia.cuvs.HnswIndex; +import com.nvidia.cuvs.HnswIndexParams; +import com.nvidia.cuvs.TieredIndex; +import com.nvidia.cuvs.spi.CuVSProvider; +import java.lang.invoke.MethodHandle; +import java.nio.file.Path; +import java.util.logging.Level; + +class FilterCuVSProvider implements CuVSProvider { + + private final CuVSProvider delegate; + + FilterCuVSProvider(CuVSProvider delegate) { + this.delegate = delegate; + } + + @Override + public Path nativeLibraryPath() { + return CuVSProvider.TMPDIR; + } + + @Override + public CuVSResources newCuVSResources(Path tempPath) throws Throwable { + return delegate.newCuVSResources(tempPath); + } + + @Override + public BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) + throws UnsupportedOperationException { + return delegate.newBruteForceIndexBuilder(cuVSResources); + } + + @Override + public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) + throws UnsupportedOperationException { + return delegate.newCagraIndexBuilder(cuVSResources); + } + + @Override + public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) + throws UnsupportedOperationException { + return delegate.newHnswIndexBuilder(cuVSResources); + } + + @Override + public CagraIndex mergeCagraIndexes(CagraIndex[] arg0) throws Throwable { + return delegate.mergeCagraIndexes(arg0); + } + + @Override + public GPUInfoProvider gpuInfoProvider() { + return delegate.gpuInfoProvider(); + } + + @Override + public Builder newHostMatrixBuilder(long rows, long cols, DataType dataType) { + return delegate.newHostMatrixBuilder(rows, cols, dataType); + } + + @Override + public Builder newHostMatrixBuilder( + long rows, long cols, int maxRows, int maxCols, DataType dataType) { + return delegate.newHostMatrixBuilder(rows, cols, maxRows, maxCols, dataType); + } + + @Override + public Builder newDeviceMatrixBuilder( + CuVSResources resources, long rows, long cols, DataType dataType) { + return delegate.newDeviceMatrixBuilder(resources, rows, cols, dataType); + } + + @Override + public Builder newDeviceMatrixBuilder( + CuVSResources resources, long rows, long cols, int maxRows, int maxCols, DataType dataType) { + return delegate.newDeviceMatrixBuilder(resources, rows, cols, maxRows, maxCols, dataType); + } + + @Override + public MethodHandle newNativeMatrixBuilder() { + return delegate.newNativeMatrixBuilder(); + } + + @Override + public MethodHandle newNativeMatrixBuilderWithStrides() { + return delegate.newNativeMatrixBuilderWithStrides(); + } + + @Override + public CuVSMatrix newMatrixFromArray(float[][] vectors) { + return delegate.newMatrixFromArray(vectors); + } + + @Override + public CuVSMatrix newMatrixFromArray(int[][] vectors) { + return delegate.newMatrixFromArray(vectors); + } + + @Override + public CuVSMatrix newMatrixFromArray(byte[][] vectors) { + return delegate.newMatrixFromArray(vectors); + } + + @Override + public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) + throws UnsupportedOperationException { + return delegate.newTieredIndexBuilder(cuVSResources); + } + + @Override + public CagraIndexParams cagraIndexParamsFromHnswParams( + long arg0, long arg1, int arg2, int arg3, HnswHeuristicType arg4, CuvsDistanceType arg5) { + return delegate.cagraIndexParamsFromHnswParams(arg0, arg1, arg2, arg3, arg4, arg5); + } + + @Override + public Level getLogLevel() { + return delegate.getLogLevel(); + } + + @Override + public void setLogLevel(Level arg0) { + delegate.setLogLevel(arg0); + } + + @Override + public HnswIndex hnswIndexFromCagra(HnswIndexParams arg0, CagraIndex arg1) throws Throwable { + return delegate.hnswIndexFromCagra(arg0, arg1); + } + + @Override + public void enableRMMManagedPooledMemory(int arg0, int arg1) { + delegate.enableRMMManagedPooledMemory(arg0, arg1); + } + + @Override + public void enableRMMPooledMemory(int arg0, int arg1) { + delegate.enableRMMPooledMemory(arg0, arg1); + } + + @Override + public void resetRMMPooledMemory() { + delegate.resetRMMPooledMemory(); + } + + @Override + public HnswIndex hnswIndexBuild(CuVSResources arg0, HnswIndexParams arg1, CuVSMatrix arg2) + throws Throwable { + return delegate.hnswIndexBuild(arg0, arg1, arg2); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java new file mode 100644 index 0000000000..1341738df8 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java @@ -0,0 +1,24 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.spi.CuVSProvider; +import com.nvidia.cuvs.spi.CuVSServiceProvider; + +/** + * A provider that creates instances of FilterCuVSProvider. + * + * @since 25.10 + */ +public class FilterCuVSServiceProvider extends CuVSServiceProvider { + + /** + * Initialize and return an CuVSProvider provided by this provider. + */ + @Override + public CuVSProvider get(CuVSProvider builtinProvider) { + return new FilterCuVSProvider(builtinProvider); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java new file mode 100644 index 0000000000..7e9f888e32 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -0,0 +1,304 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.RowView; +import java.util.ArrayList; +import java.util.List; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.NeighborArray; + +/** + * This class holds the in-memory representation of the HNSW graph + * + * @since 25.10 + */ +public class GPUBuiltHnswGraph extends HnswGraph { + + private final int size; + private final int dimensions; + private final int numLevels; + + // Store layers data - each layer has its own nodes and adjacency lists + private final List layerNodes; + private final List layerNeighbors; + + // Layer 0 is special - it contains all nodes + private final NeighborArray[] layer0Neighbors; + + /** + * Multi-layer constructor that supports arbitrary number of layers. + * + * @param size the size of the dataset + * @param dimensions the vector dimension + * @param layerNodes the nodes on the layer + * @param layerAdjacencies adjacency list + */ + public GPUBuiltHnswGraph( + int size, int dimensions, List layerNodes, List layerAdjacencies) { + + this.size = size; + this.dimensions = dimensions; + this.numLevels = layerAdjacencies.size(); + this.layerNodes = new ArrayList<>(); + this.layerNeighbors = new ArrayList<>(); + + // Process Layer 0 (base layer with all nodes) + CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); + + // Process higher layers (1 to numLevels-1) + for (int level = 1; level < numLevels; level++) { + int[] nodes = layerNodes.get(level); + CuVSMatrix adjacency = layerAdjacencies.get(level); + this.layerNodes.add(nodes); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + } + } + + /** + * Fills the neighbor array using the adjacency matrix. + * + * @param adjacency instance of adjacency CuVSMatrix + * @param size the number of nodes + * @return the NeighborArray + */ + private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + NeighborArray[] neighbors = new NeighborArray[size]; + for (int i = 0; i < size; i++) { + RowView rv = adjacency.getRow(i); + if (rv != null && rv.size() > 0) { + neighbors[i] = new NeighborArray((int) rv.size(), true); + for (int j = 0; j < rv.size(); j++) { + neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + } + } else { + neighbors[i] = new NeighborArray(0, true); + } + } + return neighbors; + } + + /** + * Get all nodes on a given level as node 0th ordinals. + */ + public NodesIterator getNodesOnLevel(int level) { + if (level == 0) { + return new Level0NodesIterator(size); + } else if (level > 0 && level < numLevels) { + int[] nodes = layerNodes.get(level - 1); + return new HigherLevelNodesIterator(nodes); + } else { + return new Level0NodesIterator(0); + } + } + + /** + * Get the neighbors for the node and the level it resides. + * + * @param level the level + * @param node the node + * @return an instance of NeighborArray + */ + public NeighborArray getNeighbors(int level, int node) { + if (level == 0 && node < size) { + return layer0Neighbors[node]; + } else if (level > 0 && level < numLevels) { + int[] nodes = layerNodes.get(level - 1); + NeighborArray[] neighbors = layerNeighbors.get(level - 1); + + // Find the index of this node in the layer + for (int i = 0; i < nodes.length; i++) { + if (nodes[i] == node) { + return neighbors[i]; + } + } + } + return null; + } + + // Implementation of abstract methods from HnswGraph + private int currentNode = -1; + private int currentLevel = -1; + private int neighborIndex = -1; + + /** + * Move the pointer to exactly the given level's target. + */ + @Override + public void seek(int level, int target) { + currentLevel = level; + currentNode = target; + neighborIndex = -1; + } + + /** + * Iterates over the neighbor list. + */ + @Override + public int nextNeighbor() { + if (currentLevel == 0 + && currentNode >= 0 + && currentNode < size + && layer0Neighbors[currentNode] != null) { + neighborIndex++; + if (neighborIndex < layer0Neighbors[currentNode].size()) { + int neighborNode = layer0Neighbors[currentNode].nodes()[neighborIndex]; + if (neighborNode >= 0 && neighborNode < size) { + return neighborNode; + } else { + return nextNeighbor(); // Skip invalid neighbor + } + } + } else if (currentLevel > 0 && currentLevel < numLevels) { + // Handle higher layers + NeighborArray neighbors = getNeighbors(currentLevel, currentNode); + if (neighbors != null) { + neighborIndex++; + if (neighborIndex < neighbors.size()) { + return neighbors.nodes()[neighborIndex]; + } + } + } + return NO_MORE_DOCS; + } + + /** + * Returns graph's entry point on the top level. + */ + @Override + public int entryNode() { + // Entry node should be from the highest layer + if (numLevels > 1) { + int topLevel = numLevels - 1; + int[] topLayerNodes = layerNodes.get(topLevel - 1); + if (topLayerNodes != null && topLayerNodes.length > 0) { + // Use random node from top layer with fixed seed for reproducibility + java.util.Random random = new java.util.Random(44); + int randomIndex = random.nextInt(topLayerNodes.length); + return topLayerNodes[randomIndex]; + } + } + return 0; // Default to node 0 for single-layer graphs + } + + /** + * returns M, the maximum number of connections for a node. + */ + @Override + public int maxConn() { + // Return the maximum degree across all nodes in layer 0 + int max = 0; + for (NeighborArray neighbor : layer0Neighbors) { + if (neighbor != null) { + max = Math.max(max, neighbor.size()); + } + } + return max; + } + + /** + * Returns the neighbor count. + */ + @Override + public int neighborCount() { + if (currentLevel == 0 + && currentNode >= 0 + && currentNode < size + && layer0Neighbors[currentNode] != null) { + return layer0Neighbors[currentNode].size(); + } else if (currentLevel > 0 && currentLevel < numLevels) { + NeighborArray neighbors = getNeighbors(currentLevel, currentNode); + return neighbors != null ? neighbors.size() : 0; + } + return 0; + } + + // NodesIterator for level 0 + private static class Level0NodesIterator extends NodesIterator { + private int current = -1; + + Level0NodesIterator(int size) { + super(size); + } + + @Override + public boolean hasNext() { + return current + 1 < size; + } + + @Override + public int nextInt() { + return ++current; + } + + @Override + public int consume(int[] dest) { + int numToCopy = Math.min(dest.length, size - (current + 1)); + for (int i = 0; i < numToCopy; i++) { + dest[i] = ++current; + } + return numToCopy; + } + } + + // NodesIterator for higher layers + private static class HigherLevelNodesIterator extends NodesIterator { + private final int[] nodeIds; + private int current = -1; + + HigherLevelNodesIterator(int[] nodeIds) { + super(nodeIds.length); + this.nodeIds = nodeIds; + } + + @Override + public boolean hasNext() { + return current + 1 < nodeIds.length; + } + + @Override + public int nextInt() { + return nodeIds[++current]; + } + + @Override + public int consume(int[] dest) { + int numToCopy = Math.min(dest.length, nodeIds.length - (current + 1)); + for (int i = 0; i < numToCopy; i++) { + dest[i] = nodeIds[++current]; + } + return numToCopy; + } + } + + /** + * Returns the number of nodes in the graph. + */ + public int size() { + return size; + } + + /** + * Returns the number of levels in the HNSW graph. + * + * @return the number of levels + */ + public int numLevels() { + return numLevels; + } + + /** + * Gets the vector dimension. + * + * @return the vector dimension + */ + public int dimensions() { + return dimensions; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java new file mode 100644 index 0000000000..d159c5c144 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java @@ -0,0 +1,117 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.util.List; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.util.RamUsageEstimator; + +/** + * cuVS based fields writer + * + * @since 25.10 + */ +/*package-private*/ class GPUFieldWriter extends KnnFieldVectorsWriter { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(GPUFieldWriter.class); + + private final FieldInfo fieldInfo; + private final FlatFieldVectorsWriter flatFieldVectorsWriter; + private int lastDocID = -1; + + public GPUFieldWriter( + FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + this.fieldInfo = fieldInfo; + this.flatFieldVectorsWriter = flatFieldVectorsWriter; + } + + /** + * Add new docID with its vector value to the given field for indexing. + */ + @Override + public void addValue(int docID, float[] vectorValue) throws IOException { + if (docID == lastDocID) { + throw new IllegalArgumentException( + "VectorValuesField \"" + + fieldInfo.name + + "\" appears more than once in this document (only one value is allowed per field)"); + } + flatFieldVectorsWriter.addValue(docID, vectorValue); + } + + /** + * Gets the list of float vectors. + * + * @return a list of float vectors + */ + List getVectors() { + return flatFieldVectorsWriter.getVectors(); + } + + /** + * Gets the vector dimension. + * + * @return the vector dimension + */ + int getVectorDimension() { + List vectors = flatFieldVectorsWriter.getVectors(); + if (vectors != null && vectors.size() > 0) { + return vectors.get(0).length; + } else { + return -1; + } + } + + /** + * Gets the field info that holds the description of the field. + * + * @return an instance of FieldInfo + */ + FieldInfo fieldInfo() { + return fieldInfo; + } + + /** + * Gets the docsWithFieldSet for the field writer. + * + * @return an instance of DocsWithFieldSet + */ + DocsWithFieldSet getDocsWithFieldSet() { + return flatFieldVectorsWriter.getDocsWithFieldSet(); + } + + /** + * Used to copy values being indexed to internal storage. + */ + @Override + public float[] copyValue(float[] vectorValue) { + throw new UnsupportedOperationException(); + } + + /** + * Returns the memory usage of this object in bytes. + */ + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed(); + } + + /** + * Returns a string containing the field name and number. + */ + @Override + public String toString() { + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(field name=").append(fieldInfo.name); + sb.append("number=").append(fieldInfo.number); + sb.append(")"); + return sb.toString(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java new file mode 100644 index 0000000000..f71c8f6f87 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java @@ -0,0 +1,148 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.BruteForceIndex; +import com.nvidia.cuvs.CagraIndex; +import java.io.Closeable; +import java.io.IOException; +import java.util.Objects; + +/** + * This class holds references to the cuVS Index (Cagra, Brute force, etc.) + * + * @since 25.10 + */ +public class GPUIndex implements Closeable { + private final CagraIndex cagraIndex; + private final BruteForceIndex bruteforceIndex; + + private int maxDocs; + private String fieldName; + private String segmentName; + private volatile boolean closed; + + /** + * Initializes an instance of {@link GPUIndex} + * + * @param segmentName the name of the segment + * @param fieldName the field name + * @param cagraIndex reference to the CagraIndex + * @param maxDocs the maximum documents + * @param bruteforceIndex reference to the BruteForceIndex + */ + public GPUIndex( + String segmentName, + String fieldName, + CagraIndex cagraIndex, + int maxDocs, + BruteForceIndex bruteforceIndex) { + this.cagraIndex = Objects.requireNonNull(cagraIndex); + this.bruteforceIndex = Objects.requireNonNull(bruteforceIndex); + this.fieldName = Objects.requireNonNull(fieldName); + this.segmentName = Objects.requireNonNull(segmentName); + if (maxDocs < 0) { + throw new IllegalArgumentException("negative maxDocs:" + maxDocs); + } + this.maxDocs = maxDocs; + } + + /** + * Initializes an instance of {@link GPUIndex} + * + * @param cagraIndex reference to the CagraIndex instance + * @param bruteforceIndex reference to the BruteForceIndex instance + */ + public GPUIndex(CagraIndex cagraIndex, BruteForceIndex bruteforceIndex) { + this.cagraIndex = cagraIndex; + this.bruteforceIndex = bruteforceIndex; + } + + /** + * Gets the reference to the CAGRA index + * + * @return an instance of CagraIndex + */ + public CagraIndex getCagraIndex() { + ensureOpen(); + return cagraIndex; + } + + /** + * Gets the reference to the Bruteforce index + * + * @return an instance of BruteForceIndex + */ + public BruteForceIndex getBruteforceIndex() { + ensureOpen(); + return bruteforceIndex; + } + + /** + * Gets the field name + * + * @return field name + */ + public String getFieldName() { + return fieldName; + } + + /** + * Gets the segment name + * + * @return segment name + */ + public String getSegmentName() { + return segmentName; + } + + /** + * Gets the max docs + * + * @return the max docs + */ + public int getMaxDocs() { + return maxDocs; + } + + /** + * Throws {@link IllegalArgumentException} if the index is closed + */ + private void ensureOpen() { + if (closed) { + throw new IllegalStateException("index is closed"); + } + } + + /** + * Closes this stream and releases any resources associated with it. + */ + @Override + public void close() throws IOException { + if (closed) { + return; + } + closed = true; + destroyIndices(); + } + + /** + * Closes the cuVS indexes. + * + * @throws IOException + */ + private void destroyIndices() throws IOException { + try { + if (cagraIndex != null) { + cagraIndex.close(); + } + if (bruteforceIndex != null) { + bruteforceIndex.close(); + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java new file mode 100644 index 0000000000..e260e26a73 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.knn.KnnCollectorManager; +import org.apache.lucene.util.Bits; + +/** + * Extends upon KnnFloatVectorQuery for GPU-only search. + * + * @since 25.10 + */ +public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { + + private final int iTopK; + private final int searchWidth; + + /** + * Initializes {@link GPUKnnFloatVectorQuery} + * + * @param field the vector field name + * @param target the vector target query + * @param k the topK value + * @param filter instance of the Query + * @param iTopK the iTopK value + * @param searchWidth the search width + */ + public GPUKnnFloatVectorQuery( + String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { + super(field, target, k, filter); + this.iTopK = iTopK; + this.searchWidth = searchWidth; + } + + @Override + protected TopDocs approximateSearch( + LeafReaderContext context, + Bits acceptDocs, + int visitedLimit, + KnnCollectorManager knnCollectorManager) + throws IOException { + + GPUPerLeafCuVSKnnCollector results = + new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth); + + LeafReader reader = context.reader(); + reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs); + return results.topDocs(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java new file mode 100644 index 0000000000..ede8d2555b --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import org.apache.lucene.search.TopKnnCollector; + +/** + * KnnCollector for cuVS used for search on the GPU. + * + * @since 25.10 + */ +class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { + + private int iTopK; + private int searchWidth; + + /** + * Initializes {@link GPUPerLeafCuVSKnnCollector} + * + * @param topK the topK value + * @param iTopK the iTopK value + * @param searchWidth the search width + */ + public GPUPerLeafCuVSKnnCollector(int topK, int visitLimit, int iTopK, int searchWidth) { + super(topK, visitLimit); + this.iTopK = iTopK > topK ? iTopK : topK; + this.searchWidth = searchWidth; + } + + public int getiTopK() { + return iTopK; + } + + public int getSearchWidth() { + return searchWidth; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java new file mode 100644 index 0000000000..a39480f0d6 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java @@ -0,0 +1,422 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSIvfPqParams; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; +import java.util.Objects; +import java.util.function.Supplier; + +public class GPUSearchParams { + + public static enum Strategy { + /* + * This strategy allows for automatic selection of the underlining CAGRA build algorithm. + * With this strategy we use NN_DESCENT for data set less then 5M vectors else we use IVF_PQ. + * Indexing parameters, especially for IVF_PQ, are heuristically identified automatically. + * + * This is the default and the recommended strategy. + */ + HEURISTIC, + /* + * This is an option when the end-user would want to use custom parameter values. + * This strategy should only be used under expert guidance. + */ + CUSTOM + } + + /* + * TODO: Update boundaries for all parameters when a consensus is reached. + * Issue: https://github.com/rapidsai/cuvs-lucene/issues/99 + */ + public static final int MIN_WRITER_THREADS = 1; + public static final int MAX_WRITER_THREADS = 512; + public static final int MIN_INT_GRAPH_DEG = 2; + public static final int MAX_INT_GRAPH_DEG = 512; + public static final int MIN_GRAPH_DEG = 1; + public static final int MAX_GRAPH_DEG = 512; + public static final int MIN_NN_DESCENT_NUM_ITERATIONS = 1; + public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; + + public static final int DEFAULT_INT_GRAPH_DEGREE = 128; + public static final int DEFAULT_GRAPH_DEGREE = 64; + public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = + CagraGraphBuildAlgo.NN_DESCENT; + public static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + public static final int DEFAULT_WRITER_THREADS = 32; + public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; + public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; + public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; + + public static final Supplier DEFAULT_IVF_PQ_PARAMS = + () -> { + return new CuVSIvfPqParams.Builder().build(); + }; + + private final int writerThreads; + private final int intermediateGraphDegree; + private final int graphdegree; + private final CagraGraphBuildAlgo cagraGraphBuildAlgo; + private final IndexType indexType; + private final CuVSIvfPqParams cuVSIvfPqParams; + private final Strategy strategy; + private final CuvsDistanceType cuvsDistanceType; + private final int nnDescentNumIterations; + + /** + * Constructs an instance of {@link GPUSearchParams} with specific parameter values. + * + * @param writerThreads Number of cuVS writer threads to use. + * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. + * @param graphdegree The graph degree to use while building the CAGRA index. + * @param cagraGraphBuildAlgo The CAGRA build algorithm to use. + * @param indexType The type of index to build - CAGRA, BRUTEFORCE, or both. + * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. + * @param strategy either HEURISTIC [Default] that automatically chooses build algorithm and its parameters based on data set size or CUSTOM that uses the parameters passed though this class. + * @param heuristicType the heuristic type. The default option is SAME_GRAPH_FOOTPRINT. + * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. + * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. + */ + private GPUSearchParams( + int writerThreads, + int intermediateGraphDegree, + int graphdegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, + IndexType indexType, + CuVSIvfPqParams cuVSIvfPqParams, + Strategy strategy, + CuvsDistanceType cuvsDistanceType, + int nnDescentNumIterations) { + super(); + this.writerThreads = writerThreads; + this.intermediateGraphDegree = intermediateGraphDegree; + this.graphdegree = graphdegree; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + this.indexType = indexType; + this.cuVSIvfPqParams = cuVSIvfPqParams; + this.strategy = strategy; + this.cuvsDistanceType = cuvsDistanceType; + this.nnDescentNumIterations = nnDescentNumIterations; + } + + /** + * Get the cuVS writer threads parameter + * + * @return cuVS writer threads parameter + */ + public int getWriterThreads() { + return writerThreads; + } + + /** + * Get the intermediate graph degree + * + * @return the graph degree parameter + */ + public int getIntermediateGraphDegree() { + return intermediateGraphDegree; + } + + /** + * Get the graph degree + * + * @return the graph degree parameter + */ + public int getGraphdegree() { + return graphdegree; + } + + /** + * Get the CAGRA build algorithm parameter value + * + * @return the CAGRA build algorithm parameter value + */ + public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { + return cagraGraphBuildAlgo; + } + + /** + * Get the index type parameter + * + * @return the index type parameter + */ + public IndexType getIndexType() { + return indexType; + } + + /** + * Get the instance of CuVSIvfPqParams + * + * @return an instance of CuVSIvfPqParams + */ + public CuVSIvfPqParams getCuVSIvfPqParams() { + return cuVSIvfPqParams; + } + + /** + * Get the chosen strategy: + * + * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set + * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally + * + * + * @return get the chosen {@link Strategy} + */ + public Strategy getStrategy() { + return strategy; + } + + /** + * Get the cuvs distance type + * + * @return the distance type + */ + public CuvsDistanceType getCuvsDistanceType() { + return cuvsDistanceType; + } + + /** + * get the number of Iterations to run if building with NN_DESCENT + * + * @return the number of iterations for NN_DESCENT + */ + public int getnNDescentNumIterations() { + return nnDescentNumIterations; + } + + @Override + public String toString() { + return "GPUSearchParams [writerThreads=" + + writerThreads + + ", intermediateGraphDegree=" + + intermediateGraphDegree + + ", graphdegree=" + + graphdegree + + ", cagraGraphBuildAlgo=" + + cagraGraphBuildAlgo + + ", indexType=" + + indexType + + ", cuVSIvfPqParams=" + + cuVSIvfPqParams + + ", strategy=" + + strategy + + ", cuvsDistanceType=" + + cuvsDistanceType + + ", nnDescentNumIterations=" + + nnDescentNumIterations + + "]"; + } + + /** + * Builder class for creating an instance of {@link GPUSearchParams} + */ + public static class Builder { + + private int writerThreads = DEFAULT_WRITER_THREADS; + private int intermediateGraphDegree = DEFAULT_INT_GRAPH_DEGREE; + private int graphdegree = DEFAULT_GRAPH_DEGREE; + private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO; + private IndexType indexType = DEFAULT_INDEX_TYPE; + private CuVSIvfPqParams cuVSIvfPqParams = null; + private Strategy strategy = DEFAULT_STRATEGY; + private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE; + private int nnDescentNumIterations = DEFAULT_NN_DESCENT_NUM_ITERATIONS; + + /** + * Set the number of cuVS writer threads while building the index + * Valid range - Minimum: {@value MIN_WRITER_THREADS}, Maximum: {@value MAX_WRITER_THREADS} + * Default value - {@value DEFAULT_WRITER_THREADS} + * + * @param writerThreads the number of cuVS writer threads + * @return instance of {@link Builder} + */ + public Builder withWriterThreads(int writerThreads) { + this.writerThreads = writerThreads; + return this; + } + + /** + * Set the intermediate graph degree to use while building CAGRA index + * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} + * Default value - {@value DEFAULT_INT_GRAPH_DEGREE} + * + * @param intermediateGraphDegree the intermediate graph degree parameter + * @return instance of {@link Builder} + */ + public Builder withIntermediateGraphDegree(int intermediateGraphDegree) { + this.intermediateGraphDegree = intermediateGraphDegree; + return this; + } + + /** + * Set the graph degree to use while building CAGRA index + * Valid range - Minimum: {@value MIN_GRAPH_DEG}, Maximum: {@value MAX_GRAPH_DEG} + * Default value - {@value DEFAULT_GRAPH_DEGREE} + * + * @param graphDegree the graph degree parameter + * @return instance of {@link Builder} + */ + public Builder withGraphDegree(int graphDegree) { + this.graphdegree = graphDegree; + return this; + } + + /** + * Set the CAGRA build algorithm. + * Cannot be null, defaults to NN_DESCENT + * + * @param cagraGraphBuildAlgo the CAGRA build algorithm to use + * @return instance of {@link Builder} + */ + public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo) { + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + return this; + } + + /** + * Set the type of index to build - CAGRA, BRUTEFORCE, or both. + * Cannot be null, defaults to CAGRA + * + * @param indexType the type of index to build + * @return instance of {@link Builder} + */ + public Builder withIndexType(IndexType indexType) { + this.indexType = indexType; + return this; + } + + /** + * Set the instance of {@link CuVSIvfPqParams} + * + * @param cuVSIvfPqParams + * @return instance of {@link Builder} + */ + public Builder withCuVSIvfPqParams(CuVSIvfPqParams cuVSIvfPqParams) { + this.cuVSIvfPqParams = cuVSIvfPqParams; + return this; + } + + /** + * Set the chosen strategy: + * + * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set + * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally + * + * Valid options - HEURISTIC, CUSTOM + * Default value - HEURISTIC + * + * @param strategy, the strategy to choose + * @return instance of {@link Builder} + */ + public Builder withStrategy(Strategy strategy) { + this.strategy = strategy; + return this; + } + + /** + * Set the CuvsDistanceType + * + * @param cuvsDistanceType the CuvsDistanceType to set + * @return instance of {@link Builder} + */ + public Builder withCuvsDistanceType(CuvsDistanceType cuvsDistanceType) { + this.cuvsDistanceType = cuvsDistanceType; + return this; + } + + /** + * Set the number of Iterations to run if building with NN_DESCENT + * + * Valid range - Minimum: {@value MIN_NN_DESCENT_NUM_ITERATIONS}, Maximum: {@value MAX_NN_DESCENT_NUM_ITERATIONS} + * Default value - {@value DEFAULT_NN_DESCENT_NUM_ITERATIONS} + * + * @param nnDescentNumIterations number of merge workers to set + * @return instance of {@link Builder} + */ + public Builder withNNDescentNumIterations(int nnDescentNumIterations) { + this.nnDescentNumIterations = nnDescentNumIterations; + return this; + } + + /** + * Validates the input parameters. + * + * @throws IllegalArgumentException + */ + private void validate() throws IllegalArgumentException { + if (writerThreads < MIN_WRITER_THREADS || writerThreads > MAX_WRITER_THREADS) { + throw new IllegalArgumentException( + "writerThreads not in valid range. Valid range: [" + + MIN_WRITER_THREADS + + ", " + + MAX_WRITER_THREADS + + "]"); + } + if (intermediateGraphDegree < MIN_INT_GRAPH_DEG + || intermediateGraphDegree > MAX_INT_GRAPH_DEG) { + throw new IllegalArgumentException( + "intermediateGraphDegree not in valid range. Valid range: [" + + MIN_INT_GRAPH_DEG + + ", " + + MAX_INT_GRAPH_DEG + + "]"); + } + if (graphdegree < MIN_GRAPH_DEG || graphdegree > MAX_GRAPH_DEG) { + throw new IllegalArgumentException( + "graphdegree not in valid range. Valid range: [" + + MIN_GRAPH_DEG + + ", " + + MAX_GRAPH_DEG + + "]"); + } + if (Objects.isNull(cagraGraphBuildAlgo)) { + throw new IllegalArgumentException("cagraGraphBuildAlgo cannot be null."); + } + if (Objects.isNull(indexType)) { + throw new IllegalArgumentException("indexType cannot be null."); + } + if (Objects.isNull(strategy)) { + throw new IllegalArgumentException("strategy cannot be null."); + } + if (Objects.isNull(cuvsDistanceType)) { + throw new IllegalArgumentException("cuvsDistanceType cannot be null."); + } + if (nnDescentNumIterations < MIN_NN_DESCENT_NUM_ITERATIONS + || nnDescentNumIterations > MAX_NN_DESCENT_NUM_ITERATIONS) { + throw new IllegalArgumentException( + "nnDescentNumIterations not in valid range. Valid range: [" + + MIN_NN_DESCENT_NUM_ITERATIONS + + ", " + + MAX_NN_DESCENT_NUM_ITERATIONS + + "]"); + } + } + + /** + * Creates and returns an instance of {@link GPUSearchParams} + * + * @return instance of {@link GPUSearchParams} + */ + public GPUSearchParams build() { + if (Objects.isNull(cuVSIvfPqParams)) { + cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS.get(); + } + validate(); + return new GPUSearchParams( + writerThreads, + intermediateGraphDegree, + graphdegree, + cagraGraphBuildAlgo, + indexType, + cuVSIvfPqParams, + strategy, + cuvsDistanceType, + nnDescentNumIterations); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java new file mode 100644 index 0000000000..0749b70a28 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java @@ -0,0 +1,63 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.io.InputStream; +import org.apache.lucene.store.IndexInput; + +/** + * InputStream for reading from an IndexInput. + * + * @since 25.10 + */ +final class IndexInputInputStream extends InputStream { + + final IndexInput in; + long pos = 0; + final long limit; + + /** + * Initializes the {@link IndexInputInputStream} + * + * @param in instance of IndexInput + */ + IndexInputInputStream(IndexInput in) { + this.in = in; + this.limit = in.length(); + } + + /** + * Reads the next byte of data from the input stream. + */ + @Override + public int read() throws IOException { + if (pos >= limit) { + return -1; + } + pos++; + return in.readByte(); + } + + /** + * Reads up to len bytes of data from the input stream into an array of bytes. + */ + @Override + public int read(byte[] b, int off, int len) throws IOException { + if (len <= 0) { + return 0; + } + if (pos >= limit) { + return -1; + } + long avail = limit - pos; + if (len > avail) { + len = (int) avail; + } + in.readBytes(b, off, len); + pos += len; + return len; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java new file mode 100644 index 0000000000..c8c9f796fa --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.io.OutputStream; +import org.apache.lucene.store.IndexOutput; + +/** + * OutputStream for writing into an IndexOutput + * + * @since 25.10 + */ +final class IndexOutputOutputStream extends OutputStream { + + static final int DEFAULT_BUFFER_SIZE = 8192; + + final IndexOutput out; + final int bufferSize; + final byte[] buffer; + int pos; + + /** + * Initializes the {@link IndexOutputOutputStream}. + * + * @param out instance of IndexOutput + */ + IndexOutputOutputStream(IndexOutput out) { + this(out, DEFAULT_BUFFER_SIZE); + } + + /** + * Initializes the {@link IndexOutputOutputStream}. + * + * @param out instance of IndexOutput + * @param bufferSize the size of buffer to use + */ + IndexOutputOutputStream(IndexOutput out, int bufferSize) { + this.out = out; + this.bufferSize = bufferSize; + this.buffer = new byte[bufferSize]; + } + + /** + * Writes the specified byte to this output stream. + */ + @Override + public void write(int b) throws IOException { + buffer[pos] = (byte) b; + pos++; + if (pos == bufferSize) { + flush(); + } + } + + /** + * Writes len bytes from the specified byte array starting at offset off to this output stream. + */ + @Override + public void write(byte[] b, int offset, int length) throws IOException { + if (pos != 0) { + flush(); + } + out.writeBytes(b, offset, length); + } + + /** + * Flushes this output stream and forces any buffered output bytes to be written out. + */ + @Override + public void flush() throws IOException { + out.writeBytes(buffer, 0, pos); + pos = 0; + } + + /** + * Closes this output stream and releases any system resources associated with this stream. + */ + @Override + public void close() throws IOException { + this.flush(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java new file mode 100644 index 0000000000..e9f4d6fead --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.LibraryException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; + +/** + * A codec that enables GPU-based accelerated HNSW capability and can be used + * to accelerated indexing using GPUs and search using CPUs. Fallbacks to CPU + * based indexing when used on a machine without a GPU and/or cuVS. + * + * @since 25.10 + */ +public class Lucene101AcceleratedHNSWCodec extends FilterCodec { + + private static final Logger log = Logger.getLogger(Lucene101AcceleratedHNSWCodec.class.getName()); + private static final String NAME = "Lucene101AcceleratedHNSWCodec"; + private KnnVectorsFormat format; + + /** + * Default constructor for {@link Lucene101AcceleratedHNSWCodec}. + * + * @throws Exception + */ + public Lucene101AcceleratedHNSWCodec() throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + } + + /** + * Constructor for {@link Lucene101AcceleratedHNSWCodec}. + * + * @param name the codec's name + * @param delegate the delegate codec to filter + */ + public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { + super(name, delegate); + initializeFormatDefaultValues(); + } + + /** + * Constructor for {@link Lucene101AcceleratedHNSWCodec}. + * + * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} + * @throws Exception exception + */ + public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) + throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + initializeFormat(acceleratedHNSWParams); + } + + /** + * Initialize an instance of {@link Lucene99AcceleratedHNSWVectorsFormat} + * with an instance of {@link AcceleratedHNSWParams} with default parameter values. + */ + private void initializeFormatDefaultValues() { + initializeFormat(new AcceleratedHNSWParams.Builder().build()); + } + + /** + * Initialize an instance of {@link Lucene99AcceleratedHNSWVectorsFormat}. + * + * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} to use + */ + private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { + try { + format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams); + setKnnFormat(format); + } catch (LibraryException ex) { + log.log( + Level.SEVERE, + "Couldn't load native library, possible classloader issue. " + ex.getMessage()); + } + } + + /** + * Get the configured {@link KnnVectorsFormat}. + * + * @return the instance of the {@link KnnVectorsFormat} + */ + @Override + public KnnVectorsFormat knnVectorsFormat() { + return format; + } + + /** + * Set the {@link KnnVectorsFormat}. + * + * @param format the {@link KnnVectorsFormat} to set + */ + public void setKnnFormat(KnnVectorsFormat format) { + this.format = format; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java new file mode 100644 index 0000000000..61c9b41c3a --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -0,0 +1,119 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import com.nvidia.cuvs.LibraryException; +import java.io.IOException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.TaskExecutor; + +/** + * cuVS based KnnVectorsFormat for indexing on GPU and searching on the CPU. + * + * @since 25.10 + */ +public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { + + private static final Logger log = + Logger.getLogger(Lucene99AcceleratedHNSWVectorsFormat.class.getName()); + private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; + private static final int MAX_DIMENSIONS = 4096; + private final AcceleratedHNSWParams acceleratedHNSWParams; + + static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; + static final String HNSW_META_CODEC_EXT = "vem"; + static final String HNSW_INDEX_CODEC_NAME = "Lucene99HnswVectorsFormatIndex"; + static final String HNSW_INDEX_EXT = "vex"; + static final LuceneProvider LUCENE_PROVIDER; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + FLAT_VECTORS_FORMAT = + LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} with an instance + * of {@link AcceleratedHNSWParams} with default parameter values. + * + * @throws LibraryException if the native library fails to load + */ + public Lucene99AcceleratedHNSWVectorsFormat() { + this(new AcceleratedHNSWParams.Builder().build()); + } + + /** + * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} with an instance + * of {@link AcceleratedHNSWParams}. + * + * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} + */ + public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNSWParams) { + super("Lucene99AcceleratedHNSWVectorsFormat"); + this.acceleratedHNSWParams = acceleratedHNSWParams; + } + + /** + * Returns a KnnVectorsWriter to write the vectors to the index. + */ + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + if (isSupported()) { + log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); + return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); + } else { + log.log( + Level.WARNING, + "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); + try { + return LUCENE_PROVIDER.getLuceneHnswVectorsWriterInstance( + state, + acceleratedHNSWParams.getMaxConn(), + acceleratedHNSWParams.getBeamWidth(), + flatWriter, + acceleratedHNSWParams.getNumMergeWorkers(), + new TaskExecutor(acceleratedHNSWParams.getMergeExec())); + } catch (Exception e) { + throw new RuntimeException(e.getMessage()); + } + } + } + + /** + * Returns a KnnVectorsReader to read the vectors from the index. + */ + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + try { + return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance( + state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + } catch (Exception e) { + throw new RuntimeException(e.getMessage()); + } + } + + /** + * Returns the maximum number of vector dimensions supported by this codec for the given field name. + */ + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMENSIONS; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java new file mode 100644 index 0000000000..0e558669e7 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -0,0 +1,342 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; + +/** + * This class extends upon the KnnVectorsWriter to + * enable the creation of GPU-based accelerated HNSW based vector search. + * + * @since 25.10 + */ +public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(Lucene99AcceleratedHNSWVectorsWriter.class); + private static final String COMPONENT = "Lucene99AcceleratedHNSWVectorsWriter"; + private static final LuceneProvider LUCENE_PROVIDER; + private static final Integer VERSION_CURRENT; + + private final AcceleratedHNSWParams acceleratedHNSWParams; + private final FlatVectorsWriter flatVectorsWriter; + private final List fields = new ArrayList<>(); + private final InfoStream infoStream; + private IndexOutput hnswMeta = null; + private IndexOutput hnswVectorIndex = null; + private String vemFileName; + private String vexFileName; + private boolean finished; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes {@link Lucene99AcceleratedHNSWVectorsWriter} + * + * @param state instance of the {@link org.apache.lucene.index.SegmentWriteState} + * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} + * @param flatVectorsWriter instance of the {@link org.apache.lucene.codecs.hnsw.FlatVectorsWriter} + * @throws IOException IOException + */ + public Lucene99AcceleratedHNSWVectorsWriter( + SegmentWriteState state, + AcceleratedHNSWParams acceleratedHNSWParams, + FlatVectorsWriter flatVectorsWriter) + throws IOException { + super(); + this.flatVectorsWriter = flatVectorsWriter; + this.infoStream = state.infoStream; + this.acceleratedHNSWParams = acceleratedHNSWParams; + vemFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); + vexFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); + boolean success = false; + try { + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + CodecUtil.writeIndexHeader( + hnswMeta, + HNSW_META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + HNSW_INDEX_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; + printInfoStream(infoStream, COMPONENT, "Lucene99AcceleratedHNSWVectorsWriter is initialized"); + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Add new field for indexing. + */ + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + var encoding = fieldInfo.getVectorEncoding(); + if (encoding != FLOAT32) { + throw new IllegalArgumentException("Expected float32, got:" + encoding); + } + var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); + var cuvsFieldWriter = new FieldWriter(QuantizationType.NONE, fieldInfo, writer); + fields.add(cuvsFieldWriter); + return writer; + } + + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param vectors vectors to index + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + if (vectors.size() < 2) { + writeSingleVectorGraph(fieldInfo, vectors); + return; + } + try { + CuVSMatrix dataset = + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + + CagraIndexParams params = + CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + + CagraIndex cagraIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + int size = (int) dataset.size(); + int dimensions = fieldInfo.getVectorDimension(); + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, + size, + dimensions, + adjacencyListMatrix, + vectors, + acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), + params, + QuantizationType.NONE); + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + cagraIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Build the indexes and writes it to the disk. + */ + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + flatVectorsWriter.flush(maxDoc, sortMap); + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } + } + + /** + * Builds the index and writes it to the disk. + * + * @param fieldData + * @throws IOException + */ + private void writeField(FieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getFloatVectors()); + } + + /** + * Builds the index and writes it to the disk. + * + * @param fieldData instance of GPUFieldWriter + * @param sortMap instance of the DocMap + * @throws IOException + */ + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; + mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); + List sortedVectors = new ArrayList(); + List floatVectors = fieldData.getFloatVectors(); + for (int i = 0; i < floatVectors.size(); i++) { + sortedVectors.add(floatVectors.get(new2OldOrd[i])); + } + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); + } + + /** + * Builds and writes a single vector graph. + * + * @param fieldInfo instance of FieldInfo + * @param vectors the list of float vectors + * @throws IOException I/O Exceptions + */ + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { + try { + int size = 1; + int dimensions = fieldInfo.getVectorDimension(); + GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Create combined data set for the merged segment and call writeFieldInternal. + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Write field for merging. + */ + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + flatVectorsWriter.mergeOneField(fieldInfo, mergeState); + vectorBasedMerge(fieldInfo, mergeState); + } + + /** + * Called once at the end before close. + */ + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + flatVectorsWriter.finish(); + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } + + /** + * Closes the resources. + */ + @Override + public void close() throws IOException { + printInfoStream(infoStream, COMPONENT, "Closing resources"); + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); + closeCuVSResourcesInstance(); + } + + /** + * Returns the memory usage of this object in bytes. + */ + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (var field : fields) { + total += field.ramBytesUsed(); + } + return total; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java new file mode 100644 index 0000000000..0b1653bc14 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java @@ -0,0 +1,61 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.LibraryException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; + +/** + * CuVS based codec for GPU based vector search + * + * @since 26.02 + */ +public class LuceneAcceleratedHNSWBinaryQuantizedCodec extends FilterCodec { + + private static final Logger log = + Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedCodec.class.getName()); + private static final String NAME = "Lucene101AcceleratedHNSWBinaryQuantizedCodec"; + + private KnnVectorsFormat format; + + public LuceneAcceleratedHNSWBinaryQuantizedCodec() throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + } + + public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) { + super(name, delegate); + initializeFormat(new AcceleratedHNSWParams.Builder().build()); + } + + public LuceneAcceleratedHNSWBinaryQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) + throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + initializeFormat(acceleratedHNSWParams); + } + + private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { + try { + format = new LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat(acceleratedHNSWParams); + setKnnFormat(format); + } catch (LibraryException ex) { + log.log( + Level.SEVERE, + "Couldn't load native library, possible classloader issue. " + ex.getMessage()); + } + } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return format; + } + + public void setKnnFormat(KnnVectorsFormat format) { + this.format = format; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java new file mode 100644 index 0000000000..d1810bc540 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -0,0 +1,117 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import com.nvidia.cuvs.LibraryException; +import java.io.IOException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** + * cuVS based Binary Quantized KnnVectorsFormat for indexing on GPU and searching on the CPU. + * + * @since 26.02 + */ +public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVectorsFormat { + + private static final Logger log = + Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class.getName()); + private static final LuceneProvider LUCENE102_PROVIDER; + private static final LuceneProvider LUCENE99_PROVIDER; + private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; + private static final int MAX_DIMENSIONS = 4096; + + private final AcceleratedHNSWParams acceleratedHNSWParams; + + static { + try { + LUCENE99_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE102_PROVIDER = LuceneProvider.getInstance("102"); + FLAT_VECTORS_FORMAT = + LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes {@link LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat} with default values. + * + * @throws LibraryException if the native library fails to load + */ + public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat() { + this(new AcceleratedHNSWParams.Builder().build()); + } + + /** + * Initializes {@link LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat} with the given threads, graph degree, etc. + * + * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} + */ + public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( + AcceleratedHNSWParams acceleratedHNSWParams) { + super("Lucene99AcceleratedHNSWBinaryQuantizedVectorsFormat"); + this.acceleratedHNSWParams = acceleratedHNSWParams; + } + + /** + * Returns a KnnVectorsWriter to write the binary quantized vectors to the index. + */ + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + if (isSupported()) { + log.log( + Level.FINE, + "cuVS is supported so using the Lucene99AcceleratedHNSWBinaryQuantizedVectorsWriter"); + return new LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( + state, acceleratedHNSWParams, flatWriter); + } else { + try { + // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat format + log.log( + Level.WARNING, + "GPU based indexing not supported, falling back to using the" + + " Lucene102HnswBinaryQuantizedVectorsFormat"); + KnnVectorsFormat fallbackFormat = + LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance( + acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); + return fallbackFormat.fieldsWriter(state); + } catch (Exception e) { + throw new RuntimeException(e.getMessage()); + } + } + } + + /** + * Returns a KnnVectorsReader to read the binary quantized vectors from the index. + */ + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + try { + return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance( + state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + } catch (Exception e) { + throw new RuntimeException(e.getMessage()); + } + } + + /** + * Returns the maximum number of vector dimensions supported by this codec for the given field name. + */ + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMENSIONS; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java new file mode 100644 index 0000000000..e00b2d07b0 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -0,0 +1,374 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; + +/** + * This class extends upon the KnnVectorsWriter to enable the creation of GPU-based accelerated + * vector search indexes. + * + * @since 26.02 + */ +public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.class); + private static final String COMPONENT = "Lucene99AcceleratedHNSWBinaryQuantizedVectorsWriter"; + + private final FlatVectorsWriter flatVectorsWriter; + private final List fields = new ArrayList<>(); + private final InfoStream infoStream; + private final AcceleratedHNSWParams acceleratedHNSWParams; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; + private boolean finished; + private String vemFileName; + private String vexFileName; + + /** + * Initializes {@link LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter} + * + * @param state instance of the {@link org.apache.lucene.index.SegmentWriteState} + * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} + * @param flatVectorsWriter instance of the {@link org.apache.lucene.codecs.hnsw.FlatVectorsWriter} + * @throws IOException IOException + */ + public LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( + SegmentWriteState state, + AcceleratedHNSWParams acceleratedHNSWParams, + FlatVectorsWriter flatVectorsWriter) + throws IOException { + super(); + this.acceleratedHNSWParams = acceleratedHNSWParams; + this.flatVectorsWriter = flatVectorsWriter; + this.infoStream = state.infoStream; + + vemFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); + + vexFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); + + boolean success = false; + try { + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + + CodecUtil.writeIndexHeader( + hnswMeta, + HNSW_META_CODEC_NAME, + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + HNSW_INDEX_CODEC_NAME, + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + + success = true; + printInfoStream( + infoStream, COMPONENT, "Lucene99AcceleratedHNSWBinaryQuantizedVectorsWriter opened"); + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Add new field for indexing. + */ + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + VectorEncoding encoding = fieldInfo.getVectorEncoding(); + if (encoding != FLOAT32) { + throw new IllegalArgumentException("expected float32, got:" + encoding); + } + var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); + var cuvsFieldWriter = new FieldWriter(QuantizationType.BINARY, fieldInfo, writer); + fields.add(cuvsFieldWriter); + return writer; + } + + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index for binary quantized vectors. + * Binary quantized vectors are stored as packed bits (1 bit per dimension, 8 dimensions per byte). + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param vectors binary quantized vectors (packed bits as bytes) + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + + try { + int dimensions = fieldInfo.getVectorDimension(); + int bytesPerVector = (dimensions + 7) / 8; + + CuVSMatrix dataset = + Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance()); + + if (dataset.size() < 2) { + writeSingleVectorGraph(fieldInfo, vectors); + return; + } + + CagraIndexParams params = + CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + + CagraIndex cagraIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + int size = (int) dataset.size(); + + // Create multi-layer HNSW graph from CAGRA + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, + size, + dimensions, + adjacencyListMatrix, + vectors, + acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), + params, + QuantizationType.BINARY); + + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + + cagraIndex.close(); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Build the indexes and writes it to the disk. + */ + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + flatVectorsWriter.flush(maxDoc, sortMap); + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } + } + + /** + * Builds the index and writes it to the disk. + * + * @param fieldData + * @throws IOException + */ + private void writeField(FieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getByteVectors()); + } + + /** + * Builds the index and writes it to the disk. + * + * @param fieldData instance of BinaryQuantizedGPUFieldWriter + * @param sortMap instance of the DocMap + * @throws IOException + */ + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord + mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); + + List sortedVectors = new ArrayList(); + List vectors = fieldData.getByteVectors(); + for (int i = 0; i < vectors.size(); i++) { + sortedVectors.add(vectors.get(new2OldOrd[i])); + } + + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); + } + + /** + * Builds and writes a single vector graph. + * + * @param fieldInfo instance of FieldInfo + * @param vectors the list of binary quantized vectors + * @throws IOException I/O Exceptions + */ + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { + // Workaround for CAGRA not supporting single vector indexes + try { + int size = 1; + int dimensions = fieldInfo.getVectorDimension(); + + // Create a dummy HNSW graph for a single vector + GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); + + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Write field for merging. + */ + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + flatVectorsWriter.mergeOneField(fieldInfo, mergeState); + vectorBasedMerge(fieldInfo, mergeState); + } + + /** + * Fallback method that rebuilds indexes from merged vectors. + * Used when native CAGRA merge() is not possible. + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + FloatVectorValues mergedVectorValues = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + + if (mergedVectorValues != null) { + List floatVectors = new ArrayList<>(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + floatVectors.add(mergedVectorValues.vectorValue(iter.index()).clone()); + } + writeFieldInternal(fieldInfo, quantizeFloatVectorsToBinary(floatVectors)); + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Called once at the end before close. + */ + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + flatVectorsWriter.finish(); + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } + + /** + * Closes the resources. + */ + @Override + public void close() throws IOException { + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); + closeCuVSResourcesInstance(); + } + + /** + * Returns the memory usage of this object in bytes. + */ + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (var field : fields) { + total += field.ramBytesUsed(); + } + return total; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java new file mode 100644 index 0000000000..0705ed0a55 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.LibraryException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; + +/** + * CuVS based codec for GPU based vector search + * + * @since 26.02 + */ +public class LuceneAcceleratedHNSWScalarQuantizedCodec extends FilterCodec { + + private static final Logger log = + Logger.getLogger(LuceneAcceleratedHNSWScalarQuantizedCodec.class.getName()); + private static final String NAME = "Lucene101AcceleratedHNSWScalarQuantizedCodec"; + + private KnnVectorsFormat format; + + public LuceneAcceleratedHNSWScalarQuantizedCodec() throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + } + + public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) { + super(name, delegate); + initializeFormatDefaultValues(); + } + + public LuceneAcceleratedHNSWScalarQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) + throws Exception { + this(NAME, LuceneProvider.getCodec("101")); + initializeFormat(acceleratedHNSWParams); + } + + private void initializeFormatDefaultValues() { + initializeFormat(new AcceleratedHNSWParams.Builder().build()); + } + + private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { + try { + format = new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(acceleratedHNSWParams); + setKnnFormat(format); + } catch (LibraryException ex) { + log.log( + Level.SEVERE, + "Couldn't load native library, possible classloader issue. " + ex.getMessage()); + } + } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return format; + } + + public void setKnnFormat(KnnVectorsFormat format) { + this.format = format; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java new file mode 100644 index 0000000000..8d599a54ef --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import com.nvidia.cuvs.LibraryException; +import java.io.IOException; +import java.util.logging.Logger; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** + * cuVS based Scalar Quantized KnnVectorsFormat for indexing on GPU and searching on the CPU. + * + * @since 26.02 + */ +public class LuceneAcceleratedHNSWScalarQuantizedVectorsFormat extends KnnVectorsFormat { + + private static final Logger log = + Logger.getLogger(LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.class.getName()); + private static final LuceneProvider LUCENE_PROVIDER; + private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; + private static final int MAX_DIMENSIONS = 4096; + + private final AcceleratedHNSWParams acceleratedHNSWParams; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneScalarQuantizedVectorsFormatInstance(); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes {@link LuceneAcceleratedHNSWScalarQuantizedVectorsFormat} with default values. + * + * @throws LibraryException if the native library fails to load + */ + public LuceneAcceleratedHNSWScalarQuantizedVectorsFormat() { + this(new AcceleratedHNSWParams.Builder().build()); + } + + /** + * Initializes {@link LuceneAcceleratedHNSWScalarQuantizedVectorsFormat} with the given threads, graph degree, etc. + * + * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} + */ + public LuceneAcceleratedHNSWScalarQuantizedVectorsFormat( + AcceleratedHNSWParams acceleratedHNSWParams) { + super("Lucene99AcceleratedHNSWScalarQuantizedVectorsFormat"); + this.acceleratedHNSWParams = acceleratedHNSWParams; + } + + /** + * Returns a KnnVectorsWriter to write the scalar quantized vectors to the index. + */ + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + if (isSupported()) { + log.info("cuVS is supported so using the Lucene99AcceleratedHNSWQuantizedVectorsWriter"); + return new LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( + state, acceleratedHNSWParams, flatWriter); + } else { + try { + // Fallback to Lucene's Lucene99HnswScalarQuantizedVectorsFormat + log.warning( + "GPU based indexing not supported, falling back to using the" + + " Lucene99HnswScalarQuantizedVectorsFormat"); + KnnVectorsFormat fallbackFormat = + LUCENE_PROVIDER.getLuceneHnswScalarQuantizedVectorsFormatInstance( + acceleratedHNSWParams.getBeamWidth(), acceleratedHNSWParams.getMaxConn()); + return fallbackFormat.fieldsWriter(state); + } catch (Exception e) { + throw new RuntimeException(e.getMessage()); + } + } + } + + /** + * Returns a KnnVectorsReader to read the scalar quantized vectors from the index. + */ + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + try { + return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance( + state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + } catch (Exception e) { + throw new RuntimeException(e.getMessage()); + } + } + + /** + * Returns the maximum number of vector dimensions supported by this Codec for the given field name. + */ + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMENSIONS; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java new file mode 100644 index 0000000000..a542cf8967 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -0,0 +1,398 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; + +/** + * This class extends upon the KnnVectorsWriter to enable the creation of GPU-based accelerated + * vector search indexes. + * + * @since 26.02 + */ +public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.class); + private static final String COMPONENT = "Lucene99AcceleratedHNSWQuantizedVectorsWriter"; + private static final LuceneProvider LUCENE_PROVIDER; + private static final Integer VERSION_CURRENT; + + private final FlatVectorsWriter flatVectorsWriter; + private final List fields = new ArrayList<>(); + private final InfoStream infoStream; + private final AcceleratedHNSWParams acceleratedHNSWParams; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; + private boolean finished; + private String vemFileName; + private String vexFileName; + + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + + /** + * Initializes {@link LuceneAcceleratedHNSWScalarQuantizedVectorsWriter} + * + * @param state instance of the {@link org.apache.lucene.index.SegmentWriteState} + * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} + * @param flatVectorsWriter instance of the {@link org.apache.lucene.codecs.hnsw.FlatVectorsWriter} + * @throws IOException IOException + */ + public LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( + SegmentWriteState state, + AcceleratedHNSWParams acceleratedHNSWParams, + FlatVectorsWriter flatVectorsWriter) + throws IOException { + super(); + this.acceleratedHNSWParams = acceleratedHNSWParams; + this.flatVectorsWriter = flatVectorsWriter; + this.infoStream = state.infoStream; + + vemFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); + + vexFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); + + boolean success = false; + try { + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + + CodecUtil.writeIndexHeader( + hnswMeta, + HNSW_META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + HNSW_INDEX_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; + printInfoStream( + infoStream, COMPONENT, "Lucene99AcceleratedHNSWQuantizedVectorsWriter opened"); + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Add new field for indexing. + */ + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + VectorEncoding encoding = fieldInfo.getVectorEncoding(); + if (encoding != FLOAT32) { + throw new IllegalArgumentException("expected float32, got:" + encoding); + } + var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); + var cuvsFieldWriter = new FieldWriter(QuantizationType.SCALAR, fieldInfo, writer); + fields.add(cuvsFieldWriter); + return writer; + } + + private static byte signedToUnsignedByte(byte signedByte) { + return (byte) (signedByte & 0xFF); + } + + private static byte[] convertSignedToUnsigned(byte[] signedVector) { + byte[] unsignedVector = new byte[signedVector.length]; + for (int i = 0; i < signedVector.length; i++) { + unsignedVector[i] = signedToUnsignedByte(signedVector[i]); + } + return unsignedVector; + } + + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param vectors quantized vectors + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + + try { + int dimensions = fieldInfo.getVectorDimension(); + + // Convert 7-bit signed bytes to 8-bit unsigned bytes for cuVS compatibility + List unsignedVectors = new ArrayList<>(vectors.size()); + for (Object signedVector : vectors) { + unsignedVectors.add(convertSignedToUnsigned((byte[]) signedVector)); + } + + // Create CuVSMatrix with BYTE data type (unsigned bytes) + CuVSMatrix dataset = + Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); + + if (dataset.size() < 2) { + writeSingleVectorGraph(fieldInfo, unsignedVectors); + return; + } + + CagraIndexParams params = + CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + + CagraIndex cagraIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + + int size = (int) dataset.size(); + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, + size, + dimensions, + adjacencyListMatrix, + unsignedVectors, + acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), + params, + QuantizationType.SCALAR); + + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + + cagraIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Build the indexes and writes it to the disk. + */ + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + flatVectorsWriter.flush(maxDoc, sortMap); + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } + } + + /** + * Builds the index and writes it to the disk. + * + * @param fieldData + * @throws IOException + */ + private void writeField(FieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getByteVectors()); + } + + /** + * Builds the index and writes it to the disk. + * + * @param fieldData instance of ScalarQuantizedGPUFieldWriter + * @param sortMap instance of the DocMap + * @throws IOException + */ + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord + mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); + + List sortedVectors = new ArrayList(); + List byteVectors = fieldData.getByteVectors(); + for (int i = 0; i < byteVectors.size(); i++) { + sortedVectors.add(byteVectors.get(new2OldOrd[i])); + } + + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); + } + + /** + * Builds and writes a single vector graph. + * + * @param fieldInfo instance of FieldInfo + * @param vectors the list of scalar quantized vectors (already converted to unsigned) + * @throws IOException I/O Exceptions + */ + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { + // Workaround for CAGRA not supporting single vector indexes + try { + int size = 1; + int dimensions = fieldInfo.getVectorDimension(); + + // Create a dummy HNSW graph for a single vector + GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); + + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Write field for merging. + */ + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + flatVectorsWriter.mergeOneField(fieldInfo, mergeState); + vectorBasedMerge(fieldInfo, mergeState); + } + + /** + * Fallback method that rebuilds indexes from merged vectors. + * Used when native CAGRA merge() is not possible. + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + FloatVectorValues mergedVectorValues = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + + if (mergedVectorValues != null) { + List floatVectors = new ArrayList<>(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + floatVectors.add(mergedVectorValues.vectorValue(iter.index()).clone()); + } + writeFieldInternal(fieldInfo, quantizeFloatVectorsToScalar(floatVectors)); + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Called once at the end before close. + */ + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + flatVectorsWriter.finish(); + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } + + /** + * Closes the resources. + */ + @Override + public void close() throws IOException { + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); + closeCuVSResourcesInstance(); + } + + /** + * Returns the memory usage of this object in bytes. + */ + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (var field : fields) { + total += field.ramBytesUsed(); + } + return total; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java new file mode 100644 index 0000000000..9caa2a6aa3 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -0,0 +1,303 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.lang.reflect.Constructor; +import java.lang.reflect.InvocationTargetException; +import java.util.List; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.TaskExecutor; + +/** + * Dynamically loads Lucene format, reader, and writer classes with a fallback mechanism. + * + * @since 25.12 + */ +public class LuceneProvider { + + static final Logger log = Logger.getLogger(LuceneProvider.class.getName()); + + private static final String BASE = "org.apache.lucene."; + private static String codecs = "codecs.lucene."; + private static String fallbackCodecs = "backward_codecs.lucene."; + + private static String luceneFlatVectorsFormat = + BASE + codecs + "LuceneFlatVectorsFormat"; + private static String luceneFlatVectorsFormatFallback = + BASE + fallbackCodecs + "LuceneFlatVectorsFormat"; + + private static String luceneHnswVectorsFormat = + BASE + codecs + "LuceneHnswVectorsFormat"; + private static String luceneHnswVectorsFormatFallback = + BASE + fallbackCodecs + "LuceneHnswVectorsFormat"; + + private static String luceneHnswVectorsReader = + BASE + codecs + "LuceneHnswVectorsReader"; + private static String luceneHnswVectorsReaderFallback = + BASE + fallbackCodecs + "LuceneHnswVectorsReader"; + + private static String luceneHnswVectorsWriter = + BASE + codecs + "LuceneHnswVectorsWriter"; + private static String luceneHnswVectorsWriterFallback = + BASE + fallbackCodecs + "LuceneHnswVectorsWriter"; + + private static String luceneBinaryQuantizedVectorsFormat = + BASE + codecs + "LuceneBinaryQuantizedVectorsFormat"; + private static String luceneBinaryQuantizedVectorsFormatFallback = + BASE + fallbackCodecs + "LuceneBinaryQuantizedVectorsFormat"; + + private static String luceneHnswBinaryQuantizedVectorsFormat = + BASE + codecs + "LuceneHnswBinaryQuantizedVectorsFormat"; + private static String luceneHnswBinaryQuantizedVectorsFormatFallback = + BASE + fallbackCodecs + "LuceneHnswBinaryQuantizedVectorsFormat"; + + private static String luceneScalarQuantizedVectorsFormat = + BASE + codecs + "LuceneScalarQuantizedVectorsFormat"; + private static String luceneScalarQuantizedVectorsFormatFallback = + BASE + fallbackCodecs + "LuceneScalarQuantizedVectorsFormat"; + + private static String luceneHnswScalarQuantizedVectorsFormat = + BASE + codecs + "LuceneHnswScalarQuantizedVectorsFormat"; + private static String luceneHnswScalarQuantizedVectorsFormatFallback = + BASE + fallbackCodecs + "LuceneHnswScalarQuantizedVectorsFormat"; + + private static String luceneCodec = BASE + codecs + "LuceneCodec"; + private static String luceneCodecFallback = BASE + fallbackCodecs + "LuceneCodec"; + + private static LuceneProvider instance; + + private static MethodHandles.Lookup lookup = MethodHandles.lookup(); + + private Class flatVectorsFormat; + private Class hnswVectorsFormat; + private Class hnswVectorsReader; + private Class hnswVectorsWriter; + private Class binaryQuantizedVectorsFormat; + private Class hnswBinaryQuantizedVectorsFormat; + private Class scalarQuantizedVectorsFormat; + private Class hnswScalarQuantizedVectorsFormat; + + public static LuceneProvider getInstance(String version) throws ClassNotFoundException { + if (instance == null) { + instance = new LuceneProvider(version); + } + return instance; + } + + private LuceneProvider(String version) throws ClassNotFoundException { + flatVectorsFormat = + loadClass( + setVersion(luceneFlatVectorsFormat, version), + setVersion(luceneFlatVectorsFormatFallback, version)); + hnswVectorsFormat = + loadClass( + setVersion(luceneHnswVectorsFormat, version), + setVersion(luceneHnswVectorsFormatFallback, version)); + hnswVectorsReader = + loadClass( + setVersion(luceneHnswVectorsReader, version), + setVersion(luceneHnswVectorsReaderFallback, version)); + hnswVectorsWriter = + loadClass( + setVersion(luceneHnswVectorsWriter, version), + setVersion(luceneHnswVectorsWriterFallback, version)); + scalarQuantizedVectorsFormat = + loadClass( + setVersion(luceneScalarQuantizedVectorsFormat, version), + setVersion(luceneScalarQuantizedVectorsFormatFallback, version)); + + hnswScalarQuantizedVectorsFormat = + loadClass( + setVersion(luceneHnswScalarQuantizedVectorsFormat, version), + setVersion(luceneHnswScalarQuantizedVectorsFormatFallback, version)); + + // TODO: Find a better way if possible, but as a separate initiative. + if ("102".equals(version)) { + binaryQuantizedVectorsFormat = + loadClass( + setVersion(luceneBinaryQuantizedVectorsFormat, version), + setVersion(luceneBinaryQuantizedVectorsFormatFallback, version)); + hnswBinaryQuantizedVectorsFormat = + loadClass( + setVersion(luceneHnswBinaryQuantizedVectorsFormat, version), + setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, version)); + } + } + + private static String setVersion(String pkg, String version) { + return pkg.replaceAll("", version); + } + + private static Class loadClass(String defaultClassName, String fallbackClassName) + throws ClassNotFoundException { + try { + return Class.forName(defaultClassName); + } catch (ClassNotFoundException e) { + // Load class from fallback package. + try { + return Class.forName(fallbackClassName); + } catch (ClassNotFoundException e1) { + // Should not reach here. + log.log(Level.SEVERE, "Unable to load class: " + fallbackClassName); + throw e1; + } + } + } + + public static Codec getCodec(String version) + throws ClassNotFoundException, + NoSuchMethodException, + SecurityException, + InstantiationException, + IllegalAccessException, + IllegalArgumentException, + InvocationTargetException { + Class codecClass = + loadClass(setVersion(luceneCodec, version), setVersion(luceneCodecFallback, version)); + Constructor codecClassConstructor = codecClass.getConstructor(); + return (Codec) codecClassConstructor.newInstance(); + } + + public FlatVectorsFormat getLuceneFlatVectorsFormatInstance(FlatVectorsScorer scorer) + throws Exception { + try { + Constructor luceneFlatVectorsFormatConstructor = + flatVectorsFormat.getConstructor(FlatVectorsScorer.class); + return (FlatVectorsFormat) luceneFlatVectorsFormatConstructor.newInstance(scorer); + } catch (Exception e) { + log.log(Level.SEVERE, "Unable to initialize LuceneFlatVectorsFormat: " + e.getMessage()); + throw e; + } + } + + public KnnVectorsReader getLuceneHnswVectorsReaderInstance( + SegmentReadState state, FlatVectorsReader reader) throws Exception { + try { + Constructor luceneHnswVectorsReaderConstructor = + hnswVectorsReader.getConstructor(SegmentReadState.class, FlatVectorsReader.class); + return (KnnVectorsReader) luceneHnswVectorsReaderConstructor.newInstance(state, reader); + } catch (Exception e) { + log.log(Level.SEVERE, "Unable to initialize LuceneHnswVectorsReader: " + e.getMessage()); + throw e; + } + } + + public KnnVectorsWriter getLuceneHnswVectorsWriterInstance( + SegmentWriteState state, + int maxConn, + int beamWidth, + FlatVectorsWriter writer, + int numMergeWorkers, + TaskExecutor executor) + throws Exception { + try { + Constructor luceneHnswVectorsWriterConstructor = + hnswVectorsWriter.getConstructor( + SegmentWriteState.class, + Integer.TYPE, + Integer.TYPE, + FlatVectorsWriter.class, + Integer.TYPE, + TaskExecutor.class); + return (KnnVectorsWriter) + luceneHnswVectorsWriterConstructor.newInstance( + state, maxConn, beamWidth, writer, numMergeWorkers, executor); + } catch (Exception e) { + log.log(Level.SEVERE, "Unable to initialize LuceneHnswVectorsWriter: " + e.getMessage()); + throw e; + } + } + + public int getStaticIntParam(String param) throws ReflectiveOperationException { + try { + VarHandle varHandle = lookup.findStaticVarHandle(hnswVectorsFormat, param, Integer.TYPE); + return (int) varHandle.get(); + } catch (NoSuchFieldException | IllegalAccessException e) { + log.log(Level.SEVERE, "Unable to get " + param + ": " + e.getMessage()); + throw e; + } + } + + public List getSimilarityFunctions() + throws ReflectiveOperationException { + try { + VarHandle varHandle = + lookup.findStaticVarHandle(hnswVectorsReader, "SIMILARITY_FUNCTIONS", List.class); + return (List) varHandle.get(); + } catch (NoSuchFieldException | IllegalAccessException e) { + log.log(Level.SEVERE, "Unable to get SIMILARITY_FUNCTIONS: " + e.getMessage()); + throw e; + } + } + + public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws Exception { + try { + Constructor luceneBinaryQuantizedVectorsFormatConstructor = + binaryQuantizedVectorsFormat.getConstructor(); + return (FlatVectorsFormat) luceneBinaryQuantizedVectorsFormatConstructor.newInstance(); + } catch (Exception e) { + log.log( + Level.SEVERE, + "Unable to initialize LuceneBinaryQuantizedVectorsFormat: " + e.getMessage()); + throw e; + } + } + + public FlatVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( + int maxConn, int beamWidth) throws Exception { + try { + Constructor luceneHnswBinaryQuantizedVectorsFormatConstructor = + hnswBinaryQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); + return (FlatVectorsFormat) + luceneHnswBinaryQuantizedVectorsFormatConstructor.newInstance(maxConn, beamWidth); + } catch (Exception e) { + log.log( + Level.SEVERE, + "Unable to initialize LuceneBinaryQuantizedVectorsFormat: " + e.getMessage()); + throw e; + } + } + + public FlatVectorsFormat getLuceneScalarQuantizedVectorsFormatInstance() throws Exception { + try { + Constructor luceneScalarQuantizedVectorsFormatConstructor = + scalarQuantizedVectorsFormat.getConstructor(); + return (FlatVectorsFormat) luceneScalarQuantizedVectorsFormatConstructor.newInstance(); + } catch (Exception e) { + log.log( + Level.SEVERE, + "Unable to initialize LuceneScalarQuantizedVectorsFormat: " + e.getMessage()); + throw e; + } + } + + public FlatVectorsFormat getLuceneHnswScalarQuantizedVectorsFormatInstance( + int beamWidth, int maxConn) throws Exception { + try { + Constructor luceneHnswScalarQuantizedVectorsFormatConstructor = + hnswScalarQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); + return (FlatVectorsFormat) + luceneHnswScalarQuantizedVectorsFormatConstructor.newInstance(beamWidth, maxConn); + } catch (Exception e) { + log.log( + Level.SEVERE, + "Unable to initialize LuceneHnswScalarQuantizedVectorsFormat: " + e.getMessage()); + throw e; + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java new file mode 100644 index 0000000000..34c95031b2 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSResources; +import java.util.logging.Level; +import java.util.logging.Logger; + +/** + * Provides a mechanism to create ThreadLocal based CuVSResource instances. + * + * @since 26.02 + */ +public class ThreadLocalCuVSResourcesProvider { + + private static final Logger log = + Logger.getLogger(ThreadLocalCuVSResourcesProvider.class.getName()); + private static final ThreadLocal cuVSResources; + + static { + cuVSResources = ThreadLocal.withInitial(() -> cuVSResourcesOrNull()); + } + + /** + * Gets an instance of CuVSResources for the accessing thread. + * + * @return an instance of CuVSResources + */ + public static CuVSResources getCuVSResourcesInstance() { + return cuVSResources.get(); + } + + /** + * Sets the instance of CuVSResources + * + * @param resources the instance of CuVSResources to set + */ + public static void setCuVSResourcesInstance(CuVSResources resources) { + cuVSResources.set(resources); + } + + private static CuVSResources cuVSResourcesOrNull() { + try { + return CuVSResources.create(); + } catch (UnsupportedOperationException uoe) { + log.log( + Level.WARNING, + "cuVS is not supported on this platform or java version: " + uoe.getMessage()); + } catch (Throwable t) { + if (t instanceof ExceptionInInitializerError ex) { + t = ex.getCause(); + } + log.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + t); + } + return null; + } + + /** + * Attempts to close the thread's {@link CuVSResources} instance. + */ + public static void closeCuVSResourcesInstance() { + CuVSResources r = cuVSResources.get(); + if (r != null) { + r.close(); + } + cuVSResources.remove(); + } + + /** + * Checks if cuVS is supported and throws {@link UnsupportedOperationException} otherwise. + * + * @throws UnsupportedOperationException + */ + public static void assertIsSupported() throws UnsupportedOperationException { + if (cuVSResources.get() == null) { + throw new UnsupportedOperationException("cuVS is not supported"); + } + } + + /** + * Checks if cuVS is supported. + * + * @return true if cuVS is supported else false + */ + public static boolean isSupported() { + return cuVSResources.get() != null; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java new file mode 100644 index 0000000000..e4a20d2b4d --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.util.InfoStream; + +/** + * This class provides common static utility methods. + * + * @since 25.10 + */ +public class Utils { + + static final Logger log = Logger.getLogger(Utils.class.getName()); + + /** + * A utility method that throws specific types of throwable objects based on types. + * + * @param t the throwable object + * @throws IOException + */ + static void handleThrowable(Throwable t) throws IOException { + switch (t) { + case IOException ioe -> throw ioe; + case Error error -> throw error; + case RuntimeException re -> throw re; + case null, default -> throw new RuntimeException("UNEXPECTED: exception type", t); + } + } + + /** + * A method to build a CuVSMatrix from a list of float vectors. + * + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. + * + * @param data The float vectors + * @param dimensions The number float elements in each vector + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix + */ + static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { + // Use Builder pattern to avoid intermediate float[][] allocation + // and copy directly from List to device memory + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, + data.size(), // rows (number of vectors) + dimensions, // columns (vector dimension) + CuVSMatrix.DataType.FLOAT); + + // Add vectors one by one - builder copies directly to device memory + for (float[] vector : data) { + builder.addVector(vector); + } + + return builder.build(); + } + + /** + * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). + * + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. + * + * @param data The byte vectors (packed bits for binary quantization) + * @param bytesPerVector The number of bytes in each vector + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix with BYTE data type + */ + static CuVSMatrix createByteMatrix( + List data, int bytesPerVector, CuVSResources resources) { + // Use Builder pattern to avoid intermediate byte[][] allocation + // and copy directly from List to device memory + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, + data.size(), // rows (number of vectors) + bytesPerVector, // columns (bytes per vector) + CuVSMatrix.DataType.BYTE); + + // Add vectors one by one - builder copies directly to device memory + for (byte[] vector : data) { + builder.addVector(vector); + } + + return builder.build(); + } + + /** + * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). + * + * @param data The 2D byte array (packed bits for binary quantization) + * @param bytesPerVector The number of bytes in each vector + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix with BYTE data type + */ + static CuVSMatrix createByteMatrixFromArray( + byte[][] data, int bytesPerVector, CuVSResources resources) { + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, + data.length, // rows (number of vectors) + bytesPerVector, // columns (bytes per vector) + CuVSMatrix.DataType.BYTE); + + // Add vectors one by one - builder copies directly to device memory + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); + } + + /** + * A utility method to convert nanoseconds to milliseconds. + * + * @param nanos + * @return milliseconds + */ + static long nanosToMillis(long nanos) { + return Duration.ofNanos(nanos).toMillis(); + } + + /** + * Creates an instance of CuVSResources. + * + * @return an instance of CuVSResources + */ + static CuVSResources cuVSResourcesOrNull() { + try { + System.loadLibrary("cudart"); + } catch (UnsatisfiedLinkError e) { + log.log(Level.WARNING, "Could not load CUDA runtime library: " + e.getMessage()); + } + try { + return CuVSResources.create(); + } catch (UnsupportedOperationException uoe) { + log.log( + Level.WARNING, + "cuVS is not supported on this platform or java version: " + uoe.getMessage()); + } catch (Throwable t) { + if (t instanceof ExceptionInInitializerError ex) { + t = ex.getCause(); + } + log.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + t); + } + return null; + } + + /** + * A utility method that conditionally ignores certain throwable objects + * + * @param t the throwable object + * @param msg the message to check + * @throws IOException + */ + static void handleThrowableWithIgnore(Throwable t, String msg) throws IOException { + if (t.getMessage().contains(msg)) { + return; + } + handleThrowable(t); + } + + /** + * Creates a list of float vectors from the input + * + * @param mergedVectorValues instance of {@link FloatVectorValues} + * @return a list of float arrays + * @throws IOException I/O Exception + */ + static List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List vectors = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + float[] vector = mergedVectorValues.vectorValue(iter.index()); + vectors.add(vector.clone()); + } + return vectors; + } + + /** + * Utility to print info/debug messages via InfoStream. + * + * @param infoStream the writer's infostream + * @param component the name of the index writer + * @param msg the log message to push via the InfoStream + */ + static void info(InfoStream infoStream, String component, String msg) { + if (infoStream.isEnabled(component)) { + infoStream.message(component, msg); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/package-info.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/package-info.java new file mode 100644 index 0000000000..f14c45183a --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/package-info.java @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * GPU-based accelerated vector search pluggable library for Apache Lucene search platforms. + * + * @since 25.10 + */ +package com.nvidia.cuvs.lucene; diff --git a/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec new file mode 100644 index 0000000000..faa0684c04 --- /dev/null +++ b/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec +com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec +com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedCodec +com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedCodec diff --git a/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat new file mode 100644 index 0000000000..6625ac72a4 --- /dev/null +++ b/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat +org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat +com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat +com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat +com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat +com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedVectorsFormat diff --git a/java/cuvs-lucene/src/main/resources/logging.properties b/java/cuvs-lucene/src/main/resources/logging.properties new file mode 100644 index 0000000000..337c59000e --- /dev/null +++ b/java/cuvs-lucene/src/main/resources/logging.properties @@ -0,0 +1,9 @@ +# Set root level +.level=WARNING + +# Define handlers +handlers=java.util.logging.ConsoleHandler + +java.util.logging.ConsoleHandler.level=WARNING +java.util.logging.ConsoleHandler.formatter=java.util.logging.SimpleFormatter +java.util.logging.SimpleFormatter.format=[%1$tF %1$tT] %4$s %2$s %5$s%6$s%n diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java new file mode 100644 index 0000000000..cd4aaa6ec8 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -0,0 +1,335 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWDeletedDocuments extends LuceneTestCase { + + protected static Logger log = + Logger.getLogger(TestAcceleratedHNSWDeletedDocuments.class.getName()); + + static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + private static Random random; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = random(); + } + + @Test + public void testVectorSearchWithDeletedDocuments() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200, 1000); // 200-1200 documents + int dimensions = random.nextInt(64, 256); // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set deletedDocs = new HashSet<>(); + + // Create index with all documents having vectors + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + + // Delete documents randomly based on probability + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Search and verify deleted documents are not returned + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + // Use a random vector for query + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we got results + assertTrue("Should have search results", hits.length > 0); + + // Verify no deleted documents in results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse( + "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); + log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); + } + + // Verify deleted documents are truly deleted + for (int deletedId : deletedDocs) { + TopDocs result = + searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + assertEquals( + "Deleted document " + deletedId + " should not be found", + 0, + result.totalHits.value()); + } + } + } + } + + @Test + public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set docsWithoutVectors = new HashSet<>(); + Set deletedDocs = new HashSet<>(); + + // Create index with mixed documents + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + // Randomly assign categories + String category = random.nextBoolean() ? "A" : "B"; + doc.add(new StringField("category", category, Field.Store.YES)); + + // Randomly decide whether to add vectors + if (random.nextFloat() < vectorProbability) { + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } else { + docsWithoutVectors.add(i); + } + writer.addDocument(doc); + } + + // Delete documents randomly + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Test vector search behavior + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse("Deleted document should not appear", deletedDocs.contains(id)); + assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); + log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); + } + + // Test filtered search with deletions + Query filter = new TermQuery(new Term("category", "A")); + Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + for (ScoreDoc hit : filteredHits) { + Document doc = reader.storedFields().document(hit.doc); + String category = doc.get("category"); + assertEquals("Should only match category A", "A", category); + int id = Integer.parseInt(doc.get("id")); + assertFalse( + "Deleted document should not appear in filtered results", deletedDocs.contains(id)); + } + } + } + } + + @Test + public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test + int dimensions = random.nextInt(128) + 32; // 32-160 dimensions + int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + + // Create and delete all documents + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + + // Delete all documents + for (int i = 0; i < datasetSize; i++) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + } + writer.commit(); + writer.forceMerge(1); // Force merge to apply deletions + } + + // Verify search returns no results + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + TopDocs results = searcher.search(query, topK); + + assertEquals( + "Should return no results when all documents are deleted", + 0, + results.totalHits.value()); + } + } + } + + @Test + public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + List activeDocIds = new ArrayList<>(); + + // Initial indexing + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + for (int i = 0; i < initialDocs; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + + // Delete some documents randomly + List candidatesForDeletion = new ArrayList<>(activeDocIds); + for (int docId : candidatesForDeletion) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(docId))); + activeDocIds.remove(Integer.valueOf(docId)); + } + } + + // Add new documents with higher IDs + for (int i = initialDocs; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + writer.commit(); + } + + // Verify search behavior after deletions and additions + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + Set resultIds = new HashSet<>(); + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + resultIds.add(id); + assertTrue("Result should be from active documents", activeDocIds.contains(id)); + } + + log.log( + Level.FINE, + "Search returned " + + hits.length + + " results from " + + activeDocIds.size() + + " active documents"); + } + } + } + + private RandomIndexWriter createWriter(Directory directory) throws IOException { + return new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + } + + private IndexWriterConfig createWriterConfig() { + return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java new file mode 100644 index 0000000000..03b5633a87 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -0,0 +1,194 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_BEAM_WIDTH; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CUVS_DISTANCE_TYPE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_INT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_MAX_CONN; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_NN_DESCENT_NUM_ITERATIONS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_NUM_MERGE_WORKERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_STRATEGY; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_BEAM_WIDTH; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_INT_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_MAX_CONN; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_NN_DESCENT_NUM_ITERATIONS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_NUM_MERGE_WORKERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_BEAM_WIDTH; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_INT_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_MAX_CONN; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_NN_DESCENT_NUM_ITERATIONS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_NUM_MERGE_WORKERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_WRITER_THREADS; +import static java.lang.Integer.MAX_VALUE; +import static java.lang.Integer.MIN_VALUE; + +import java.util.Random; +import java.util.logging.Logger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWParams extends LuceneTestCase { + + @SuppressWarnings("unused") + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWParams.class.getName()); + + private static Random random; + + @Test + public void testAcceleratedHNSWParamsDefaultValues() { + AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); + assertEquals(DEFAULT_BEAM_WIDTH, params.getBeamWidth()); + assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); + assertEquals(DEFAULT_HNSW_LAYERS, params.getHnswLayers()); + assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); + assertEquals(DEFAULT_MAX_CONN, params.getMaxConn()); + assertEquals(DEFAULT_WRITER_THREADS, params.getWriterThreads()); + assertEquals(DEFAULT_NUM_MERGE_WORKERS, params.getNumMergeWorkers()); + assertEquals(DEFAULT_CAGRA_GRAPH_BUILD_ALGO, params.getCagraGraphBuildAlgo()); + assertEquals(DEFAULT_STRATEGY, params.getStrategy()); + assertEquals(DEFAULT_CUVS_DISTANCE_TYPE, params.getCuvsDistanceType()); + assertEquals(DEFAULT_NN_DESCENT_NUM_ITERATIONS, params.getNNDescentNumIterations()); + } + + @Test + public void testAcceleratedHNSWParamsInvalidBeamWidth() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_BEAM_WIDTH), random.nextInt(MAX_BEAM_WIDTH + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withBeamWidth(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidGraphDegree() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_GRAPH_DEG), random.nextInt(MAX_GRAPH_DEG + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withGraphDegree(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidHNSWLayers() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_HNSW_LAYERS), random.nextInt(MAX_HNSW_LAYERS + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withHNSWLayer(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidIntGraphDegree() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_INT_GRAPH_DEG), + random.nextInt(MAX_INT_GRAPH_DEG + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withIntermediateGraphDegree(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidMaxConn() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_MAX_CONN), random.nextInt(MAX_MAX_CONN + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withMaxConn(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidWriterThreads() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_WRITER_THREADS), + random.nextInt(MAX_WRITER_THREADS + 1, Integer.MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withWriterThreads(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidNumMergeWorkers() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_NUM_MERGE_WORKERS), + random.nextInt(MAX_NUM_MERGE_WORKERS + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withNumMergeWorkers(v).build()); + } + } + + @Test + public void testAcceleratedHNSWParamsInvalidCagraGraphBuildAlgo() { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withCagraGraphBuildAlgo(null).build()); + } + + @Test + public void testAcceleratedHNSWParamsInvalidStrategy() { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withStrategy(null).build()); + } + + @Test + public void testAcceleratedHNSWParamsInvalidCuvsDistanceType() { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withCuvsDistanceType(null).build()); + } + + @Test + public void testAcceleratedHNSWParamsInvalidNNDescentNumIterations() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, (int) MIN_NN_DESCENT_NUM_ITERATIONS), + random.nextInt((int) (MAX_NN_DESCENT_NUM_ITERATIONS + 1), MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withNNDescentNumIterations(v).build()); + } + } + + @BeforeClass + public static void beforeClass() { + random = random(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java new file mode 100644 index 0000000000..d638180b06 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertTrue; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.junit.Test; + +/** + * Tests the backward compatibility mechanism. + * + * @since 25.12 + */ +public class TestBackCompat { + + @Test + public void testFallback() throws Exception { + // Lucene99Codec exists in the org.apache.lucene.backward_codecs.lucene99 + Codec c = LuceneProvider.getCodec("99"); + assertEquals(c.getName(), "Lucene99"); + } + + @Test(expected = ClassNotFoundException.class) + public void testNonexistentCodec() throws Exception { + LuceneProvider.getCodec("0"); + } + + @Test + public void testExistingComponents() throws Exception { + LuceneProvider provider = LuceneProvider.getInstance("99"); + assertTrue(provider.getLuceneFlatVectorsFormatInstance(null) instanceof FlatVectorsFormat); + assertEquals(provider.getStaticIntParam("VERSION_CURRENT"), 0); + assertNotEquals(provider.getSimilarityFunctions().size(), 0); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBogusAddPropagation.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBogusAddPropagation.java new file mode 100644 index 0000000000..2f096c65fe --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBogusAddPropagation.java @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.Test; + +/** + * Verifies end-to-end build propagation from the cuVS C layer, through the cuvs-java bindings, to + * cuvs-lucene: the bogus {@code cuvsBogusAddV1} C function (added in {@code c/src/core/c_api.cpp}) is + * exposed as {@code CuVSProvider.bogusAddV1} in cuvs-java and is called from the cuvs-lucene {@link + * BogusPropagation#bogusAddV1(int, int)} method exercised here. If any layer of the build fails to + * propagate, this test will fail to compile or fail at runtime. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestBogusAddPropagation extends LuceneTestCase { + + @Test + public void testBogusAddReachableFromLucene() { + assertEquals(5, BogusPropagation.bogusAddV1(2, 3)); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java new file mode 100644 index 0000000000..67e6b0ad30 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -0,0 +1,202 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { + + private static Logger log = + Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); + private static Random random; + private static Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + // Fixed seed so that we can validate against the same result. + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @Test + public void testCagraToHnswSerializationAndSearch() throws Exception { + AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.log(Level.FINE, "Successfully opened index"); + + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(knnValues); + log.log( + Level.FINE, + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + } + assertTrue("Dataset size mismatch", vectorCount == numDocs); + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + Integer[] expected = new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}; + HashSet expectedIds = new HashSet(Arrays.asList(expected)); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue( + "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + } + } + + @Test + public void testSingleVectorIndex() throws Exception { + // Test single vector index support with dummy HNSW graph + // TODO: This test can be removed once https://github.com/rapidsai/cuvs/pull/1256 is merged + // and CAGRA natively supports single vector indexes + Codec codec = new Lucene101AcceleratedHNSWCodec(); + + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int dimension = 32; + float[] vector = generateDataset(random, 1, dimension)[0]; + + // Index a single document with a vector - this should now work with dummy HNSW graph + try (Directory indexDirectory = newDirectory()) { + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + try (IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, "0", Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + indexWriter.addDocument(document); + + // This should now succeed by creating a dummy HNSW graph for the single vector + indexWriter.commit(); + } + + // Verify the index can be opened and searched + try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + assertEquals(1, reader.numDocs()); + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(knnValues); + assertEquals(1, knnValues.size()); + assertEquals(dimension, knnValues.dimension()); + + // Test search functionality + IndexSearcher searcher = new IndexSearcher(reader); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, vector, 1); + TopDocs results = searcher.search(query, 1); + assertEquals(1, results.totalHits.value()); + assertEquals(1, results.scoreDocs.length); + assertEquals(0, results.scoreDocs[0].doc); + } + } + } + + @After + public void afterTest() throws Exception { + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java new file mode 100644 index 0000000000..8999d8e81e --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -0,0 +1,173 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.setCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraToHnswSerializationAndSearchWithFallbackWriter extends LuceneTestCase { + + private static Logger log = + Logger.getLogger(TestCagraToHnswSerializationAndSearchWithFallbackWriter.class.getName()); + + private static Random random; + private static Path indexDirPath; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + // Set resources to null to simulate that cuVS is not supported. + setCuVSResourcesInstance(null); + // Fixed seed so that we can validate against the same result. + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @Test + public void testCagraToHnswSerializationAndSearchWithFallbackWriter() throws Exception { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withHNSWLayer(3) + .withMaxConn(16) + .withBeamWidth(100) + .build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath)) { + try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.log(Level.FINE, "Successfully opened index"); + + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(knnValues); + log.log( + Level.FINE, + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + } + assertTrue("Dataset size mismatch", vectorCount == numDocs); + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + Integer[] expected = new Integer[] {1869, 1411, 1497, 351, 554}; + HashSet expectedIds = new HashSet(Arrays.asList(expected)); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue( + "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + + } catch (Exception e) { + e.printStackTrace(); + } + } + } + + @AfterClass + public static void afterClass() throws Exception { + // Reset resources for other tests to work + setCuVSResourcesInstance(cuVSResourcesOrNull()); + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java new file mode 100644 index 0000000000..10d3ceeec2 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java @@ -0,0 +1,335 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSAcceleratedHNSWDeletedDocuments extends LuceneTestCase { + + protected static Logger log = + Logger.getLogger(TestCuVSAcceleratedHNSWDeletedDocuments.class.getName()); + + static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + private static Random random; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = random(); + } + + @Test + public void testVectorSearchWithDeletedDocuments() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200, 1000); // 200-1200 documents + int dimensions = random.nextInt(64, 256); // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set deletedDocs = new HashSet<>(); + + // Create index with all documents having vectors + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + + // Delete documents randomly based on probability + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Search and verify deleted documents are not returned + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + // Use a random vector for query + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we got results + assertTrue("Should have search results", hits.length > 0); + + // Verify no deleted documents in results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse( + "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); + log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); + } + + // Verify deleted documents are truly deleted + for (int deletedId : deletedDocs) { + TopDocs result = + searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + assertEquals( + "Deleted document " + deletedId + " should not be found", + 0, + result.totalHits.value()); + } + } + } + } + + @Test + public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set docsWithoutVectors = new HashSet<>(); + Set deletedDocs = new HashSet<>(); + + // Create index with mixed documents + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + // Randomly assign categories + String category = random.nextBoolean() ? "A" : "B"; + doc.add(new StringField("category", category, Field.Store.YES)); + + // Randomly decide whether to add vectors + if (random.nextFloat() < vectorProbability) { + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } else { + docsWithoutVectors.add(i); + } + writer.addDocument(doc); + } + + // Delete documents randomly + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Test vector search behavior + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse("Deleted document should not appear", deletedDocs.contains(id)); + assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); + log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); + } + + // Test filtered search with deletions + Query filter = new TermQuery(new Term("category", "A")); + Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + for (ScoreDoc hit : filteredHits) { + Document doc = reader.storedFields().document(hit.doc); + String category = doc.get("category"); + assertEquals("Should only match category A", "A", category); + int id = Integer.parseInt(doc.get("id")); + assertFalse( + "Deleted document should not appear in filtered results", deletedDocs.contains(id)); + } + } + } + } + + @Test + public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test + int dimensions = random.nextInt(128) + 32; // 32-160 dimensions + int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + + // Create and delete all documents + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + + // Delete all documents + for (int i = 0; i < datasetSize; i++) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + } + writer.commit(); + writer.forceMerge(1); // Force merge to apply deletions + } + + // Verify search returns no results + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + TopDocs results = searcher.search(query, topK); + + assertEquals( + "Should return no results when all documents are deleted", + 0, + results.totalHits.value()); + } + } + } + + @Test + public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + List activeDocIds = new ArrayList<>(); + + // Initial indexing + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + for (int i = 0; i < initialDocs; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + + // Delete some documents randomly + List candidatesForDeletion = new ArrayList<>(activeDocIds); + for (int docId : candidatesForDeletion) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(docId))); + activeDocIds.remove(Integer.valueOf(docId)); + } + } + + // Add new documents with higher IDs + for (int i = initialDocs; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + writer.commit(); + } + + // Verify search behavior after deletions and additions + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + Set resultIds = new HashSet<>(); + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + resultIds.add(id); + assertTrue("Result should be from active documents", activeDocIds.contains(id)); + } + + log.log( + Level.FINE, + "Search returned " + + hits.length + + " results from " + + activeDocIds.size() + + " active documents"); + } + } + } + + private RandomIndexWriter createWriter(Directory directory) throws IOException { + return new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + } + + private IndexWriterConfig createWriterConfig() { + return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java new file mode 100644 index 0000000000..fa0936e67e --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSAcceleratedHNSWGaps extends LuceneTestCase { + + protected static Logger log = Logger.getLogger(TestCuVSAcceleratedHNSWGaps.class.getName()); + + static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + static IndexSearcher searcher; + static IndexReader reader; + static Directory directory; + static Random random; + + static int DATASET_SIZE_LIMIT = 1000; + static int DIMENSIONS_LIMIT = 2048; + static int NUM_QUERIES_LIMIT = 10; + static int TOP_K_LIMIT = 64; + + static int datasetSize; + static int dimension; + static float[][] dataset; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + directory = newDirectory(); + random = random(); + + RandomIndexWriter writer = + new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + + log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); + + datasetSize = random.nextInt(100, DATASET_SIZE_LIMIT); + dimension = random.nextInt(8, DIMENSIONS_LIMIT); + dataset = generateDataset(random, datasetSize, dimension); + + // Create documents where only even-numbered documents have vectors + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); + + // Only add vectors to even-numbered documents + if (i % 2 == 0) { + doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.log(Level.FINE, "Test finished"); + } + + @Test + public void testVectorSearchWithAlternatingDocuments() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + + // Use the first vector (from document 0) as query + float[] queryVector = dataset[0]; + int topK = random.nextInt(5, TOP_K_LIMIT); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we get exactly TOP_K results + assertEquals("Should return exactly " + topK + " results", topK, hits.length); + + // Verify all returned documents have vectors (even-numbered IDs) + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); + log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); + } + + // Verify the results match expected top-k based on Euclidean distance + List expectedIds = calculateExpectedTopK(queryVector, topK, dataset); + for (int i = 0; i < hits.length; i++) { + String docId = reader.storedFields().document(hits[i].doc).get("id"); + int id = Integer.parseInt(docId); + assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); + } + + log.log(Level.FINE, "Alternating document test passed with " + hits.length + " results"); + } + + @Test + public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + + // Use the first vector (from document 0) as query + float[] queryVector = dataset[0]; + int topK = random.nextInt(5, TOP_K_LIMIT); + + // Create a filter that only matches documents with ID less than 10 + // This should further restrict our results to even numbers 0, 2, 4, 6, 8 + Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 + + Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Should only get document 8 (the only one that matches the filter and has a vector) + assertEquals("Should return exactly 1 result", 1, filteredHits.length); + + String docId = reader.storedFields().document(filteredHits[0].doc).get("id"); + assertEquals("Should only return document 8", "8", docId); + + log.log( + Level.FINE, + "Filtered alternating document test passed with " + filteredHits.length + " results"); + } + + public static List calculateExpectedTopK(float[] query, int topK, float[][] dataset) { + Map distances = new TreeMap<>(); + + // Calculate distances only for documents that have vectors (even-numbered) + for (int i = 0; i < dataset.length; i += 2) { + double distance = 0; + for (int j = 0; j < dataset[0].length; j++) { + distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); + } + distances.put(i, distance); + } + + // Sort by distance and return top-k + return distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .limit(topK) + .toList(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java new file mode 100644 index 0000000000..7303d04d2f --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -0,0 +1,337 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSDeletedDocuments extends LuceneTestCase { + + protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + private static Random random; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = random(); + } + + @Test + public void testVectorSearchWithDeletedDocuments() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200, 1000); // 200-1200 documents + int dimensions = random.nextInt(64, 256); // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set deletedDocs = new HashSet<>(); + + // Create index with all documents having vectors + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + + // Delete documents randomly based on probability + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Search and verify deleted documents are not returned + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + // Use a random vector for query + float[] queryVector = generateRandomVector(dimensions, random); + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we got results + assertTrue("Should have search results", hits.length > 0); + + // Verify no deleted documents in results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse( + "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); + log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); + } + + // Verify deleted documents are truly deleted + for (int deletedId : deletedDocs) { + TopDocs result = + searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + assertEquals( + "Deleted document " + deletedId + " should not be found", + 0, + result.totalHits.value()); + } + } + } + } + + @Test + public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set docsWithoutVectors = new HashSet<>(); + Set deletedDocs = new HashSet<>(); + + // Create index with mixed documents + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + // Randomly assign categories + String category = random.nextBoolean() ? "A" : "B"; + doc.add(new StringField("category", category, Field.Store.YES)); + + // Randomly decide whether to add vectors + if (random.nextFloat() < vectorProbability) { + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } else { + docsWithoutVectors.add(i); + } + writer.addDocument(doc); + } + + // Delete documents randomly + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Test vector search behavior + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse("Deleted document should not appear", deletedDocs.contains(id)); + assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); + log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); + } + + // Test filtered search with deletions + Query filter = new TermQuery(new Term("category", "A")); + GPUKnnFloatVectorQuery filteredQuery = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + for (ScoreDoc hit : filteredHits) { + Document doc = reader.storedFields().document(hit.doc); + String category = doc.get("category"); + assertEquals("Should only match category A", "A", category); + int id = Integer.parseInt(doc.get("id")); + assertFalse( + "Deleted document should not appear in filtered results", deletedDocs.contains(id)); + } + } + } + } + + @Test + public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test + int dimensions = random.nextInt(128) + 32; // 32-160 dimensions + int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + + // Create and delete all documents + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + + // Delete all documents + for (int i = 0; i < datasetSize; i++) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + } + writer.commit(); + writer.forceMerge(1); // Force merge to apply deletions + } + + // Verify search returns no results + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + TopDocs results = searcher.search(query, topK); + + assertEquals( + "Should return no results when all documents are deleted", + 0, + results.totalHits.value()); + } + } + } + + @Test + public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + List activeDocIds = new ArrayList<>(); + + // Initial indexing + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + for (int i = 0; i < initialDocs; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + + // Delete some documents randomly + List candidatesForDeletion = new ArrayList<>(activeDocIds); + for (int docId : candidatesForDeletion) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(docId))); + activeDocIds.remove(Integer.valueOf(docId)); + } + } + + // Add new documents with higher IDs + for (int i = initialDocs; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + writer.commit(); + } + + // Verify search behavior after deletions and additions + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + Set resultIds = new HashSet<>(); + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + resultIds.add(id); + assertTrue("Result should be from active documents", activeDocIds.contains(id)); + } + + log.log( + Level.FINE, + "Search returned " + + hits.length + + " results from " + + activeDocIds.size() + + " active documents"); + } + } + } + + private RandomIndexWriter createWriter(Directory directory) throws IOException { + return new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + } + + private IndexWriterConfig createWriterConfig() { + return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java new file mode 100644 index 0000000000..31958ef0db --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSGaps extends LuceneTestCase { + + protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + static IndexSearcher searcher; + static IndexReader reader; + static Directory directory; + static Random random; + + static int DATASET_SIZE_LIMIT = 1000; + static int DIMENSIONS_LIMIT = 2048; + static int NUM_QUERIES_LIMIT = 10; + static int TOP_K_LIMIT = 64; + + static int datasetSize; + static int dimension; + static float[][] dataset; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + directory = newDirectory(); + random = random(); + + RandomIndexWriter writer = + new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + + log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); + + datasetSize = random.nextInt(100, DATASET_SIZE_LIMIT); + dimension = random.nextInt(8, DIMENSIONS_LIMIT); + dataset = generateDataset(random, datasetSize, dimension); + + // Create documents where only even-numbered documents have vectors + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); + + // Only add vectors to even-numbered documents + if (i % 2 == 0) { + doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.log(Level.FINE, "Test finished"); + } + + @Test + public void testVectorSearchWithAlternatingDocuments() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + + // Use the first vector (from document 0) as query + float[] queryVector = dataset[0]; + int topK = random.nextInt(5, TOP_K_LIMIT); + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we get exactly TOP_K results + assertEquals("Should return exactly " + topK + " results", topK, hits.length); + + // Verify all returned documents have vectors (even-numbered IDs) + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); + log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); + } + + // Verify the results match expected top-k based on Euclidean distance + List expectedIds = calculateExpectedTopK(queryVector, topK, dataset); + for (int i = 0; i < hits.length; i++) { + String docId = reader.storedFields().document(hits[i].doc).get("id"); + int id = Integer.parseInt(docId); + assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); + } + + log.log(Level.FINE, "Alternating document test passed with " + hits.length + " results"); + } + + @Test + public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + + // Use the first vector (from document 0) as query + float[] queryVector = dataset[0]; + int topK = random.nextInt(5, TOP_K_LIMIT); + + // Create a filter that only matches documents with ID less than 10 + // This should further restrict our results to even numbers 0, 2, 4, 6, 8 + Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 + + GPUKnnFloatVectorQuery filteredQuery = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Should only get document 8 (the only one that matches the filter and has a vector) + assertEquals("Should return exactly 1 result", 1, filteredHits.length); + + String docId = reader.storedFields().document(filteredHits[0].doc).get("id"); + assertEquals("Should only return document 8", "8", docId); + + log.log( + Level.FINE, + "Filtered alternating document test passed with " + filteredHits.length + " results"); + } + + public static List calculateExpectedTopK(float[] query, int topK, float[][] dataset) { + Map distances = new TreeMap<>(); + + // Calculate distances only for documents that have vectors (even-numbered) + for (int i = 0; i < dataset.length; i += 2) { + double distance = 0; + for (int j = 0; j < dataset[0].length; j++) { + distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); + } + distances.put(i, distance); + } + + // Sort by distance and return top-k + return distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .limit(topK) + .toList(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java new file mode 100644 index 0000000000..e630a0c3a1 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java @@ -0,0 +1,218 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSRandomizedHNSWVectorSearch extends LuceneTestCase { + + protected static Logger log = + Logger.getLogger(TestCuVSRandomizedHNSWVectorSearch.class.getName()); + + static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + static IndexSearcher searcher; + static IndexReader reader; + static Directory directory; + + static int DATASET_SIZE_LIMIT = 1000; + static int DIMENSIONS_LIMIT = 2048; + static int NUM_QUERIES_LIMIT = 10; + static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 + static float[][] dataset; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + directory = newDirectory(); + + RandomIndexWriter writer = + new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + + log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); + + Random random = random(); + int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; + int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; + dataset = generateDataset(random, datasetSize, dimensions); + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); + boolean skipVector = + random.nextInt(10) + < 4; // some documents won't have vectors to test deleted/missing vectors + if (!skipVector + || datasetSize < 100) { // about 10th of the documents shouldn't have a single vector + doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField("vector2", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.log(Level.FINE, "Test finished"); + } + + @Test + public void testVectorSearch() throws IOException { + Random random = random(); + int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; + int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); + + if (dataset.length < topK) topK = dataset.length; + + float[][] queries = generateQueries(random, dataset[0].length, numQueries); + List> expected = generateExpectedResults(topK, dataset, queries); + + log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); + log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); + log.log(Level.FINE, "TopK: " + topK); + + Query query = new KnnFloatVectorQuery("vector", queries[0], topK); + int correct[] = new int[topK]; + for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); + + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); + log.log(Level.FINE, "EXPECTD: " + expected.get(0)); + + for (ScoreDoc hit : hits) { + log.log( + Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); + } + + for (ScoreDoc hit : hits) { + int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); + assertTrue("Result returned was not in topk*2: " + doc, expected.get(0).contains(doc)); + } + } + + private static List> generateExpectedResults( + int topK, float[][] dataset, float[][] queries) { + List> neighborsResult = new ArrayList<>(); + int dimensions = dataset[0].length; + + for (float[] query : queries) { + Map distances = new TreeMap<>(); + for (int j = 0; j < dataset.length; j++) { + double distance = 0; + for (int k = 0; k < dimensions; k++) { + distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); + } + distances.put(j, (distance)); + } + + Map sorted = new TreeMap(distances); + log.log(Level.FINE, "EXPECTED: " + sorted); + + // Sort by distance and select the topK nearest neighbors + List neighbors = + distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .toList(); + neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); + } + + log.log(Level.FINE, "Expected results generated successfully."); + return neighborsResult; + } + + @Test + public void testVectorSearchWithFilter() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + + Random random = random(); + int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); + + if (dataset.length < topK) topK = dataset.length; + + // Find a document that has a vector by doing a search first + Query unfiltered = new KnnFloatVectorQuery("vector", dataset[0], 1); + ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; + + // Skip test if no vectors found at all + assumeTrue( + "Need at least one document with vector for filtering test", unfilteredHits.length > 0); + + String targetDocId = reader.storedFields().document(unfilteredHits[0].doc).get("id"); + float[] queryVector = dataset[0]; + + // Create a filter that matches only the document we know has a vector + Query filter = new TermQuery(new Term("id", targetDocId)); + + // Test the new constructor with filter + Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); + + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Ensure we got some results + assertTrue("Should have at least one result", filteredHits.length > 0); + + // Verify that all results match the filter + for (ScoreDoc hit : filteredHits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + assertEquals("All results should match the filter", targetDocId, docId); + } + + log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java new file mode 100644 index 0000000000..85720c3025 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -0,0 +1,218 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { + + protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + static IndexSearcher searcher; + static IndexReader reader; + static Directory directory; + + static int DATASET_SIZE_LIMIT = 1000; + static int DIMENSIONS_LIMIT = 2048; + static int NUM_QUERIES_LIMIT = 10; + static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 + static float[][] dataset; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + directory = newDirectory(); + + RandomIndexWriter writer = + new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + + log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); + + Random random = random(); + int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; + int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; + dataset = generateDataset(random, datasetSize, dimensions); + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); + boolean skipVector = + random.nextInt(10) + < 4; // some documents won't have vectors to test deleted/missing vectors + if (!skipVector + || datasetSize < 100) { // about 10th of the documents shouldn't have a single vector + doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField("vector2", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.log(Level.FINE, "Test finished"); + } + + @Test + public void testVectorSearch() throws IOException { + Random random = random(); + int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; + int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); + + if (dataset.length < topK) topK = dataset.length; + + float[][] queries = generateQueries(random, dataset[0].length, numQueries); + List> expected = generateExpectedResults(topK, dataset, queries); + + log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); + log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); + log.log(Level.FINE, "TopK: " + topK); + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queries[0], topK, null, topK, 1); + int correct[] = new int[topK]; + for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); + + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); + log.log(Level.FINE, "EXPECTD: " + expected.get(0)); + + for (ScoreDoc hit : hits) { + log.log( + Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); + } + + for (ScoreDoc hit : hits) { + int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); + assertTrue("Result returned was not in topk*2: " + doc, expected.get(0).contains(doc)); + } + } + + private static List> generateExpectedResults( + int topK, float[][] dataset, float[][] queries) { + List> neighborsResult = new ArrayList<>(); + int dimensions = dataset[0].length; + + for (float[] query : queries) { + Map distances = new TreeMap<>(); + for (int j = 0; j < dataset.length; j++) { + double distance = 0; + for (int k = 0; k < dimensions; k++) { + distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); + } + distances.put(j, (distance)); + } + + Map sorted = new TreeMap(distances); + log.log(Level.FINE, "EXPECTED: " + sorted); + + // Sort by distance and select the topK nearest neighbors + List neighbors = + distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .toList(); + neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); + } + + log.log(Level.FINE, "Expected results generated successfully."); + return neighborsResult; + } + + @Test + public void testVectorSearchWithFilter() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + + Random random = random(); + int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); + + if (dataset.length < topK) topK = dataset.length; + + // Find a document that has a vector by doing a search first + GPUKnnFloatVectorQuery unfiltered = + new GPUKnnFloatVectorQuery("vector", dataset[0], 1, null, 1, 1); + ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; + + // Skip test if no vectors found at all + assumeTrue( + "Need at least one document with vector for filtering test", unfilteredHits.length > 0); + + String targetDocId = reader.storedFields().document(unfilteredHits[0].doc).get("id"); + float[] queryVector = dataset[0]; + + // Create a filter that matches only the document we know has a vector + Query filter = new TermQuery(new Term("id", targetDocId)); + + // Test the new constructor with filter + GPUKnnFloatVectorQuery filteredQuery = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Ensure we got some results + assertTrue("Should have at least one result", filteredHits.length > 0); + + // Verify that all results match the filter + for (ScoreDoc hit : filteredHits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + assertEquals("All results should match the filter", targetDocId, docId); + } + + log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java new file mode 100644 index 0000000000..2b960bf3cd --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.util.List; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Ignore; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS is not supported", isSupported()); + } + + @Override + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + } + + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { + float[][] f = new float[][] {randomVector(384), randomVector(384)}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc1 = new Document(); + doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); + w.addDocument(doc1); + w.commit(); + Document doc2 = new Document(); + doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); + w.addDocument(doc2); + w.flush(); + w.commit(); + + // sanity - verify one doc per leaf + try (DirectoryReader reader = DirectoryReader.open(w)) { + List subReaders = reader.leaves(); + assertEquals(2, subReaders.size()); + assertEquals(1, subReaders.get(0).reader().getFloatVectorValues("f").size()); + assertEquals(1, subReaders.get(1).reader().getFloatVectorValues("f").size()); + } + + // now merge to a single segment + w.forceMerge(1); + + // verify merged content + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues values = r.getFloatVectorValues("f"); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(f[0], values.vectorValue(0), 0.0f); + assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + } + } + } + + // Basic test for multiple vectors fields per document + public void testTwoVectorFieldsPerDoc() throws Exception { + float[][] f1 = new float[][] {randomVector(384), randomVector(384)}; + float[][] f2 = new float[][] {randomVector(384), randomVector(384)}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc1 = new Document(); + doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); + doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); + w.addDocument(doc1); + Document doc2 = new Document(); + doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); + doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); + w.addDocument(doc2); + w.forceMerge(1); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues values = r.getFloatVectorValues("f1"); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); + assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); + + values = r.getFloatVectorValues("f2"); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); + assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); + + // opportunistically check boundary condition - search with a 0 topK + var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); + assertEquals(0, topDocs.scoreDocs.length); + assertEquals(0, topDocs.totalHits.value()); + } + } + } + + @Override + // Overriding this method from superclass for the tests to only use float vector encoding + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testByteVectorScorerIteration() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testEmptyByteVectorData() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testMergingWithDifferentByteKnnFields() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testMismatchedFields() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testRandomBytes() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testSortedIndexBytes() {} +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java new file mode 100644 index 0000000000..32e5c81426 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java @@ -0,0 +1,138 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_CUVS_DISTANCE_TYPE; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_INDEX_TYPE; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_INT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_NN_DESCENT_NUM_ITERATIONS; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_STRATEGY; +import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MAX_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MAX_INT_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MAX_NN_DESCENT_NUM_ITERATIONS; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MAX_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MIN_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MIN_INT_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MIN_NN_DESCENT_NUM_ITERATIONS; +import static com.nvidia.cuvs.lucene.GPUSearchParams.MIN_WRITER_THREADS; +import static java.lang.Integer.MAX_VALUE; +import static java.lang.Integer.MIN_VALUE; + +import java.util.Random; +import java.util.logging.Logger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestGPUSearchParams extends LuceneTestCase { + + @SuppressWarnings("unused") + private static final Logger log = Logger.getLogger(TestGPUSearchParams.class.getName()); + + private static Random random; + + @Test + public void testGPUSearchParamsDefaultValues() { + GPUSearchParams params = new GPUSearchParams.Builder().build(); + assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); + assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); + assertEquals(DEFAULT_WRITER_THREADS, params.getWriterThreads()); + assertEquals(DEFAULT_CAGRA_GRAPH_BUILD_ALGO, params.getCagraGraphBuildAlgo()); + assertEquals(DEFAULT_INDEX_TYPE, params.getIndexType()); + assertEquals(DEFAULT_STRATEGY, params.getStrategy()); + assertEquals(DEFAULT_CUVS_DISTANCE_TYPE, params.getCuvsDistanceType()); + assertEquals(DEFAULT_NN_DESCENT_NUM_ITERATIONS, params.getnNDescentNumIterations()); + } + + @Test + public void testGPUSearchParamsInvalidGraphDegree() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_GRAPH_DEG), random.nextInt(MAX_GRAPH_DEG + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withGraphDegree(v).build()); + } + } + + @Test + public void testGPUSearchParamsInvalidIntermediateGraphDegree() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_INT_GRAPH_DEG), + random.nextInt(MAX_INT_GRAPH_DEG + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withIntermediateGraphDegree(v).build()); + } + } + + @Test + public void testGPUSearchParamsInvalidWriterThreads() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_WRITER_THREADS), + random.nextInt(MAX_WRITER_THREADS + 1, MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withWriterThreads(v).build()); + } + } + + @Test + public void testGPUSearchParamsInvalidCagraGraphBuildAlgo() { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withCagraGraphBuildAlgo(null).build()); + } + + @Test + public void testGPUSearchParamsInvalidIndexType() { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withIndexType(null).build()); + } + + @Test + public void testGPUSearchParamsInvalidStrategy() { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withStrategy(null).build()); + } + + @Test + public void testGPUSearchParamsInvalidCuvsDistanceType() { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withCuvsDistanceType(null).build()); + } + + @Test + public void testGPUSearchParamsInvalidNNDescentNumIterations() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, (int) MIN_NN_DESCENT_NUM_ITERATIONS), + random.nextInt((int) (MAX_NN_DESCENT_NUM_ITERATIONS + 1), MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new GPUSearchParams.Builder().withNNDescentNumIterations(v).build()); + } + } + + @BeforeClass + public static void beforeClass() { + random = random(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java new file mode 100644 index 0000000000..39fc9b8c67 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.util.ArrayUtil.copyOfSubArray; + +import java.io.IOException; +import java.util.Random; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; + +@SuppressSysoutChecks(bugUrl = "") +public class TestIndexOutputOutputStream extends LuceneTestCase { + + public void testBasic() throws IOException { + try (var dir = newDirectory()) { + try (var indexOut = dir.createOutput("test", IOContext.DEFAULT)) { + var out = new IndexOutputOutputStream(indexOut); + out.write(0x56); + out.write(new byte[] {0x10, 0x11, 0x12, 0x13, 0x14}); + out.close(); + } + + try (var indexInput = dir.openInput("test", IOContext.DEFAULT); + var in = new IndexInputInputStream(indexInput)) { + byte[] ba = new byte[6]; + assertEquals(6, in.read(ba)); + assertArrayEquals(new byte[] {0x56, 0x10, 0x11, 0x12, 0x13, 0x14}, ba); + } + } + } + + public void testGetFilePointer() throws IOException { + try (var dir = newDirectory()) { + try (var indexOut = dir.createOutput("test", IOContext.DEFAULT)) { + var out = new IndexOutputOutputStream(indexOut); + out.write(0x56); + out.write(new byte[] {0x10, 0x11, 0x12}); + assertEquals(4, indexOut.getFilePointer()); + out.close(); + } + } + } + + public void testWithRandom() throws IOException { + byte[] data = new byte[Math.min(atLeast(10_000), 20_000)]; + Random random = random(); + random.nextBytes(data); + + try (var dir = newDirectory()) { + try (var indexOut = dir.createOutput("test", IOContext.DEFAULT)) { + var out = new IndexOutputOutputStream(indexOut); + int i = 0; + while (i < data.length) { + if (random.nextBoolean()) { + out.write(data[i]); + i++; + } else { + int numBytes = random.nextInt(Math.min(data.length - i, 100)); + out.write(data, i, numBytes); + i += numBytes; + } + } + out.close(); + } + + try (var indexInput = dir.openInput("test", IOContext.DEFAULT); + var in = new IndexInputInputStream(indexInput); ) { + int i = 0; + while (i < data.length) { + if (random.nextBoolean()) { + int b = in.read(); + assertEquals(data[i], b); + i++; + } else { + int numBytes = random.nextInt(Math.min(data.length - i, 100)); + byte[] ba = new byte[numBytes]; + in.read(ba, 0, numBytes); + assertArrayEquals(copyOfSubArray(data, i, i + numBytes), ba); + i += numBytes; + } + } + assertEquals(-1, in.read()); + assertEquals(-1, in.read(new byte[2])); + } + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java new file mode 100644 index 0000000000..9bb140228b --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.util.List; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Ignore; + +@SuppressSysoutChecks(bugUrl = "") +public class TestLucene99AcceleratedHNSWVectorsFormat extends BaseKnnVectorsFormatTestCase { + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS is not supported", isSupported()); + } + + @Override + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + } + + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { + float[][] f = new float[][] {randomVector(384), randomVector(384)}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc1 = new Document(); + doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); + w.addDocument(doc1); + w.commit(); + Document doc2 = new Document(); + doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); + w.addDocument(doc2); + w.flush(); + w.commit(); + + // sanity - verify one doc per leaf + try (DirectoryReader reader = DirectoryReader.open(w)) { + List subReaders = reader.leaves(); + assertEquals(2, subReaders.size()); + assertEquals(1, subReaders.get(0).reader().getFloatVectorValues("f").size()); + assertEquals(1, subReaders.get(1).reader().getFloatVectorValues("f").size()); + } + + // now merge to a single segment + w.forceMerge(1); + + // verify merged content + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues values = r.getFloatVectorValues("f"); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(f[0], values.vectorValue(0), 0.0f); + assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + } + } + } + + // Basic test for multiple vectors fields per document + public void testTwoVectorFieldsPerDoc() throws Exception { + float[][] f1 = new float[][] {randomVector(384), randomVector(384)}; + float[][] f2 = new float[][] {randomVector(384), randomVector(384)}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc1 = new Document(); + doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); + doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); + w.addDocument(doc1); + Document doc2 = new Document(); + doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); + doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); + w.addDocument(doc2); + w.forceMerge(1); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues values = r.getFloatVectorValues("f1"); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); + assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); + + values = r.getFloatVectorValues("f2"); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); + assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); + + // opportunistically check boundary condition - search with a 0 topK + var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); + assertEquals(0, topDocs.scoreDocs.length); + assertEquals(0, topDocs.totalHits.value()); + } + } + } + + @Override + // Overriding this method from superclass for the tests to only use float vector encoding + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testByteVectorScorerIteration() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testEmptyByteVectorData() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testMergingWithDifferentByteKnnFields() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testMismatchedFields() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testRandomBytes() {} + + @Ignore + @Override + // Ignoring this test from superclass as we do not support byte vectors + public void testSortedIndexBytes() {} +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java new file mode 100644 index 0000000000..a42f55b3fb --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -0,0 +1,1260 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.SortedDocValuesField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.SortedDocValues; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.Sort; +import org.apache.lucene.search.SortField; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.util.BytesRef; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * Comprehensive tests for merge functionality with CuVS indexes. + * Tests merge operations across different index types including brute force, + * CAGRA, and combined index configurations to ensure proper vector handling + * and search functionality after segment merging. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestMerge extends LuceneTestCase { + + private static final Logger log = Logger.getLogger(TestMerge.class.getName()); + + private static final int MIN_VECTOR_DIMENSION = 64; + private static final int MAX_VECTOR_DIMENSION = 256; + private static final int TOP_K_LIMIT = 64; + + private static CagraGraphBuildAlgo cagraGraphBuildAlgo; + + public TestMerge(@Name("cagraGraphBuildAlgo") CagraGraphBuildAlgo cagraGraphBuildAlgo) { + TestMerge.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + } + + @ParametersFactory + public static List parameters() { + return Arrays.asList( + new Object[][] {{CagraGraphBuildAlgo.NN_DESCENT}, {CagraGraphBuildAlgo.IVF_PQ}}); + } + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS is not supported", isSupported()); + } + + private Directory directory; + private int vectorDimension; + + @Before + public void setUp() throws Exception { + super.setUp(); + directory = newDirectory(); + + // Randomize vector dimension for each test + vectorDimension = + MIN_VECTOR_DIMENSION + random().nextInt(MAX_VECTOR_DIMENSION - MIN_VECTOR_DIMENSION + 1); + // Ensure dimension is multiple of 4 for better performance + vectorDimension = (vectorDimension / 4) * 4; + + log.log(Level.FINE, "Using randomized vector dimension: " + vectorDimension); + } + + @After + public void tearDown() throws Exception { + if (directory != null) { + directory.close(); + } + super.tearDown(); + } + + /** + * Test merging many documents across multiple segments + **/ + @Test + public void testMergeManyDocumentsMultipleSegments() throws IOException { + log.log( + Level.FINE, + "Starting testMergeManyDocumentsMultipleSegments with CagraGraphBuildAlgo: " + + cagraGraphBuildAlgo); + + // Randomize configuration parameters + int maxBufferedDocs = 5 + random().nextInt(16); // 5-20 docs per buffer + int totalBatches = 8 + random().nextInt(8); // 8-15 batches + int docsPerBatch = 15 + random().nextInt(11); // 15-25 docs per batch + int totalDocuments = totalBatches * docsPerBatch; + + // Randomize vector presence probability (60-85%) + double vectorProbability = 0.6 + (random().nextDouble() * 0.25); + + log.log( + Level.FINE, + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", totalBatches=" + + totalBatches + + ", docsPerBatch=" + + docsPerBatch + + ", totalDocuments=" + + totalDocuments + + ", vectorProbability=" + + vectorProbability); + + GPUSearchParams params = + new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); + + CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(format)) + .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + List expectedVectors = new ArrayList<>(); + List expectedDocIds = new ArrayList<>(); + int documentsWithVectors = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Add documents in multiple batches to create many segments + for (int batch = 0; batch < totalBatches; batch++) { + for (int i = 0; i < docsPerBatch; i++) { + int docId = batch * docsPerBatch + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new NumericDocValuesField("batch", batch)); + + // Randomly decide if document has vector + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + expectedVectors.add(vector); + expectedDocIds.add(docId); + documentsWithVectors++; + } + + writer.addDocument(doc); + } + writer.commit(); // Create a new segment + } + + int documentsWithoutVectors = totalDocuments - documentsWithVectors; + log.log( + Level.FINE, "Created " + totalDocuments + " documents in " + totalBatches + " segments"); + log.log(Level.FINE, "Documents with vectors: " + documentsWithVectors); + log.log(Level.FINE, "Documents without vectors: " + documentsWithoutVectors); + + // Force merge to trigger merge logic + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge to single segment completed"); + } + + // Verify the merged index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Verify vector search works correctly after merge + if (documentsWithVectors > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Randomize search parameters + int searchK = + Math.min(5 + random().nextInt(10), Math.min(documentsWithVectors, TOP_K_LIMIT)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs results = searcher.search(query, searchK); + + assertTrue("Should find some results after merge", results.scoreDocs.length > 0); + assertTrue( + "Should find reasonable number of results", + results.scoreDocs.length <= documentsWithVectors); + + log.log( + Level.FINE, + "Vector search returned " + + results.scoreDocs.length + + " results out of " + + documentsWithVectors + + " documents with vectors"); + + // Verify all returned documents have valid IDs + for (ScoreDoc scoreDoc : results.scoreDocs) { + int docId = Integer.parseInt(searcher.storedFields().document(scoreDoc.doc).get("id")); + assertTrue("Document ID should be valid", docId >= 0 && docId < totalDocuments); + } + } else { + log.log(Level.FINE, "No documents with vectors - skipping vector search verification"); + } + + log.log(Level.FINE, "Merge verification completed successfully"); + } + } + + /** + * Test merging with index sorting enabled using text-based sorting and SortingMergePolicy + **/ + @Test + public void testMergeWithIndexSorting() throws IOException { + log.log( + Level.FINE, + "Starting testMergeWithIndexSorting with text-based sorting with CagraGraphBuildAlgo: " + + cagraGraphBuildAlgo); + + // Randomize sort field type + SortField.Type sortType = random().nextBoolean() ? SortField.Type.STRING : SortField.Type.LONG; + String sortFieldName = sortType == SortField.Type.STRING ? "text_sort_key" : "numeric_sort_key"; + + // Configure index sorting by a randomized field + Sort indexSort = new Sort(new SortField(sortFieldName, sortType)); + + // Randomize merge policy parameters + TieredMergePolicy mergePolicy = new TieredMergePolicy(); + mergePolicy.setMaxMergedSegmentMB(128 + random().nextInt(257)); // 128-384 MB + mergePolicy.setSegmentsPerTier(3 + random().nextInt(4)); // 3-6 segments per tier + + // Randomize writer configuration parameters + int maxBufferedDocs = 10 + random().nextInt(16); // 10-25 docs per buffer + int totalDocuments = 80 + random().nextInt(81); // 80-160 documents + int segmentSize = 15 + random().nextInt(11); // 15-25 docs per segment + double vectorProbability = 0.65 + (random().nextDouble() * 0.25); // 65-90% have vectors + + log.log( + Level.FINE, + "Randomized sorting parameters: sortType=" + sortType + ", sortFieldName=" + sortFieldName); + log.log( + Level.FINE, + "Randomized config: maxBufferedDocs=" + + maxBufferedDocs + + ", totalDocuments=" + + totalDocuments + + ", segmentSize=" + + segmentSize + + ", vectorProbability=" + + vectorProbability); + + GPUSearchParams params = + new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); + + CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(format)) + .setIndexSort(indexSort) // This automatically enables sorting during merges + .setMergePolicy(mergePolicy) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + // List documents = new ArrayList<>(); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create documents with randomized sort keys + for (int i = 0; i < totalDocuments; i++) { + float[] vector = null; + + // Randomly decide if document has vector + if (random().nextDouble() < vectorProbability) { + vector = generateRandomVector(vectorDimension, random()); + } + + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(new StringField("original_order", String.valueOf(i), Field.Store.YES)); + + // Add sort field based on randomized type + if (sortType == SortField.Type.STRING) { + // Randomize text sort key length (4-12 characters) + int keyLength = 4 + random().nextInt(9); + String textSortKey = generateRandomText(random(), keyLength); + doc.add(new SortedDocValuesField(sortFieldName, new BytesRef(textSortKey))); + doc.add(new StringField(sortFieldName + "_stored", textSortKey, Field.Store.YES)); + } else { + // Use numeric sort key with wider range + long numericSortKey = random().nextLong() % 100000; // Can be negative for more variety + doc.add(new NumericDocValuesField(sortFieldName, numericSortKey)); + doc.add( + new StringField( + sortFieldName + "_stored", String.valueOf(numericSortKey), Field.Store.YES)); + } + + if (vector != null) { + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + } + + writer.addDocument(doc); + + // Commit based on randomized segment size + if ((i + 1) % segmentSize == 0) { + writer.commit(); + log.log( + Level.FINE, + "Committed segment " + + ((i + 1) / segmentSize) + + " with " + + (i + 1) + + " total documents"); + } + } + + log.log(Level.FINE, "Created " + totalDocuments + " documents with text-based index sorting"); + + // Force merge with sorting - this will use the sorting merge policy + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge with text-based sorting completed"); + } + + // Verify the merged and sorted index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Verify documents are sorted correctly by the randomized sort field + log.log( + Level.FINE, + "Verifying document sorting order using sortType: " + + sortType + + ", field: " + + sortFieldName); + + if (sortType == SortField.Type.STRING) { + // Verify string-based sorting + String previousSortKey = ""; + SortedDocValues sortedValues = leafReader.getSortedDocValues(sortFieldName); + + for (int docId = 0; docId < leafReader.maxDoc(); docId++) { + String currentSortKey = ""; + if (sortedValues != null && sortedValues.advanceExact(docId)) { + currentSortKey = sortedValues.lookupOrd(sortedValues.ordValue()).utf8ToString(); + } + + assertTrue( + "Documents should be sorted by " + + sortFieldName + + ": '" + + previousSortKey + + "' should be <= '" + + currentSortKey + + "'", + previousSortKey.compareTo(currentSortKey) <= 0); + previousSortKey = currentSortKey; + + // Log first 10 documents to verify sorting + if (docId < 10) { + IndexSearcher searcher = new IndexSearcher(reader); + String originalOrder = searcher.storedFields().document(docId).get("original_order"); + log.log( + Level.FINE, + "DocId: " + + docId + + ", OriginalOrder: " + + originalOrder + + ", SortKey: '" + + currentSortKey + + "'"); + } + } + } else { + // Verify numeric-based sorting + long previousSortKey = Long.MIN_VALUE; + var numericValues = leafReader.getNumericDocValues(sortFieldName); + + for (int docId = 0; docId < leafReader.maxDoc(); docId++) { + long currentSortKey = Long.MIN_VALUE; + if (numericValues != null && numericValues.advanceExact(docId)) { + currentSortKey = numericValues.longValue(); + } + + assertTrue( + "Documents should be sorted by " + + sortFieldName + + ": " + + previousSortKey + + " should be <= " + + currentSortKey, + previousSortKey <= currentSortKey); + previousSortKey = currentSortKey; + + // Log first 10 documents to verify sorting + if (docId < 10) { + IndexSearcher searcher = new IndexSearcher(reader); + String originalOrder = searcher.storedFields().document(docId).get("original_order"); + log.log( + Level.FINE, + "DocId: " + + docId + + ", OriginalOrder: " + + originalOrder + + ", SortKey: " + + currentSortKey); + } + } + } + + // Count total vectors by checking if vector field exists and has values + var vectorValues = leafReader.getFloatVectorValues("vector"); + int documentsWithVectors = vectorValues != null ? vectorValues.size() : 0; + + log.log( + Level.FINE, + "Found " + documentsWithVectors + " documents with vectors after sorted merge"); + + // Test vector search on sorted index + if (documentsWithVectors > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + KnnFloatVectorQuery query = + new KnnFloatVectorQuery("vector", queryVector, Math.min(10, documentsWithVectors)); + TopDocs results = searcher.search(query, 10); + + assertTrue("Should find results in sorted index", results.scoreDocs.length > 0); + log.log( + Level.FINE, + "Vector search on sorted index returned " + results.scoreDocs.length + " results"); + + // Verify that returned documents maintain sort order if we check their sort keys + log.log(Level.FINE, "Verifying vector search results maintain sorting consistency..."); + for (int i = 0; i < Math.min(3, results.scoreDocs.length); i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + String originalOrder = + searcher.storedFields().document(scoreDoc.doc).get("original_order"); + String sortKey = + searcher.storedFields().document(scoreDoc.doc).get(sortFieldName + "_stored"); + log.log( + Level.FINE, + "Result " + + i + + ": DocId=" + + scoreDoc.doc + + ", OriginalOrder=" + + originalOrder + + ", SortKey='" + + sortKey + + "', Score=" + + scoreDoc.score); + } + } + + log.log(Level.FINE, "Text-based index sorting verification completed successfully"); + } + } + + /** + * Test merging segments with various patterns of missing vectors + **/ + @Test + public void testMergeWithMissingVectors() throws IOException { + log.log( + Level.FINE, + "Starting testMergeWithMissingVectors with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); + + // Randomize configuration + int maxBufferedDocs = 10 + random().nextInt(11); // 10-20 docs per buffer + int numSegments = 3 + random().nextInt(3); // 3-5 segments + + GPUSearchParams params = + new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); + + CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(format)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + log.log( + Level.FINE, + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments); + + int totalExpectedVectors = 0; + int totalDocuments = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int seg = 0; seg < numSegments; seg++) { + // Randomize segment characteristics + int docsInSegment = 15 + random().nextInt(16); // 15-30 docs per segment + double vectorProbability = random().nextDouble(); // 0-100% vector probability + String segmentType = "seg_" + seg + "_prob_" + String.format("%.2f", vectorProbability); + + int segmentVectorCount = 0; + + for (int i = 0; i < docsInSegment; i++) { + Document doc = new Document(); + doc.add(new StringField("id", "seg" + seg + "_" + i, Field.Store.YES)); + doc.add(new StringField("segment", segmentType, Field.Store.YES)); + doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new NumericDocValuesField("doc_in_segment", i)); + + // Randomly add vector based on segment's probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + segmentVectorCount++; + } + + writer.addDocument(doc); + } + + writer.commit(); + totalDocuments += docsInSegment; + totalExpectedVectors += segmentVectorCount; + + log.log( + Level.FINE, + "Created segment " + + seg + + ": " + + docsInSegment + + " documents, " + + segmentVectorCount + + " with vectors (probability: " + + String.format("%.2f", vectorProbability) + + ")"); + } + + // Force merge all segments + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge of " + numSegments + " segments completed"); + } + + // Verify the merged index handles missing vectors correctly + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + log.log( + Level.FINE, + "Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); + + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); + + // Test vector search if we have vectors + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Randomize search parameters + int searchK = Math.min(5 + random().nextInt(10), Math.min(actualVectorCount, TOP_K_LIMIT)); + + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + + assertTrue("Should find some vector results", vectorResults.scoreDocs.length > 0); + assertTrue( + "Should not find more vectors than exist", + vectorResults.scoreDocs.length <= actualVectorCount); + + log.log( + Level.FINE, + "Found " + + vectorResults.scoreDocs.length + + " vector results out of " + + actualVectorCount + + " available"); + } else { + log.log(Level.FINE, "No vectors in merged index - skipping vector search"); + } + + log.log(Level.FINE, "Missing vectors test completed successfully"); + } + } + + /** + * Test merge behavior with document deletions + **/ + @Test + public void testMergeWithDeletions() throws IOException { + log.log( + Level.FINE, + "Starting testMergeWithDeletions with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); + + // Randomize configuration parameters + int maxBufferedDocs = 15 + random().nextInt(11); // 15-25 docs per buffer + int numSegments = 3 + random().nextInt(4); // 3-6 segments + int docsPerSegment = 20 + random().nextInt(21); // 20-40 docs per segment + double vectorProbability = 0.7 + (random().nextDouble() * 0.25); // 70-95% have vectors + double deletionProbability = 0.2 + (random().nextDouble() * 0.3); // 20-50% deletion rate + + log.log( + Level.FINE, + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments + + ", docsPerSegment=" + + docsPerSegment + + ", vectorProbability=" + + vectorProbability + + ", deletionProbability=" + + deletionProbability); + + GPUSearchParams params = + new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); + + CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(format)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + List expectedRemainingDocs = new ArrayList<>(); + List deletedDocs = new ArrayList<>(); + int totalDocuments = numSegments * docsPerSegment; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create multiple segments with documents + for (int seg = 0; seg < numSegments; seg++) { + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new StringField("segment", "seg_" + seg, Field.Store.YES)); + doc.add(new NumericDocValuesField("doc_num", docId)); + doc.add(new NumericDocValuesField("segment_num", seg)); + + // Randomly add vectors + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + } + + writer.addDocument(doc); + } + writer.commit(); + } + + log.log( + Level.FINE, + "Created " + + numSegments + + " segments with " + + docsPerSegment + + " documents each (" + + totalDocuments + + " total)"); + + // Delete documents randomly and track which ones are deleted + int deletedCount = 0; + for (int docId = 0; docId < totalDocuments; docId++) { + if (random().nextDouble() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(docId))); + deletedDocs.add(docId); + deletedCount++; + } else { + expectedRemainingDocs.add(docId); + } + } + + log.log( + Level.FINE, + "Deleted " + + deletedCount + + " documents (" + + String.format("%.1f", (100.0 * deletedCount / totalDocuments)) + + "%), remaining: " + + expectedRemainingDocs.size()); + + writer.commit(); + + // Force merge to apply deletions + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge with deletions completed"); + } + + // Verify the merged index correctly handles deletions + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + int expectedRemaining = expectedRemainingDocs.size(); + assertEquals( + "Should have correct number of documents after deletions", + expectedRemaining, + leafReader.maxDoc()); + + // Verify that deleted documents are not present + IndexSearcher searcher = new IndexSearcher(reader); + + // Test that we can find expected remaining documents + for (int i = 0; i < Math.min(10, expectedRemainingDocs.size()); i++) { + int docId = expectedRemainingDocs.get(i); + TopDocs result = searcher.search(new TermQuery(new Term("id", String.valueOf(docId))), 1); + assertEquals("Should find remaining document " + docId, 1, (int) result.totalHits.value()); + } + + // Test that actually deleted documents are not found + int deletedDocsToCheck = Math.min(10, deletedDocs.size()); // Check up to 10 deleted docs + for (int i = 0; i < deletedDocsToCheck; i++) { + int docId = deletedDocs.get(i); + TopDocs result = searcher.search(new TermQuery(new Term("id", String.valueOf(docId))), 1); + assertEquals( + "Should not find deleted document " + docId, 0, (int) result.totalHits.value()); + } + + // Test vector search works after deletions + float[] queryVector = generateRandomVector(vectorDimension, random()); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, 10); + TopDocs vectorResults = searcher.search(vectorQuery, 10); + + assertTrue( + "Should find some vector results after deletions", vectorResults.scoreDocs.length > 0); + + log.log( + Level.FINE, + "Found " + vectorResults.scoreDocs.length + " vector results after deletions"); + log.log(Level.FINE, "Deletion merge verification completed successfully"); + } + } + + /** + * Test merging segments for {@link IndexType#BRUTE_FORCE} + * */ + @Test + public void testMergeBruteForceIndex() throws IOException { + log.log( + Level.FINE, + "Starting testMergeBruteForceIndex with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); + + // Randomize configuration parameters + int maxBufferedDocs = 8 + random().nextInt(8); // 8-15 docs per buffer + int numSegments = 3 + random().nextInt(3); // 3-5 segments + int docsPerSegment = 12 + random().nextInt(9); // 12-20 docs per segment + double vectorProbability = 0.8 + (random().nextDouble() * 0.2); // 80-100% have vectors + + log.log( + Level.FINE, + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments + + ", docsPerSegment=" + + docsPerSegment + + ", vectorProbability=" + + vectorProbability); + + GPUSearchParams params = + new GPUSearchParams.Builder() + .withCagraGraphBuildAlgo(cagraGraphBuildAlgo) + .withIndexType(IndexType.BRUTE_FORCE) + .build(); + + // Configure with brute force index type + CuVS2510GPUVectorsFormat bruteForceFormat = + new CuVS2510GPUVectorsFormat(params); // Use brute force index + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(bruteForceFormat)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + int totalDocuments = numSegments * docsPerSegment; + int totalExpectedVectors = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create multiple segments with brute force index + for (int seg = 0; seg < numSegments; seg++) { + int segmentVectorCount = 0; + + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new StringField("segment", "seg_" + seg, Field.Store.YES)); + doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new NumericDocValuesField("doc_in_segment", i)); + + // Randomly add vectors based on probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + segmentVectorCount++; + } + + writer.addDocument(doc); + } + + writer.commit(); + totalExpectedVectors += segmentVectorCount; + + log.log( + Level.FINE, + "Created brute force segment " + + seg + + ": " + + docsPerSegment + + " documents, " + + segmentVectorCount + + " with vectors"); + } + + log.log( + Level.FINE, + "Created " + + numSegments + + " brute force segments with " + + totalDocuments + + " total documents and " + + totalExpectedVectors + + " vectors"); + + // Force merge all brute force segments + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge of brute force segments completed"); + } + + // Verify the merged brute force index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + log.log( + Level.FINE, + "Brute force merge results: Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); + + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); + + // Test brute force vector search (exact search) + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Search for reasonable number of results + int searchK = Math.min(8 + random().nextInt(8), Math.min(actualVectorCount, TOP_K_LIMIT)); + + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + + assertTrue( + "Should find some vector results in brute force index", + vectorResults.scoreDocs.length > 0); + assertTrue( + "Should not find more vectors than exist", + vectorResults.scoreDocs.length <= actualVectorCount); + + log.log( + Level.FINE, + "Brute force search found " + + vectorResults.scoreDocs.length + + " results out of " + + actualVectorCount + + " available vectors"); + + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + String docId = searcher.storedFields().document(scoreDoc.doc).get("id"); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); + } + } else { + log.log(Level.FINE, "No vectors in brute force merged index - skipping vector search"); + } + + log.log(Level.FINE, "Brute force merge verification completed successfully"); + } + } + + /** + * Test merging segments for {@link IndexType#CAGRA_AND_BRUTE_FORCE} + * */ + @Test + public void testMergeCagraAndBruteForceIndex() throws IOException { + log.log( + Level.FINE, + "Starting testMergeCagraAndBruteForceIndex with CagraGraphBuildAlgo: " + + cagraGraphBuildAlgo); + + // Use moderate dataset size + int maxBufferedDocs = 15 + random().nextInt(10); // 15-24 docs per buffer + int numSegments = + 4; // Fixed 4 segments: alternating CAGRA vs small segments (brute force fallback) + int docsPerSegment = 20 + random().nextInt(11); // 20-30 docs per segment + double vectorProbability = 0.9 + (random().nextDouble() * 0.1); // 90-100% have vectors + + log.log( + Level.FINE, + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments + + ", docsPerSegment=" + + docsPerSegment + + ", vectorProbability=" + + vectorProbability); + + // Configure with CAGRA + brute force combined index type + GPUSearchParams params = + new GPUSearchParams.Builder() + .withCagraGraphBuildAlgo(cagraGraphBuildAlgo) + .withIndexType(IndexType.CAGRA_AND_BRUTE_FORCE) + .build(); + + CuVS2510GPUVectorsFormat combinedFormat = + new CuVS2510GPUVectorsFormat(params); // Use combined CAGRA + brute force + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(combinedFormat)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + int totalDocuments = numSegments * docsPerSegment; + int totalExpectedVectors = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create segments that will result in mixed index types during merge + for (int seg = 0; seg < numSegments; seg++) { + int segmentVectorCount = 0; + + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new StringField("segment", "mixed_seg_" + seg, Field.Store.YES)); + doc.add(new StringField("index_type", "cagra_and_brute_force", Field.Store.YES)); + doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new NumericDocValuesField("doc_in_segment", i)); + + // Add vectors based on probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + segmentVectorCount++; + } + + writer.addDocument(doc); + } + + writer.commit(); + totalExpectedVectors += segmentVectorCount; + + log.log( + Level.FINE, + "Created CAGRA+brute force segment " + + seg + + ": " + + docsPerSegment + + " documents, " + + segmentVectorCount + + " with vectors"); + } + + log.log( + Level.FINE, + "Created " + + numSegments + + " CAGRA+brute force segments with " + + totalDocuments + + " total documents and " + + totalExpectedVectors + + " vectors"); + + // Force merge all CAGRA+brute force segments + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge of CAGRA+brute force segments completed"); + } + + // Verify the merged CAGRA+brute force index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + log.log( + Level.FINE, + "CAGRA+brute force merge results: Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); + + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); + + // Test CAGRA+brute force index vector search + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Search for reasonable number of results + int searchK = Math.min(12 + random().nextInt(8), Math.min(actualVectorCount, TOP_K_LIMIT)); + + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + + assertTrue( + "Should find some vector results in CAGRA+brute force index", + vectorResults.scoreDocs.length > 0); + assertTrue( + "Should not find more vectors than exist", + vectorResults.scoreDocs.length <= actualVectorCount); + + log.log( + Level.FINE, + "CAGRA+brute force index search found " + + vectorResults.scoreDocs.length + + " results out of " + + actualVectorCount + + " available vectors"); + + // Verify all returned documents are valid and have expected metadata + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document resultDoc = searcher.storedFields().document(scoreDoc.doc); + String docId = resultDoc.get("id"); + String indexType = resultDoc.get("index_type"); + + assertNotNull("Document should have valid ID", docId); + assertEquals( + "Document should be marked as CAGRA+brute force index type", + "cagra_and_brute_force", + indexType); + assertTrue("Score should be positive", scoreDoc.score > 0); + } + + // Test that the CAGRA+brute force index handles both approximate and exact search + // consistently + for (int trial = 0; trial < 3; trial++) { + float[] trialQueryVector = generateRandomVector(vectorDimension, random()); + KnnFloatVectorQuery trialQuery = + new KnnFloatVectorQuery("vector", trialQueryVector, Math.min(5, actualVectorCount)); + TopDocs trialResults = searcher.search(trialQuery, Math.min(5, actualVectorCount)); + + assertTrue("Trial " + trial + " should find results", trialResults.scoreDocs.length > 0); + log.log( + Level.FINE, + "Trial " + trial + " found " + trialResults.scoreDocs.length + " results"); + } + } else { + log.log( + Level.FINE, "No vectors in CAGRA+brute force merged index - skipping vector search"); + } + + log.log(Level.FINE, "CAGRA+brute force merge verification completed successfully"); + } + } + + /** + * Test large scale merge to stress test the system + **/ + @Test + public void testLargeScaleMerge() throws IOException { + assumeTrue( + "testLargeScaleMerge requires -DlargeScale=true", + Boolean.parseBoolean(System.getProperty("largeScale", "false"))); + + log.log(Level.FINE, "Starting testLargeScaleMerge"); + + // Randomize large scale parameters + int maxBufferedDocs = 40 + random().nextInt(21); // 40-60 docs per buffer + int segmentCount = 15 + random().nextInt(11); // 15-25 segments + int docsPerSegment = 30 + random().nextInt(21); // 30-50 docs per segment + int totalDocuments = segmentCount * docsPerSegment; + + log.log( + Level.FINE, + "Randomized large scale parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", segmentCount=" + + segmentCount + + ", docsPerSegment=" + + docsPerSegment + + ", totalDocuments=" + + totalDocuments); + + GPUSearchParams params = + new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); + + CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(format)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int seg = 0; seg < segmentCount; seg++) { + log.log(Level.FINE, "Creating segment " + (seg + 1) + "/" + segmentCount); + + // Randomize vector probability per segment + double vectorProbability = + 0.5 + (random().nextDouble() * 0.4); // 50-90% vectors per segment + + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new NumericDocValuesField("segment", seg)); + doc.add(new NumericDocValuesField("position", i)); + + // Add vector based on segment's randomized probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + } + + writer.addDocument(doc); + } + writer.commit(); + } + + log.log( + Level.FINE, + "Created " + segmentCount + " segments with " + totalDocuments + " total documents"); + + // Force merge all segments + long startTime = System.currentTimeMillis(); + writer.forceMerge(1); + long mergeTime = System.currentTimeMillis() - startTime; + + log.log(Level.FINE, "Large scale merge completed in " + mergeTime + "ms"); + } + + // Verify the large merged index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Test vector search performance + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Randomize search parameters for large scale test + int searchK = + Math.min(20 + random().nextInt(31), Math.min(actualVectorCount, TOP_K_LIMIT)); // 20-50 + + long searchStart = System.currentTimeMillis(); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + long searchTime = System.currentTimeMillis() - searchStart; + + assertTrue("Should find vector results in large index", vectorResults.scoreDocs.length > 0); + log.log( + Level.FINE, + "Vector search in large index returned " + + vectorResults.scoreDocs.length + + " results out of " + + actualVectorCount + + " vectors in " + + searchTime + + "ms"); + } else { + log.log(Level.FINE, "No vectors in large merged index - skipping vector search"); + } + + log.log(Level.FINE, "Large scale merge verification completed successfully"); + } + } + + /** Helper method to generate random vectors */ + private float[] generateRandomVector(int dimension, Random random) { + float[] vector = new float[dimension]; + for (int i = 0; i < dimension; i++) { + vector[i] = (float) random().nextGaussian(); + } + // Normalize the vector + float norm = 0.0f; + for (float v : vector) { + norm += v * v; + } + norm = (float) Math.sqrt(norm); + if (norm > 0) { + for (int i = 0; i < dimension; i++) { + vector[i] /= norm; + } + } + return vector; + } + + /** Helper method to generate random text strings for sorting */ + private String generateRandomText(Random random, int length) { + StringBuilder sb = new StringBuilder(length); + String chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; + for (int i = 0; i < length; i++) { + sb.append(chars.charAt(random().nextInt(chars.length()))); + } + return sb.toString(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java new file mode 100644 index 0000000000..914fa963a8 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java @@ -0,0 +1,141 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Random; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestMultithreadedCuVSGPUSearch extends LuceneTestCase { + + private static final Logger log = + Logger.getLogger(TestMultithreadedCuVSGPUSearch.class.getName()); + private static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + private static final String VECTOR_FIELD = "vectors"; + + private static Directory directory; + private static Random random; + private static BlockingQueue queries; + private static int numQueries; + private static int topK; + private static int numThreads; + + @BeforeClass + public static void beforeClass() throws IOException { + assumeTrue("cuVS not supported", isSupported()); + random = random(); + directory = newDirectory(new ByteBuffersDirectory()); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec); + IndexWriter writer = new IndexWriter(directory, config); + + int datasetSize = random.nextInt(500, 2000); + int dimensions = random.nextInt(64, 256); + topK = random.nextInt(2, 30); + log.log(Level.FINE, "Using topK as: " + topK); + numThreads = random.nextInt(2, 8); + log.log(Level.FINE, "Generating a dataset with " + datasetSize + " vectors"); + float[][] dataset = generateDataset(random, datasetSize, dimensions); + numQueries = random.nextInt(100, 500); + log.log(Level.FINE, "Generating a query set with " + numQueries + " queries"); + float[][] queryVectors = generateDataset(random, numQueries, dimensions); + queries = new ArrayBlockingQueue<>(numQueries, true, Arrays.asList(queryVectors)); + + log.log(Level.FINE, "Indexing " + datasetSize + " vectors"); + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(newTextField("text_field", English.intToEnglish(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + writer.close(); + } + + @Test + public void testMultithreadedCuVSGPUSearch() throws Exception { + DirectoryReader reader = DirectoryReader.open(directory); + ExecutorService executorService = Executors.newFixedThreadPool(numThreads); + IndexSearcher searcher = new IndexSearcher(reader); + CountDownLatch latch = new CountDownLatch(numThreads); + AtomicInteger totalSuccessfulQueries = new AtomicInteger(); + log.log(Level.FINE, "Using " + numThreads + " threads"); + for (int i = 0; i < numThreads; i++) { + executorService.execute( + new Runnable() { + public void run() { + try { + float[] queryVector; + String threadName = Thread.currentThread().getName(); + while ((queryVector = queries.poll()) != null) { + log.log(Level.FINE, "Thread: " + threadName + ", queue size: " + queries.size()); + log.log(Level.FINER, "Query: " + Arrays.toString(queryVector)); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK, null, topK, 1); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + totalSuccessfulQueries.addAndGet(hits.length == topK ? 1 : 0); + } + } catch (Exception e) { + e.printStackTrace(); + } finally { + latch.countDown(); + } + } + }); + } + latch.await(); + executorService.shutdown(); + executorService.close(); + reader.close(); + log.log( + Level.FINE, + "Number queries that returned topK values: " + totalSuccessfulQueries.intValue()); + assertEquals( + "All search queries did not return topK results", + totalSuccessfulQueries.intValue(), + numQueries); + } + + @AfterClass + public static void afterClass() throws IOException { + if (directory != null) directory.close(); + directory = null; + log.log(Level.FINE, "Test finished"); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java new file mode 100644 index 0000000000..5140782fc7 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -0,0 +1,212 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.COSINE; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import java.util.Arrays; +import java.util.List; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Ignore; + +@SuppressSysoutChecks(bugUrl = "") +public class TestQuantizedVectorsFormats extends BaseKnnVectorsFormatTestCase { + + private static final Logger log = Logger.getLogger(TestQuantizedVectorsFormats.class.getName()); + + private static KnnVectorsFormat knnVectorsFormat; + + public TestQuantizedVectorsFormats(@Name("knnVectorsWriter") KnnVectorsFormat knnVectorsFormat) { + TestQuantizedVectorsFormats.knnVectorsFormat = knnVectorsFormat; + } + + @ParametersFactory + public static List parameters() { + return Arrays.asList( + new Object[][] { + {new LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat()}, + {new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat()} + }); + } + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS is not supported so skipping these tests", isSupported()); + } + + @Override + protected Codec getCodec() { + log.log(Level.FINE, "Running tests for: " + knnVectorsFormat.getName()); + return TestUtil.alwaysKnnVectorsFormat(knnVectorsFormat); + } + + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { + final int R = 2, D = 128; + float[][] f = new float[R][D]; + final String F = "f"; + for (int i = 0; i < R; i++) { + f[i] = randomVector(D); + } + + try (Directory dir = newDirectory(new ByteBuffersDirectory()); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + for (int i = 0; i < R; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new KnnFloatVectorField(F, f[i], EUCLIDEAN)); + w.addDocument(doc); + w.commit(); + } + w.flush(); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + List subReaders = reader.leaves(); + assertEquals(2, subReaders.size()); + for (int i = 0; i < R; i++) { + assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(F).size()); + } + } + w.forceMerge(1); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues values = r.getFloatVectorValues(F); + assertNotNull(values); + assertEquals(R, values.size()); + for (int i = 0; i < R; i++) { + assertArrayEquals(f[i], values.vectorValue(i), 0.0f); + } + } + } + } + + public void testTwoVectorFieldsPerDoc() throws Exception { + final int R = 2, D = 128; + final String F1 = "f1", F2 = "f2"; + float[][] f1 = new float[R][D]; + float[][] f2 = new float[R][D]; + + for (int i = 0; i < R; i++) { + f1[i] = randomVector(D); + f2[i] = randomVector(D); + } + + try (Directory dir = newDirectory(new ByteBuffersDirectory()); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + + for (int i = 0; i < R; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new KnnFloatVectorField(F1, f1[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(F2, f2[i], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues values = r.getFloatVectorValues(F1); + assertNotNull(values); + assertEquals(R, values.size()); + for (int i = 0; i < R; i++) { + assertArrayEquals(f1[i], values.vectorValue(i), 0.0f); + } + + values = r.getFloatVectorValues(F2); + assertNotNull(values); + assertEquals(R, values.size()); + for (int i = 0; i < R; i++) { + assertArrayEquals(f2[i], values.vectorValue(i), 0.0f); + } + } + } + } + + public void testCosineSimilarity() throws Exception { + final int R = 2, D = 128; + final String F = "f"; + float[][] f = new float[R][D]; + for (int i = 0; i < R; i++) { + f[i] = randomVector(D); + } + + try (Directory dir = newDirectory(new ByteBuffersDirectory()); + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + + for (int i = 0; i < R; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new KnnFloatVectorField(F, f[i], COSINE)); + w.addDocument(doc); + } + w.forceMerge(1); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + + FloatVectorValues values = r.getFloatVectorValues(F); + assertNotNull(values); + assertEquals(R, values.size()); + + float[] queryVector = randomVector(D); + var topDocs = r.searchNearestVectors(F, queryVector, 2, null, 10); + assertTrue("Should return at least one result", topDocs.scoreDocs.length > 0); + assertTrue("Scores should be non-negative", topDocs.scoreDocs[0].score >= 0); + } + } + } + + @Override + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Ignore + @Override + public void testByteVectorScorerIteration() {} + + @Ignore + @Override + public void testEmptyByteVectorData() {} + + @Ignore + @Override + public void testMergingWithDifferentByteKnnFields() {} + + @Ignore + @Override + public void testMismatchedFields() {} + + @Ignore + @Override + public void testRandomBytes() {} + + @Ignore + @Override + public void testSortedIndexBytes() {} +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java new file mode 100644 index 0000000000..88ca3150fd --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.Random; + +public class TestUtils { + + public static float[][] generateDataset(Random random, int size, int dimensions) { + float[][] dataset = new float[size][dimensions]; + for (int i = 0; i < size; i++) { + for (int j = 0; j < dimensions; j++) { + dataset[i][j] = random.nextFloat() * 100; + } + } + return dataset; + } + + public static float[] generateRandomVector(int dimensions, Random random) { + float[] vector = new float[dimensions]; + for (int i = 0; i < dimensions; i++) { + vector[i] = random.nextFloat() * 100; + } + return vector; + } + + public static float[][] generateQueries(Random random, int dimensions, int numQueries) { + // Generate random query vectors + float[][] queries = new float[numQueries][dimensions]; + for (int i = 0; i < numQueries; i++) { + for (int j = 0; j < dimensions; j++) { + queries[i][j] = random.nextFloat() * 100; + } + } + return queries; + } +} diff --git a/pyproject.toml b/pyproject.toml index 50755ed387..121866a8e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 [tool.ruff] @@ -47,6 +47,6 @@ follow_imports = "skip" skip = "./.git,./.github,./cpp/build,.*egg-info.*,./.mypy_cache,.*_skbuild" # ignore short words, and typename parameters like OffsetT ignore-regex = "\\b(.{1,4}|[A-Z]\\w*T)\\b" -ignore-words-list = "inout,unparseable,numer,rocess,thirdparty,couldn,subtile" +ignore-words-list = "inout,unparseable,numer,rocess,thirdparty,couldn,subtile,expectd,atleast" builtin = "clear" quiet-level = 3