diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 6c18e9f97a..af13b353e7 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -61,34 +61,54 @@ jobs: contents: read uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 with: - build_type: pull-request + build_type: ${{ inputs.build_type || 'branch' }} matrix_name: conda-cpp-build rocky8-clib-standalone-build: - needs: [build-details, rocky8-clib-standalone-build-matrix] + needs: [rocky8-clib-standalone-build-matrix] + runs-on: linux-${{ matrix.ARCH }}-cpu16 + env: + RAPIDS_BUILD_TYPE: ${{ inputs.build_type || 'branch' }} 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@release/26.10 strategy: fail-fast: false matrix: ${{ fromJSON(needs.rocky8-clib-standalone-build-matrix.outputs.matrix) }} - with: - build_type: ${{ inputs.build_type || 'branch' }} - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - branch: ${{ inputs.branch }} - arch: "${{matrix.ARCH}}" - date: ${{ inputs.date }} - container_image: "rapidsai/ci-wheel:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" - node_type: "cpu16" - requires_license_builder: true - script: "ci/build_standalone_c.sh" - artifact-name: "libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz" - file_to_upload: "libcuvs_c.tar.gz" - sha: ${{ inputs.sha }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + ref: ${{ inputs.sha || github.sha }} + fetch-depth: 0 + persist-credentials: false + - name: Standardize repository information + uses: rapidsai/shared-actions/rapids-github-info@main + with: + branch: ${{ inputs.branch }} + date: ${{ inputs.date }} + sha: ${{ inputs.sha }} + - uses: aws-actions/configure-aws-credentials@cbe3b392738ccf3f987d68400dafcf4b0624a56c # v6.2.4 + with: + role-to-assume: ${{ vars.AWS_ROLE_ARN }} + aws-region: ${{ vars.AWS_REGION }} + role-duration-seconds: 10800 + - name: Setup sccache-dist + uses: rapidsai/shared-actions/setup-sccache-dist@main + with: + request-timeout: 7140 + - name: Build standalone C tarball + env: + CUVS_TARBALL_CUDA_VERSION: ${{ matrix.CUDA_VER }} + CUVS_TARBALL_PYTHON_VERSION: ${{ matrix.PY_VER }} + run: | + ./build.sh tarball --tarball-build-tests + - name: Upload artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz + path: libcuvs_c.tar.gz rust-build-matrix: needs: cpp-build permissions: @@ -128,7 +148,7 @@ jobs: contents: read uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 with: - build_type: pull-request + build_type: ${{ inputs.build_type || 'branch' }} matrix_name: conda-cpp-build matrix_filter: map(select(.ARCH == "amd64")) go-build: @@ -161,7 +181,7 @@ jobs: contents: read uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 with: - build_type: pull-request + build_type: ${{ inputs.build_type || 'branch' }} matrix_name: conda-cpp-build matrix_filter: map(select(.ARCH == "amd64")) java-build: @@ -190,7 +210,8 @@ jobs: file_to_upload: "java/cuvs-java/target/" sha: ${{ inputs.sha }} lucene-build: - # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against. + # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against, and + # reuses its matrix so the two always agree on the set of CUDA versions. needs: [java-build-matrix, java-build] permissions: actions: read diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index c4b576d6a4..e42c1d7d61 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -27,7 +27,10 @@ jobs: - rocky8-clib-tests - conda-java-build-and-tests-matrix - conda-java-build-and-tests + - conda-java-tests-other-arch-matrix + - conda-java-tests-other-arch - conda-lucene-build-and-tests + - conda-lucene-tests-other-arch - rust-build-matrix - rust-build - go-build-matrix @@ -563,30 +566,45 @@ jobs: build_type: pull-request matrix_name: conda-cpp-build rocky8-clib-standalone-build: - needs: [build-details, rocky8-clib-standalone-build-matrix] + needs: [rocky8-clib-standalone-build-matrix] + runs-on: linux-${{ matrix.ARCH }}-cpu16 + env: + RAPIDS_BUILD_TYPE: pull-request 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@release/26.10 strategy: fail-fast: false matrix: ${{ fromJSON(needs.rocky8-clib-standalone-build-matrix.outputs.matrix) }} - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - arch: "${{matrix.arch}}" - date: ${{ inputs.date }}_c - container_image: "rapidsai/ci-wheel:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" - node_type: "cpu16" - requires_license_builder: true - script: "ci/build_standalone_c.sh --build-tests" - artifact-name: "libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz" - file_to_upload: "libcuvs_c.tar.gz" - sha: ${{ inputs.sha }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + ref: ${{ inputs.sha || github.sha }} + fetch-depth: 0 + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@cbe3b392738ccf3f987d68400dafcf4b0624a56c # v6.2.4 + with: + role-to-assume: ${{ vars.AWS_ROLE_ARN }} + aws-region: ${{ vars.AWS_REGION }} + role-duration-seconds: 10800 + - name: Setup sccache-dist + uses: rapidsai/shared-actions/setup-sccache-dist@main + with: + request-timeout: 7140 + - name: Build standalone C tarball + env: + CUVS_TARBALL_CUDA_VERSION: ${{ matrix.CUDA_VER }} + CUVS_TARBALL_PYTHON_VERSION: ${{ matrix.PY_VER }} + run: | + ./build.sh tarball --tarball-build-tests + - name: Upload artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz + path: libcuvs_c.tar.gz rocky8-clib-tests-matrix: needs: [rocky8-clib-standalone-build, changed-files] permissions: @@ -650,6 +668,39 @@ jobs: script: "ci/test_java.sh" artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" + conda-java-tests-other-arch-matrix: + # All architectures the C++ core builds for, except amd64. + needs: [conda-cpp-build, changed-files] + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 + 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-java-tests-other-arch: + # Verifies that the amd64-built cuvs-java jar (jextract-generated Panama bindings and + # all) also works correctly on every non-amd64 architecture, by running its own IT test + # suite against the amd64-built classes without recompiling -- see ci/test_java_prebuilt.sh. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests] + 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@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-matrix.outputs.matrix) }} + with: + build_type: pull-request + node_type: "gpu-l4-latest-1" + arch: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_java_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }}" conda-lucene-build-and-tests: # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against, and # reuses its matrix so the two always agree on the set of CUDA versions. @@ -675,6 +726,29 @@ jobs: script: "ci/test_lucene.sh cuvs-java-cuda${{ matrix.CUDA_VER }}" artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-lucene/target/" + conda-lucene-tests-other-arch: + # Verifies that cuvs-lucene, built against the amd64-built cuvs-java jar, also works + # correctly on every non-amd64 architecture, by running its test suite against the + # amd64-built classes without recompiling -- see ci/test_lucene_prebuilt.sh. Reuses + # conda-java-tests-other-arch-matrix so the two always agree on which architectures to test. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests, conda-lucene-build-and-tests] + 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@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-matrix.outputs.matrix) }} + with: + build_type: pull-request + node_type: "gpu-l4-latest-1" + arch: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_lucene_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }} cuvs-lucene-cuda${{ matrix.CUDA_VER }}" rust-build-matrix: needs: [conda-cpp-build, changed-files] permissions: diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 16d967973b..98e51d8365 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -104,6 +104,41 @@ jobs: script: "ci/test_java.sh" artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" + conda-java-tests-other-arch-matrix: + # All architectures the C++ core builds for, except amd64 -- so a newly added + # architecture is automatically picked up here without a workflow change. + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 + with: + build_type: ${{ inputs.build_type }} + matrix_name: conda-cpp-build + matrix_filter: map(select(.ARCH != "amd64")) + conda-java-tests-other-arch: + # Verifies that the amd64-built cuvs-java jar (jextract-generated Panama bindings and + # all) also works correctly on every non-amd64 architecture, by running its own IT test + # suite against the amd64-built classes without recompiling -- see ci/test_java_prebuilt.sh. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests] + 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@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-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: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_java_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }}" conda-lucene-build-and-tests: # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against, and # reuses its matrix so the two always agree on the set of CUDA versions. @@ -132,6 +167,32 @@ jobs: script: "ci/test_lucene.sh cuvs-java-cuda${{ matrix.CUDA_VER }}" artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-lucene/target/" + conda-lucene-tests-other-arch: + # Verifies that cuvs-lucene, built against the amd64-built cuvs-java jar, also works + # correctly on every non-amd64 architecture, by running its test suite against the + # amd64-built classes without recompiling -- see ci/test_lucene_prebuilt.sh. Reuses + # conda-java-tests-other-arch-matrix so the two always agree on which architectures to test. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests, conda-lucene-build-and-tests] + 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@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-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: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_lucene_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }} cuvs-lucene-cuda${{ matrix.CUDA_VER }}" wheel-tests-cuvs: permissions: actions: read diff --git a/.gitignore b/.gitignore index 3627558ff5..7357b2a81e 100644 --- a/.gitignore +++ b/.gitignore @@ -19,6 +19,7 @@ dist/ python/**/**/*.cpp python/cuvs/record.txt log +infrequent_licenses/ .ipynb_checkpoints .DS_Store dask-worker-space/ @@ -29,6 +30,7 @@ temporary_*.json rust/target/ rust/Cargo.lock rmm_log.txt +*.tar.gz ## example notebooks notebooks/simplewiki-2020-11-01-nq-distilbert-base-v1.pt diff --git a/Dockerfile.standalone b/Dockerfile.standalone new file mode 100644 index 0000000000..11a4580d81 --- /dev/null +++ b/Dockerfile.standalone @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Environment and runner for the standalone C build. +# +# End-user instructions: https://docs.nvidia.com/cuvs/installation#build-the-standalone-c-library-with-docker + +ARG CUDA_VERSION=notset +ARG PYTHON_VERSION=notset +ARG RAPIDS_VERSION=notset +FROM rapidsai/ci-wheel:${RAPIDS_VERSION}-cuda${CUDA_VERSION}-rockylinux8-py${PYTHON_VERSION} + +# Output directory for the standalone archive. Bind-mount a host folder here +# (e.g. -v $(pwd)/build:/build) so libcuvs_c.tar.gz is written to the host. +ENV CUVS_TARBALL_BUILD_OUTPUT_DIR=/build +ENV CUVS_TARBALL_IN_CONTAINER=1 + +# Run from repo root; the repo is expected to be bind-mounted at /workspace. +WORKDIR /workspace + +# Run the build script, which writes the standalone tarball to the mounted output dir. +# Example: docker run -v $(pwd):/workspace -v $(pwd)/build:/build cuvs-standalone-c [--tarball-build-tests] +ENTRYPOINT ["/bin/bash", "-c", "ci/standalone_c/build_standalone_c.sh \"$@\"", "--"] diff --git a/README.md b/README.md index 77d2466c6a..fa25907ca6 100755 --- a/README.md +++ b/README.md @@ -84,6 +84,10 @@ cuVS comes with pre-built packages that can be installed through [conda](https:/ Please see the [Build and Install Guide](https://docs.rapids.ai/api/cuvs/nightly/build/) for more information on installing the available cuVS packages and building from source. +### Standalone C library (Docker) + +To build the standalone C library tarball (`libcuvs_c.tar.gz`) for use in your own C/C++ projects, see **[Building the standalone C library with Docker](https://docs.nvidia.com/cuvs/installation#build-the-standalone-c-library-with-docker)**. + ## Getting Started The following code snippets train an approximate nearest neighbors index for the CAGRA algorithm in the various different languages supported by cuVS. diff --git a/build.sh b/build.sh index da8adb793f..dadbba9ce7 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 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 --build-java-examples" +VALIDARGS="clean libcuvs python rust go java lucene docs tests bench-ann examples tarball --tarball-build-tests --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 --build-java-examples" 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) @@ -34,6 +34,7 @@ HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool= is: -v - verbose build mode @@ -59,6 +60,7 @@ HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool=\\\" - pass arbitrary list of CMake configuration options (escape all quotes in argument) @@ -607,3 +609,56 @@ if hasArg examples; then ./build.sh popd fi + +################################################################################ +# Build the standalone C library tarball (if requested) + +if hasArg tarball; then + if [[ "${CUVS_TARBALL_IN_CONTAINER:-0}" == "1" ]]; then + CUVS_TARBALL_BUILD_OUTPUT_DIR="${CUVS_TARBALL_BUILD_OUTPUT_DIR:-${REPODIR}}" + tar czf "${CUVS_TARBALL_BUILD_OUTPUT_DIR}/libcuvs_c.tar.gz" -C "${REPODIR}/c/build/install" . + else + CUVS_TARBALL_CUDA_VERSION="${CUVS_TARBALL_CUDA_VERSION:-13.3.0}" + CUVS_TARBALL_PYTHON_VERSION="${CUVS_TARBALL_PYTHON_VERSION:-3.14}" + CUVS_TARBALL_BUILD_OUTPUT_DIR="${CUVS_TARBALL_BUILD_OUTPUT_DIR:-${REPODIR}/build}" + CUVS_TARBALL_IMAGE_NAME="nvidia/cuvs-standalone-c:local-cuda${CUVS_TARBALL_CUDA_VERSION}-py${CUVS_TARBALL_PYTHON_VERSION}" + + mkdir -p "${CUVS_TARBALL_BUILD_OUTPUT_DIR}" + BUILD_OUTPUT_DIR_ABS=$(realpath "${CUVS_TARBALL_BUILD_OUTPUT_DIR}") + + echo "Building Docker image ${CUVS_TARBALL_IMAGE_NAME} (CUDA ${CUVS_TARBALL_CUDA_VERSION}, Python ${CUVS_TARBALL_PYTHON_VERSION})..." + docker build \ + -f "${REPODIR}/Dockerfile.standalone" \ + --build-arg CUDA_VERSION="${CUVS_TARBALL_CUDA_VERSION}" \ + --build-arg PYTHON_VERSION="${CUVS_TARBALL_PYTHON_VERSION}" \ + --build-arg RAPIDS_VERSION="${RAPIDS_VERSION_MAJOR_MINOR}" \ + -t "${CUVS_TARBALL_IMAGE_NAME}" \ + "${REPODIR}" + + # optionally pass additional arguments through to the container's entrypoint + DOCKER_ENTRYPOINT_ARGS=() + if hasArg --tarball-build-tests; then + DOCKER_ENTRYPOINT_ARGS+=(--tarball-build-tests) + fi + + echo "Running standalone C build in container..." + # NOTE: the '--env-file' trick for AWS credentials keeps them out of 'ps aux' / 'docker ps' output + docker run \ + --rm \ + -v "${REPODIR}:/workspace:rw" \ + -v "${BUILD_OUTPUT_DIR_ABS}:/build:rw" \ + --env CI="${CI:-false}" \ + --env PARALLEL_LEVEL="${PARALLEL_LEVEL}" \ + --env RAPIDS_BUILD_TYPE="${RAPIDS_BUILD_TYPE:-}" \ + --env-file <(env | grep -E '^AWS_(ACCESS_KEY_ID|SECRET_ACCESS_KEY|SESSION_TOKEN)=') \ + --env-file <(env | grep -E '^SCCACHE_.*=') \ + "${CUVS_TARBALL_IMAGE_NAME}" \ + "${DOCKER_ENTRYPOINT_ARGS[@]}" + + cp -v "${BUILD_OUTPUT_DIR_ABS}/libcuvs_c.tar.gz" "${REPODIR}/libcuvs_c.tar.gz" + echo "Copied libcuvs_c.tar.gz to ${REPODIR}/libcuvs_c.tar.gz" + fi + + # print contents of the tarball + ls -lh "${CUVS_TARBALL_BUILD_OUTPUT_DIR}/libcuvs_c.tar.gz" +fi diff --git a/ci/build_standalone_c.sh b/ci/standalone_c/build_standalone_c.sh similarity index 64% rename from ci/build_standalone_c.sh rename to ci/standalone_c/build_standalone_c.sh index e1267f71b7..9b7a4d4a47 100755 --- a/ci/build_standalone_c.sh +++ b/ci/standalone_c/build_standalone_c.sh @@ -1,67 +1,47 @@ #!/bin/bash # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Build script for the standalone C library. +# +# Use 'Dockerfile.standalone' to build an image with all the prerequisites +# and run this in a container. set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + TOOLSET_VERSION=14 -NINJA_VERSION=v1.13.1 BUILD_C_LIB_TESTS="OFF" -if [[ "${1:-}" == "--build-tests" ]]; then +if [[ "${1:-}" == "--tarball-build-tests" ]]; then BUILD_C_LIB_TESTS="ON" fi -dnf install -y \ - patch \ - tar \ - unzip \ - wget - -if ! command -V ninja >/dev/null 2>&1; then - case "$(uname -m)" in - x86_64) - wget --no-hsts -q -O /tmp/ninja-linux.zip "https://github.com/ninja-build/ninja/releases/download/${NINJA_VERSION}/ninja-linux.zip"; - ;; - aarch64) - wget --no-hsts -q -O /tmp/ninja-linux.zip "https://github.com/ninja-build/ninja/releases/download/${NINJA_VERSION}/ninja-linux-aarch64.zip"; - ;; - *) - echo "Unrecognized platform '$(uname -m)'" >&2 - exit 1 - ;; - esac - unzip -d /usr/bin /tmp/ninja-linux.zip - chmod +x /usr/bin/ninja - rm /tmp/ninja-linux.zip -fi - source rapids-install-sccache source rapids-configure-sccache -source rapids-datetime-string -rapids-pip-retry install cmake +PIP_PACKAGES=( + 'cmake>=4.0' + 'git+https://github.com/rapidsai/spdx-license-builder.git' + 'ninja>=1.13' +) RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" if [[ "${RAPIDS_CUDA_MAJOR}" == "13" ]]; then - rapids-pip-retry install cuda-tile "cuda-toolkit[tileiras]==${RAPIDS_CUDA_VERSION%.*}.*" + PIP_PACKAGES+=( + cuda-tile + ) fi +rapids-pip-retry install "${PIP_PACKAGES[@]}" pyenv rehash -rapids-print-env - rapids-logger "Begin cpp build" sccache --stop-server 2>/dev/null || true -RAPIDS_PACKAGE_VERSION=$(rapids-generate-version) -export RAPIDS_PACKAGE_VERSION - -RAPIDS_ARTIFACTS_DIR=${RAPIDS_ARTIFACTS_DIR:-"${PWD}/artifacts"} -mkdir -p "${RAPIDS_ARTIFACTS_DIR}" -export RAPIDS_ARTIFACTS_DIR - scl enable gcc-toolset-${TOOLSET_VERSION} -- \ cmake -S cpp -B cpp/build/ -GNinja \ -DCMAKE_CUDA_HOST_COMPILER=/opt/rh/gcc-toolset-${TOOLSET_VERSION}/root/usr/bin/gcc \ @@ -107,11 +87,8 @@ if [ "${BUILD_C_LIB_TESTS}" != "OFF" ]; then cmake --install c/build --prefix c/build/install --component testing fi - rapids-logger "Begin gathering licenses" -rapids-pip-retry install git+https://github.com/rapidsai/spdx-license-builder.git license-builder . --output-json c/build/install/licenses.json --output-txt c/build/install/LICENSE rapids-logger "Begin c tarball creation" -tar czf libcuvs_c.tar.gz -C c/build/install/ . -ls -lh libcuvs_c.tar.gz +"${REPO_ROOT}/build.sh" tarball diff --git a/ci/standalone_c/dlsym_smoke.c b/ci/standalone_c/dlsym_smoke.c new file mode 100644 index 0000000000..e9148a5012 --- /dev/null +++ b/ci/standalone_c/dlsym_smoke.c @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +int main(int argc, char** argv) +{ + if (argc != 2) { + fprintf(stderr, "Usage: %s \n", argv[0]); + return 2; + } + + void* handle = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL); + if (handle == NULL) { + fprintf(stderr, "Failed to load %s: %s\n", argv[1], dlerror()); + return 1; + } + + dlerror(); + void* symbol = dlsym(handle, "cuvsResourcesCreate"); + const char* error = dlerror(); + if (error != NULL || symbol == NULL) { + fprintf(stderr, + "Failed to resolve cuvsResourcesCreate: %s\n", + error != NULL ? error : "unknown error"); + dlclose(handle); + return 1; + } + + printf("Successfully loaded %s and resolved cuvsResourcesCreate\n", argv[1]); + dlclose(handle); + return 0; +} diff --git a/ci/test_java_prebuilt.sh b/ci/test_java_prebuilt.sh new file mode 100755 index 0000000000..5f811aa9b5 --- /dev/null +++ b/ci/test_java_prebuilt.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# Runs cuvs-java's own IT test suite against the classes from a prior amd64 build job, +# without recompiling (or running jextract) on this host. cuvs-java is always built on +# amd64; this verifies that the resulting jextract-generated Panama bindings also work +# correctly against a native libcuvs_c.so on the given (non-amd64) architecture. +# +# Takes the target architecture (as used by the GitHub Actions job matrix, e.g. "arm64") +# and the name of the cuvs-java artifact uploaded by the amd64 Java job. + +ARCH="${1:?Usage: $0 }" +CUVS_JAVA_ARTIFACT="${2:?Usage: $0 }" + +rapids-logger "Testing the amd64-built cuvs-java artifact on ${ARCH}" + +case "${ARCH}" in + amd64) CONDA_ARCH="x86_64" ;; + arm64) CONDA_ARCH="aarch64" ;; + *) CONDA_ARCH="${ARCH}" ;; +esac + +if [ -e "/opt/conda/etc/profile.d/conda.sh" ]; then + . /opt/conda/etc/profile.d/conda.sh +fi + +rapids-logger "Check GPU usage" +nvidia-smi + +rapids-logger "Configuring conda strict channel priority" +conda config --set channel_priority strict + +rapids-logger "Downloading artifacts from previous jobs" +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") +CUVS_JAVA_DIR=$(rapids-download-from-github "${CUVS_JAVA_ARTIFACT}") + +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=${CONDA_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 + +# libcuvs comes from the conda environment here, matching the architecture this script is +# running on. +export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + +rapids-logger "Restore the amd64-built target/ directory so Maven can run the tests without recompiling" + +rm -rf java/cuvs-java/target +mkdir -p java/cuvs-java/target +cp -a "${CUVS_JAVA_DIR}/." java/cuvs-java/target/ + +EXITCODE=0 +trap "EXITCODE=1" ERR +set +e + +rapids-logger "Run cuvs-java IT tests against the amd64-built classes" + +# -Dskip.compile activates the pom's "skip-compile" profile, which disables all +# compilation for this run, forcing test to use amd64-compiled jar instead of +# local code. +pushd java/cuvs-java +mvn --batch-mode verify -Dskip.compile=true +popd + +rapids-logger "Test script exiting with value: $EXITCODE" +exit ${EXITCODE} diff --git a/ci/test_lucene_prebuilt.sh b/ci/test_lucene_prebuilt.sh new file mode 100755 index 0000000000..477ff75ad3 --- /dev/null +++ b/ci/test_lucene_prebuilt.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# Runs cuvs-lucene's test suite against the classes from a prior amd64 build job, without +# recompiling on this host. cuvs-lucene is always built on amd64, against the amd64-built +# cuvs-java jar; this verifies that pairing also works correctly on the given (non-amd64) +# architecture. Since cuvs-lucene depends on the plain (no bundled natives) cuvs-java jar, +# this doubles as a cross-arch check that the jextract-generated Panama bindings baked +# into the amd64 jar work unmodified against a native libcuvs_c.so on that architecture. +# +# Takes the target architecture (as used by the GitHub Actions job matrix, e.g. "arm64") +# and the names of the cuvs-java and cuvs-lucene artifacts uploaded by the amd64 jobs. + +ARCH="${1:?Usage: $0 }" +CUVS_JAVA_ARTIFACT="${2:?Usage: $0 }" +CUVS_LUCENE_ARTIFACT="${3:?Usage: $0 }" + +rapids-logger "Testing the amd64-built cuvs-java/cuvs-lucene artifacts on ${ARCH}" + +case "${ARCH}" in + amd64) CONDA_ARCH="x86_64" ;; + arm64) CONDA_ARCH="aarch64" ;; + *) CONDA_ARCH="${ARCH}" ;; +esac + +if [ -e "/opt/conda/etc/profile.d/conda.sh" ]; then + . /opt/conda/etc/profile.d/conda.sh +fi + +rapids-logger "Check GPU usage" +nvidia-smi + +rapids-logger "Configuring conda strict channel priority" +conda config --set channel_priority strict + +rapids-logger "Downloading artifacts from previous jobs" +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") +CUVS_JAVA_DIR=$(rapids-download-from-github "${CUVS_JAVA_ARTIFACT}") +CUVS_LUCENE_DIR=$(rapids-download-from-github "${CUVS_LUCENE_ARTIFACT}") + +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=${CONDA_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 + +# libcuvs comes from the conda environment here, matching the architecture this script is +# running on. cuvs-lucene depends on the plain cuvs-java jar, which bundles no native +# libraries, so the JVM resolves libcuvs_c.so through the dynamic loader. +export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + +rapids-logger "Install the amd64-built cuvs-java artifact into the local Maven repository" + +# cuvs-lucene resolves cuvs-java from the local Maven repository. Its pom.xml travels with +# the artifact and supplies the coordinates, so no version needs to be hardcoded. +CUVS_JAVA_POM="${CUVS_JAVA_DIR}/pom.xml" +if [ ! -f "${CUVS_JAVA_POM}" ]; then + echo "Could not find pom.xml in the cuvs-java artifact at ${CUVS_JAVA_DIR}" >&2 + exit 1 +fi + +# The artifact also carries the per-architecture native jar and the sources/javadoc/test jars; +# cuvs-lucene depends on the plain one. +mapfile -t CUVS_JAVA_JARS < <(find "${CUVS_JAVA_DIR}" -maxdepth 1 -name 'cuvs-java-*.jar' \ + ! -name '*-sources.jar' ! -name '*-javadoc.jar' ! -name '*-tests.jar' ! -name '*-cuda*.jar') +if [ "${#CUVS_JAVA_JARS[@]}" -ne 1 ]; then + echo "Expected exactly one cuvs-java jar in ${CUVS_JAVA_DIR}, found: ${CUVS_JAVA_JARS[*]:-none}" >&2 + exit 1 +fi + +# Install cuvs-java jar into .m2, cd is needed to pick up pom.xml in order to avoid rate limit of main maven repo +pushd java/cuvs-lucene +mvn --batch-mode install:install-file -Dfile="${CUVS_JAVA_JARS[0]}" -DpomFile="${CUVS_JAVA_POM}" +popd + +rapids-logger "Restore the amd64-built cuvs-lucene target/ directory so Maven can run the tests without recompiling" + +rm -rf java/cuvs-lucene/target +mkdir -p java/cuvs-lucene/target +cp -a "${CUVS_LUCENE_DIR}/." java/cuvs-lucene/target/ + +EXITCODE=0 +trap "EXITCODE=1" ERR +set +e + +rapids-logger "Run cuvs-lucene tests against the amd64-built classes" + +# -Dskip.compile activates the pom's "skip-compile" profile, which disables all +# compilation for this run, forcing test to use amd64-compiled jar instead of +# local code. +pushd java/cuvs-lucene +mvn --batch-mode test -Dskip.compile=true +popd + +rapids-logger "Test script exiting with value: $EXITCODE" +exit ${EXITCODE} diff --git a/ci/test_standalone_c.sh b/ci/test_standalone_c.sh index 894d687bff..519fa2470d 100755 --- a/ci/test_standalone_c.sh +++ b/ci/test_standalone_c.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -7,6 +7,7 @@ set -euo pipefail rapids-pip-retry install cmake pyenv rehash +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" INSTALL_PREFIX="${PWD}/libcuvs_c_install" mkdir -p "${INSTALL_PREFIX}" @@ -24,6 +25,25 @@ DOWNLOAD_LOCATION=$(rapids-download-from-github "${payload_name}") # Extract the artifact to a staging directory tar -xf "${DOWNLOAD_LOCATION}/${pkg_name}" -C "${INSTALL_PREFIX}" +rapids-logger "Validate C API shared library" +C_API_LIBRARY="" +for C_API_LIBRARY_DIR in "${INSTALL_PREFIX}/lib" "${INSTALL_PREFIX}/lib64"; do + if [[ -f "${C_API_LIBRARY_DIR}/libcuvs_c.so" ]]; then + C_API_LIBRARY="${C_API_LIBRARY_DIR}/libcuvs_c.so" + break + fi +done + +if [[ -z "${C_API_LIBRARY}" ]]; then + echo "Error: C API shared library not found under ${INSTALL_PREFIX}/lib or ${INSTALL_PREFIX}/lib64" >&2 + exit 1 +fi + +C_API_SMOKE_TEST="${INSTALL_PREFIX}/bin/cuvs_c_dlsym_smoke" +"${CC:-cc}" -std=c11 -Wall -Wextra -Werror \ + "${SCRIPT_DIR}/standalone_c/dlsym_smoke.c" -ldl -o "${C_API_SMOKE_TEST}" +LD_LIBRARY_PATH="$(dirname "${C_API_LIBRARY}"):${LD_LIBRARY_PATH:-}" \ + "${C_API_SMOKE_TEST}" "${C_API_LIBRARY}" rapids-logger "Run C API tests" ls -l "${INSTALL_PREFIX}" diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 90d23d9aef..9e3b69fdbe 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -181,6 +181,7 @@ function(ConfigureAnnBench) sqlite3 Threads::Threads $<$:CUDA::cudart_static> + $<$:rmm::rmm> $ $ ) @@ -426,8 +427,8 @@ if(CUVS_ANN_BENCH_SINGLE_EXE) target_include_directories(ANN_BENCH PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_link_libraries( - ANN_BENCH PRIVATE raft::raft nlohmann_json::nlohmann_json sqlite3 benchmark::benchmark dl - $<$:CUDA::nvtx3> + ANN_BENCH PRIVATE $ nlohmann_json::nlohmann_json sqlite3 + benchmark::benchmark dl $<$:CUDA::nvtx3> ) set_target_properties( ANN_BENCH diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp index 6e4c57b35e..069d257717 100644 --- a/cpp/bench/ann/src/common/util.hpp +++ b/cpp/bench/ann/src/common/util.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -7,6 +7,10 @@ #include "ann_types.hpp" #include "cuda_stub.hpp" // cuda-related utils +#ifndef BUILD_CPU_ONLY +#include +#endif + #if __has_include() #define ANN_BENCH_NVTX3_HEADERS_FOUND #include @@ -145,25 +149,8 @@ struct cuda_timer { }; #ifndef BUILD_CPU_ONLY -// ATM, rmm::stream does not support passing in flags; hence this helper type. -struct non_blocking_stream { - non_blocking_stream() { cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking); } - ~non_blocking_stream() noexcept - { - if (stream_ != nullptr) { cudaStreamDestroy(stream_); } - } - non_blocking_stream(non_blocking_stream const&) = delete; - non_blocking_stream(non_blocking_stream&& other) noexcept { std::swap(stream_, other.stream_); } - auto operator=(non_blocking_stream const&) -> non_blocking_stream& = delete; - auto operator=(non_blocking_stream&&) -> non_blocking_stream& = delete; - [[nodiscard]] auto view() const noexcept -> cudaStream_t { return stream_; } - - private: - cudaStream_t stream_{nullptr}; -}; - namespace detail { -inline std::vector global_stream_pool(0); +inline std::vector global_stream_pool(0); inline std::mutex gsp_mutex; } // namespace detail #endif @@ -180,7 +167,9 @@ inline auto get_stream_from_global_pool() -> cudaStream_t #ifndef BUILD_CPU_ONLY std::lock_guard guard(detail::gsp_mutex); if (static_cast(detail::global_stream_pool.size()) < benchmark_n_threads) { - detail::global_stream_pool.resize(benchmark_n_threads); + while (static_cast(detail::global_stream_pool.size()) < benchmark_n_threads) { + detail::global_stream_pool.emplace_back(rmm::cuda_stream::flags::non_blocking); + } } return detail::global_stream_pool[benchmark_thread_id].view(); #else diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h index 1a276e8cc8..f699bdad1a 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,9 +15,11 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -127,6 +129,7 @@ class configured_raft_resources { { raft::resource::set_large_workspace_resource( *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()}); + raft::resource::set_cuda_stream_pool(*res_, std::make_shared(1)); } /** Default constructor creates all resources anew. */ diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h index ae334f70d7..8ed5adee26 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,7 +15,6 @@ #include #include #include -#include #include #include @@ -88,9 +87,6 @@ class cuvs_ivf_flat : public algo, public algo_gpu { template void cuvs_ivf_flat::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); index_ = std::make_shared>( std::move(cuvs::neighbors::ivf_flat::build( handle_, diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h index 161563d8fe..363d065e91 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -20,7 +20,6 @@ #include #include #include -#include #include @@ -118,9 +117,6 @@ void cuvs_ivf_pq::load(const std::string& file) template void cuvs_ivf_pq::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); std::make_shared>( std::move(cuvs::neighbors::ivf_pq::build(handle_, index_params_, dataset_v))) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index ca8f77b808..0c41ee1c96 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -18,7 +18,6 @@ #include #include #include -#include #include @@ -98,9 +97,6 @@ void cuvs_ivf_rabitq::load(const std::string& file) template void cuvs_ivf_rabitq::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); std::make_shared>( std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v))) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h index 1503e6bb84..5bf0098eaa 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -12,9 +12,7 @@ #include #include #include -#include #include -#include #include #include @@ -83,8 +81,6 @@ class cuvs_ivf_sq : public algo, public algo_gpu { template void cuvs_ivf_sq::build(const T* dataset, size_t nrow) { - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); index_ = std::make_shared>( std::move(cuvs::neighbors::ivf_sq::build( handle_, diff --git a/dependencies.yaml b/dependencies.yaml index 08b1209411..374949ff64 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -429,6 +429,8 @@ dependencies: - libcusolver-dev - libcusparse-dev - libnvjitlink-dev + # NOTE: 'tileiras' is needed too, but in CI that's provided by a system-installed CTK, so + # we do not use the [tileiras] extra recommended at https://pypi.org/project/cuda-tile/ cutile_python: specific: - output_types: conda @@ -436,16 +438,6 @@ dependencies: - matrix: cuda: "12.*" packages: - - matrix: - cuda: "13.3" - packages: - - cutile-python - - cuda-tileiras - - matrix: - cuda: "13.*" - packages: - - cutile-python - - cuda-tileiras - matrix: packages: - cutile-python @@ -455,22 +447,15 @@ dependencies: - matrix: cuda: "12.*" packages: - - matrix: - cuda: "13.3" - packages: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* - matrix: cuda: "13.*" packages: - &cutile_python_cu13 cuda-tile - - &cutile_toolkit_cu13 cuda-toolkit[tileiras]==13.* # if no matching matrix selectors passed, list the CUDA 13 requirement # (as a source of documentation in the generated pyproject.toml) - matrix: packages: - *cutile_python_cu13 - - *cutile_toolkit_cu13 cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing diff --git a/examples/java/cuvs-lucene/README.md b/examples/java/cuvs-lucene/README.md index 4113b00adc..5a3c3675e0 100644 --- a/examples/java/cuvs-lucene/README.md +++ b/examples/java/cuvs-lucene/README.md @@ -32,17 +32,3 @@ To run the Index and Search on GPU example do: ```sh mvn clean install && 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 ``` - -To run the optimized CAGRA-HNSW build example (reference pattern for efficiently building an -accelerated HNSW index from a large `.fbin` with every ingest-side knob on — open the file once and -stream sequential prefetched chunks that overlap the disk read with indexing, hold at most two chunks -in memory, reuse a single vector array, size a native flat buffer per segment, auto-select the CAGRA -graph-build algorithm, and optionally partition into K segments built sequentially or overlapped) do: - -```sh -mvn clean install && 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.OptimizedCagraHnswBuildExample -``` - -With no arguments it generates and indexes a small demo `.fbin` as a single segment; pass a real file, -chunk size, segment count, and overlap flag as -`... OptimizedCagraHnswBuildExample `. diff --git a/examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java deleted file mode 100644 index 6cfcfafa61..0000000000 --- a/examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java +++ /dev/null @@ -1,268 +0,0 @@ -/* - * 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 com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; -import com.nvidia.cuvs.lucene.CagraHnswBulkIndexWriter; -import com.nvidia.cuvs.lucene.FbinVectorSource; -import com.nvidia.cuvs.spi.CuVSProvider; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.nio.file.StandardOpenOption; -import java.util.Random; -import java.util.UUID; -import java.util.logging.Logger; -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.IndexWriterConfig; -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.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; - -/** - * Reference usage of {@link CagraHnswBulkIndexWriter}: builds an accelerated HNSW index (whose - * graph is built on the GPU with CAGRA) from a {@code .fbin} vector file on local disk. - * - *

Demonstrates both ways to use {@link CagraHnswBulkIndexWriter}: - * - *

    - *
  • {@link #main} — the one-shot convenience path. All of the bulk-build mechanics — - * prefetched streaming reads, native flat buffering, the {@code IndexWriterConfig} tuning - * that guarantees a single unmerged segment per slice, K-segment partitioning, and combining - * the result by hardlink — are owned by {@link CagraHnswBulkIndexWriter} itself; see its - * Javadoc for how each of those works and the tradeoffs of {@code numSegments} and {@code - * overlap}. This example only wires up what's genuinely application-specific: where the - * vectors come from ({@link FbinVectorSource}, or your own {@link - * com.nvidia.cuvs.lucene.VectorSource} for a different data source), the graph-build quality - * knobs ({@link AcceleratedHNSWParams}), and — via a {@link - * CagraHnswBulkIndexWriter.FieldCallback} — any per-vector metadata to attach. - *
  • {@link #runManualExample} — the manual, direct-instance path: construct a {@link - * CagraHnswBulkIndexWriter} yourself and drive {@code addDocument}/{@code close} exactly - * like a plain Lucene {@code IndexWriter}, building each {@link Document} (metadata included) - * yourself instead of going through a callback. - *
- * - *

Usage: {@code OptimizedCagraHnswBuildExample [] [] [] - * []}. With no arguments a small demo {@code .fbin} is generated and indexed as a - * single segment. - */ -public class OptimizedCagraHnswBuildExample { - - private static final Logger log = - Logger.getLogger(OptimizedCagraHnswBuildExample.class.getName()); - private static final String ID_FIELD = "id"; - private static final String CATEGORY_FIELD = "category"; - private static final String VECTOR_FIELD = "vector_field"; - - public static void main(String[] args) throws Exception { - // It is recommended to enable RMM allocation mode at application start, before constructing - // any CagraHnswBulkIndexWriter, to avoid device-wide sync from the default allocator. - CuVSProvider.provider().enableRMMAsyncMemory(); - - int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; - int numSegments = args.length >= 3 ? Math.max(1, Integer.parseInt(args[2])) : 1; - boolean overlap = args.length >= 4 && Boolean.parseBoolean(args[3]); - Path indexDirPath = Paths.get(UUID.randomUUID().toString()); - - Path fbinPath; - boolean generated = false; - if (args.length >= 1) { - fbinPath = Paths.get(args[0]); - } else { - fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); - writeDemoFbin(fbinPath, 5000, 32, new Random(222)); - generated = true; - log.info("No .fbin provided; generated a demo file at " + fbinPath); - } - - try { - int dim; - try (FbinVectorSource probe = new FbinVectorSource(fbinPath, 1)) { - dim = probe.dimensions(); - } - - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dim, VectorSimilarityFunction.EUCLIDEAN) - .idField(ID_FIELD) - .graphBuild( - new AcceleratedHNSWParams.Builder() - // HEURISTIC lets cuVS pick the build algorithm and auto-tune its parameters - // based on maxConn and beamWidth below. - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) - // Primary recall/graph-size knobs. Higher values improve recall at the cost - // of a larger graph and longer build. Match to your dataset and recall - // target. - .withMaxConn(32) - .withBeamWidth(32) - // Must match the distance metric used when querying the index. - .withCuvsDistanceType(CuvsDistanceType.L2Expanded) - // Starting point: one thread per logical CPU. Profile and tune for your - // hardware. - .withWriterThreads(Runtime.getRuntime().availableProcessors()) - .build()) - .segments(numSegments, overlap) - .targetDirectory(indexDirPath) - .build(); - - log.info( - "Indexing " - + fbinPath - + " (" - + dim - + "-dim) into " - + numSegments - + " segment(s), " - + (overlap && numSegments > 1 ? "overlapped" : "sequential") - + " build, " - + chunkSizeMB - + " MB prefetched chunks"); - - // FieldCallback lets the one-shot path attach metadata per vector: indexFbin/build build the - // id+vector fields internally (they own the loop), so this is how a caller reaches the - // Document to add anything else -- here, an illustrative "even"/"odd" category by id. - CagraHnswBulkIndexWriter.indexFbin( - fbinPath, - config, - (doc, id) -> - doc.add( - new StringField(CATEGORY_FIELD, id % 2 == 0 ? "even" : "odd", Field.Store.YES)), - chunkSizeMB); - log.info("Index build complete: " + indexDirPath); - - runSampleSearch(indexDirPath, fbinPath, 5); - } finally { - FileUtils.deleteDirectory(indexDirPath.toFile()); - if (generated) { - Files.deleteIfExists(fbinPath); - } - } - - runManualExample(); - } - - /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ - private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { - float[] queryVector; - try (FbinVectorSource reader = new FbinVectorSource(fbinPath, 1)) { - queryVector = reader.get(0); - } - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); - log.info("Sample search returned " + results.scoreDocs.length + " hits:"); - for (int i = 0; i < results.scoreDocs.length; i++) { - ScoreDoc sd = results.scoreDocs[i]; - Document hit = searcher.storedFields().document(sd.doc); - log.info( - " rank " - + (i + 1) - + ": id=" - + hit.get(ID_FIELD) - + " category=" - + hit.get(CATEGORY_FIELD) - + " score=" - + sd.score); - } - } - } - - /** - * Short demonstration of the manual, direct-instance API: {@link CagraHnswBulkIndexWriter} is - * constructed directly and driven with {@code addDocument}/{@code close}, the same shape as a - * plain Lucene {@code IndexWriter} — the caller builds each {@link Document} itself, including - * whatever metadata it wants, with no callback needed since it already owns the loop. Unlike the - * one-shot path above, this only ever builds a single segment; K-segment partitioning and - * overlap are only available via {@link CagraHnswBulkIndexWriter#indexFbin}/{@link - * CagraHnswBulkIndexWriter#build}. - */ - private static void runManualExample() throws Exception { - int numDocs = 200; - int dim = 16; - Random random = new Random(7); - Path manualIndexDirPath = Paths.get("manual-" + UUID.randomUUID()); - - try { - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dim, VectorSimilarityFunction.EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - float[][] vectors = new float[numDocs][dim]; - try (Directory dir = FSDirectory.open(manualIndexDirPath); - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, numDocs)) { - for (int i = 0; i < numDocs; i++) { - for (int j = 0; j < dim; j++) { - vectors[i][j] = random.nextFloat() * 100; - } - Document doc = new Document(); - doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - doc.add(new StringField(CATEGORY_FIELD, i % 2 == 0 ? "even" : "odd", Field.Store.YES)); - doc.add( - new KnnFloatVectorField( - VECTOR_FIELD, vectors[i], VectorSimilarityFunction.EUCLIDEAN)); - writer.addDocument(doc); // same call shape as a plain IndexWriter - } - } // close() runs the single native-buffered flush (the GPU CAGRA build happens here) - - try (Directory dir = FSDirectory.open(manualIndexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[0], 1), 1); - Document hit = searcher.storedFields().document(results.scoreDocs[0].doc); - log.info( - "Manual example: nearest neighbor of vector 0 is id=" - + hit.get(ID_FIELD) - + " category=" - + hit.get(CATEGORY_FIELD)); - } - } finally { - FileUtils.deleteDirectory(manualIndexDirPath.toFile()); - } - } - - /** Writes a small random {@code .fbin} so the example is runnable without external data. */ - private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) - throws IOException { - ByteBuffer buf = - ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); - buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] - buf.putInt(dim); - for (int i = 0; i < numVectors; i++) { - for (int j = 0; j < dim; j++) { - buf.putFloat(random.nextFloat() * 100); - } - } - buf.flip(); - try (FileChannel ch = - FileChannel.open( - path, - StandardOpenOption.CREATE, - StandardOpenOption.WRITE, - StandardOpenOption.TRUNCATE_EXISTING)) { - while (buf.hasRemaining()) { - ch.write(buf); - } - } - } -} diff --git a/fern/pages/build.md b/fern/pages/build.md index e8fabd8d27..4cd49bba2b 100644 --- a/fern/pages/build.md +++ b/fern/pages/build.md @@ -41,6 +41,118 @@ conda activate cuvs You may prefer `mamba` over `conda` for faster environment solves. The `conda/environments` directory also contains language-specific environment YAML files for narrower development environments. Conda is not required, but if you do not use it, install all required build dependencies explicitly before running `build.sh`. +## Build the Standalone C Library with Docker + + +The standalone tarball is built with Docker so that it uses the supported toolchain versions and the build remains portable and reproducible across all supported installation platforms. + + +Use the standalone Docker build when you want a `libcuvs_c.tar.gz` archive that you can unpack and use to build your own C or C++ binaries for deployment or integration. + +### Prerequisites + +- Docker with support for the target platform: x86_64 or aarch64. +- At least 16 GB of memory and 20 GB of free disk space available to Docker. +- NVIDIA Container Toolkit and a GPU if you want to run GPU-dependent steps. The image is based on CUDA and may require GPU support at runtime. + +### Use the Helper Script + +From the repository root, run: + +```bash +# (optional) clean old build directories +rm -rf ./{build,c/build/,cpp/build} + +# build +./build.sh tarball +``` + +The script builds the Docker image, runs the build in a container, writes the tarball to `./build/libcuvs_c.tar.gz`, and copies it to `./libcuvs_c.tar.gz` for CI artifact upload and convenience. + +The helper accepts the following environment variables: + +| Variable | Default | Purpose | +| --- | --- | --- | +| `CUVS_TARBALL_CUDA_VERSION` | `13.3.0` | CUDA version for the `rapidsai/ci-wheel` base image. | +| `CUVS_TARBALL_PYTHON_VERSION` | `3.14` | Python version for the `rapidsai/ci-wheel` base image. | +| `CUVS_TARBALL_BUILD_OUTPUT_DIR` | `./build` | Host directory where the tarball is written. | + +CUDA and Python versions should match an existing `rapidsai/ci-wheel` tag. +See https://hub.docker.com/r/rapidsai/ci-wheel/tags + +For example: + +```bash +CUVS_TARBALL_CUDA_VERSION=13.3.0 \ +CUVS_TARBALL_PYTHON_VERSION=3.14 \ + ./build.sh tarball +``` + +To write the tarball to another directory, set `CUVS_TARBALL_BUILD_OUTPUT_DIR`: + +```console +$ CUVS_TARBALL_BUILD_OUTPUT_DIR="${PWD}/dist" ./build.sh tarball +$ find . -name 'libcuvs_c.tar.gz' +./dist/libcuvs_c.tar.gz +./libcuvs_c.tar.gz +``` + +To build and install the C library tests in the archive, pass `--tarball-build-tests`: + +```bash +./build.sh tarball --tarball-build-tests +``` + +### Tarball Contents + +The archive contains the headers, libraries, CMake configuration, and license information needed to compile and link C or C++ applications against the standalone NVIDIA cuVS libraries. + +### Build and Run the Docker Image Manually + +If you do not want to use the helper script, build the image directly from the repository root: + +```bash +docker build \ + -f Dockerfile.standalone \ + --build-arg CUDA_VERSION="13.3.0" \ + --build-arg PYTHON_VERSION="3.14" \ + --build-arg RAPIDS_VERSION="$(head -1 ./VERSION | cut -d. -f1,2 )" \ + -t cuvs-standalone-c:local \ + . +``` + +This command builds a local image and tags it as `cuvs-standalone-c:local`. + +Run the build in a container using that image and mount the repository plus an output directory: + +```bash +mkdir -p build +docker run --rm \ + -v "${PWD}:/workspace" \ + -v "${PWD}/build:/build" \ + cuvs-standalone-c:local +``` + +Mount another host directory at `/build` to change the output location: + +```bash +mkdir -p "${PWD}/dist" +docker run --rm \ + -v "${PWD}:/workspace" \ + -v "${PWD}/dist:/build" \ + cuvs-standalone-c:local +``` + +Pass `--tarball-build-tests` to include the C library tests: + +```bash +mkdir -p build +docker run --rm \ + -v "${PWD}:/workspace" \ + -v "${PWD}/build:/build" \ + cuvs-standalone-c:local --tarball-build-tests +``` + ## Documentation Preview The NVIDIA cuVS documentation is a Fern project in the repository's `fern` directory. Fern requires Node.js 22 or newer. If the docs fail with an error such as `SyntaxError: Unexpected token '.'`, check `node --version` and activate a newer Node.js runtime. diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md index f19b50fb6b..1e5018c3e7 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md @@ -412,13 +412,6 @@ Switch RMM allocations (used internally by various cuVS algorithms and by the de `CuVSDeviceMatrix`) to use pooled memory. This operation has a global effect, and will affect all resources on the current device. -Ownership: this and the other `enableRMM*`/`resetRMMPooledMemory` methods are -not called automatically by any cuVS/cuvs-lucene class. Because the effect is global to the -current device, it is application code's responsibility to call one of them, if desired, once -during startup/initialization — before any GPU resource is created — rather than have an -individual class call it lazily on first use, which cannot be made race-free against -concurrent construction of that class from multiple threads. - **Parameters** | Name | Description | @@ -426,7 +419,7 @@ concurrent construction of that class from multiple threads. | `initialPoolSizePercent` | The initial pool size, in percentage of the total GPU memory | | `maxPoolSizePercent` | The maximum pool size, in percentage of the total GPU memory | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:250`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:243`_ ### enableRMMManagedPooledMemory @@ -438,8 +431,6 @@ Switch RMM allocations (used internally by various cuVS algorithms and by the de `CuVSDeviceMatrix`) to use pooled memory. This operation has a global effect, and will affect all resources on the current device. -Ownership: see `#enableRMMPooledMemory`. - **Parameters** | Name | Description | @@ -447,7 +438,7 @@ Ownership: see `#enableRMMPooledMemory`. | `initialPoolSizePercent` | The initial pool size, in percentage of the total GPU memory | | `maxPoolSizePercent` | The maximum pool size, in percentage of the total GPU memory | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:262`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:253`_ ### enableRMMAsyncMemory @@ -462,9 +453,7 @@ on deallocation. This is especially beneficial when multiple CAGRA searches run on separate CUDA streams, because internal workspace allocations no longer serialize kernel launches. This operation has a global effect and will affect all resources on the current device. -Ownership: see `#enableRMMPooledMemory`. - -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:274`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:263`_ ### resetRMMPooledMemory @@ -474,9 +463,7 @@ void resetRMMPooledMemory() Disables pooled memory on the current device, reverting back to the default setting. -Ownership: see `#enableRMMPooledMemory`. - -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:281`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:266`_ ### provider @@ -486,7 +473,7 @@ static CuVSProvider provider() Retrieves the system-wide provider. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:284`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:269`_ ### cagraIndexParamsFromHnswParams @@ -516,7 +503,7 @@ may be shifted along the curve right or left. See the heuristics descriptions fo A new CAGRA index parameters object -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:304`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:289`_ ### cagraIndexParamsFromDataset @@ -540,6 +527,6 @@ Create CAGRA index parameters heuristically tuned for a dataset. A new CAGRA index parameters object -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:322`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:307`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:18`_ diff --git a/java/cuvs-java/pom.xml b/java/cuvs-java/pom.xml index b8e969d07b..d57cb3ccac 100644 --- a/java/cuvs-java/pom.xml +++ b/java/cuvs-java/pom.xml @@ -285,9 +285,46 @@ SPDX-License-Identifier: Apache-2.0 + + + skip-compile + + + skip.compile + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + default-compile + none + + + compile-java-22 + none + + + default-testCompile + none + + + + + + x86_64-cuda12 + + amd64 + cuda.version 12 @@ -371,6 +408,9 @@ SPDX-License-Identifier: Apache-2.0 x86_64-cuda13 + + amd64 + cuda.version 13 @@ -450,5 +490,177 @@ SPDX-License-Identifier: Apache-2.0 + + + aarch64-cuda12 + + + aarch64 + + + cuda.version + 12 + + + + aarch64-cuda12 + ${native.build.path}/cuda12 + + + + + org.apache.maven.plugins + maven-resources-plugin + 3.3.1 + + + copy-native-libs + prepare-package + + copy-resources + + + true + ${project.build.directory}/native-libs/${os.arch}/${os.name} + + + ${native.lib.path} + + libcuvs.so + libcuvs_c.so + + + + ${native.lib.path}/_deps/rmm-build + + librmm.so + + + + ${native.lib.path}/_deps/rapids_logger-build + + librapids_logger.so + + + + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.4.2 + + + src/assembly/native-with-deps.xml + + + + true + true + 12 + + + + + + assemble-native + package + + single + + + + + + + + + + aarch64-cuda13 + + + aarch64 + + + cuda.version + 13 + + + + aarch64-cuda13 + ${native.build.path}/cuda13 + + + + + org.apache.maven.plugins + maven-resources-plugin + 3.3.1 + + + copy-native-libs + prepare-package + + copy-resources + + + true + ${project.build.directory}/native-libs/${os.arch}/${os.name} + + + ${native.lib.path} + + libcuvs.so + libcuvs_c.so + + + + ${native.lib.path}/_deps/rmm-build + + librmm.so + + + + ${native.lib.path}/_deps/rapids_logger-build + + librapids_logger.so + + + + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.4.2 + + + src/assembly/native-with-deps.xml + + + + true + true + 13 + + + + + + assemble-native + package + + single + + + + + + + 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 cb44f9b672..44706d0b6a 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 @@ -238,13 +238,6 @@ MultiPartitionSearchResults searchCagraMultiPartition( * {@link CuVSDeviceMatrix}) to use pooled memory. * This operation has a global effect, and will affect all resources on the current device. * - *

Ownership: this and the other {@code enableRMM*}/{@code resetRMMPooledMemory} methods are - * not called automatically by any cuVS/cuvs-lucene class. Because the effect is global to the - * current device, it is application code's responsibility to call one of them, if desired, once - * during startup/initialization — before any GPU resource is created — rather than have an - * individual class call it lazily on first use, which cannot be made race-free against - * concurrent construction of that class from multiple threads. - * * @param initialPoolSizePercent The initial pool size, in percentage of the total GPU memory * @param maxPoolSizePercent The maximum pool size, in percentage of the total GPU memory */ @@ -255,8 +248,6 @@ MultiPartitionSearchResults searchCagraMultiPartition( * {@link CuVSDeviceMatrix}) to use pooled memory. * This operation has a global effect, and will affect all resources on the current device. * - *

Ownership: see {@link #enableRMMPooledMemory}. - * * @param initialPoolSizePercent The initial pool size, in percentage of the total GPU memory * @param maxPoolSizePercent The maximum pool size, in percentage of the total GPU memory */ @@ -269,16 +260,10 @@ MultiPartitionSearchResults searchCagraMultiPartition( * on deallocation. This is especially beneficial when multiple CAGRA searches run concurrently * on separate CUDA streams, because internal workspace allocations no longer serialize kernel * launches. This operation has a global effect and will affect all resources on the current device. - * - *

Ownership: see {@link #enableRMMPooledMemory}. */ void enableRMMAsyncMemory(); - /** - * Disables pooled memory on the current device, reverting back to the default setting. - * - *

Ownership: see {@link #enableRMMPooledMemory}. - */ + /** Disables pooled memory on the current device, reverting back to the default setting. */ void resetRMMPooledMemory(); /** Retrieves the system-wide provider. */ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java index 874f35ee20..b4f16a084c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java @@ -39,7 +39,7 @@ static CuVSProvider builtinProvider() { var osArch = System.getProperty("os.arch"); var supportedJavaRuntime = javaRuntimeVersion > 21; var supportedOs = osName.startsWith("Linux"); - var supportedArchitecture = osArch.equals("amd64"); + var supportedArchitecture = osArch.equals("amd64") || osArch.equals("aarch64"); if (supportedJavaRuntime && supportedOs && supportedArchitecture) { try { var cls = Class.forName("com.nvidia.cuvs.spi.JDKProvider"); @@ -63,7 +63,8 @@ static CuVSProvider builtinProvider() { unsupportedReasons.add("cuvs-java supports only Linux, but os.name is " + osName); } if (!supportedArchitecture) { - unsupportedReasons.add("cuvs-java supports only x86-64 (amd64), but os.arch is " + osArch); + unsupportedReasons.add( + "cuvs-java supports only x86-64 (amd64) and arm64 (aarch64), but os.arch is " + osArch); } return new UnsupportedProvider(String.join("; ", unsupportedReasons)); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java index ab59395dd2..1f28c45950 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java @@ -23,7 +23,7 @@ public class BruteForceAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); CuVSProvider.provider().enableRMMPooledMemory(10, 60); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java index 22af6fb4c3..96c10a0936 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -25,7 +25,7 @@ public class BruteForceRandomizedIT extends CuVSTestCase { @Before public void setup() { - assumeTrue(isLinuxAmd64()); + assumeTrue(isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); CuVSProvider.provider().enableRMMPooledMemory(10, 60); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java index f762334619..423c58db35 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java @@ -36,7 +36,7 @@ public class CagraAceBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index d90b975846..a4f08c47d4 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -47,7 +47,7 @@ public class CagraBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java index 371e842973..17790eb2fc 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java @@ -44,7 +44,7 @@ private interface ResourcesSupplier { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Multi-threaded stability test initialized"); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index 2dcb8b1cd0..9110f6ddc2 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -24,7 +24,7 @@ public class CagraRandomizedIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java index cd3cf43e55..954a7d1e4f 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -25,7 +25,7 @@ public class CuVSMatrixIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java index 00753b42c4..632eefc6a1 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java @@ -20,7 +20,7 @@ public class CuVSResourcesIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java index 62e636eb82..53cbe47392 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java @@ -155,9 +155,10 @@ public boolean equals(Object o) { assertEquals(sortedExpected, sortedActual); } - protected static boolean isLinuxAmd64() { + protected static boolean isLinuxSupportedArch() { String name = System.getProperty("os.name"); - return (name.startsWith("Linux")) && System.getProperty("os.arch").equals("amd64"); + String arch = System.getProperty("os.arch"); + return (name.startsWith("Linux")) && (arch.equals("amd64") || arch.equals("aarch64")); } protected static int[][] createIntMatrix() { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java index 2abb153207..d7464ca97e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java @@ -52,7 +52,7 @@ public class FilterBitsetHandleIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java index fbfab24518..4b9fc045a3 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -19,7 +19,7 @@ public class GPUInfoIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java index 835fcd6e23..351873ce9a 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java @@ -40,7 +40,7 @@ public class HnswAceBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java index 7cb4f0c3fa..78f86bd964 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java @@ -28,7 +28,7 @@ public class HnswBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } // Sample data and query diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java index fe5ca88073..4ace8fc92e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java @@ -28,7 +28,7 @@ public class HnswRandomizedIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java index b80a789025..1d9d40ab27 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java @@ -18,7 +18,7 @@ public class MemoryTrackingResourcesIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java index 62e13b733e..2d1e33484a 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java @@ -40,7 +40,7 @@ public class MultiPartitionCagraSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java index 7965ca34b6..6a47e506ca 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -31,7 +31,7 @@ public class TieredIndexIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); CuVSProvider.provider().enableRMMPooledMemory(10, 60); log.debug("Random context initialized for test"); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java index b6a5664831..3bd8d81464 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal.common; @@ -22,7 +22,7 @@ public class UtilIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java index c40c42204e..076db40553 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.spi; @@ -18,7 +18,7 @@ public class CuVSProviderIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); // Clear sysprop from previous runs/command line System.clearProperty("cuvs.max_version"); } diff --git a/java/cuvs-lucene/pom.xml b/java/cuvs-lucene/pom.xml index 925bd377da..ef31cd37af 100644 --- a/java/cuvs-lucene/pom.xml +++ b/java/cuvs-lucene/pom.xml @@ -123,11 +123,6 @@ SPDX-License-Identifier: Apache-2.0 lucene-backward-codecs 10.2.0 - - org.apache.lucene - lucene-misc - 10.2.0 - org.apache.lucene lucene-test-framework @@ -256,4 +251,37 @@ SPDX-License-Identifier: Apache-2.0 + + + + + skip-compile + + + skip.compile + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + default-compile + none + + + default-testCompile + none + + + + + + + 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 index b57614539a..a5f164b70b 100644 --- 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 @@ -59,7 +59,7 @@ public static enum Strategy { 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.AUTO_SELECT; + 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; @@ -97,22 +97,17 @@ public static enum Strategy { * * @param writerThreads Number of cuVS writer threads to use. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. - * Only consulted under the {@link Strategy#CUSTOM} strategy. - * @param graphdegree The graph degree to use while building the CAGRA index. Only consulted - * under the {@link Strategy#CUSTOM} strategy. + * @param graphdegree The graph degree to use while building the CAGRA index. * @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]. Only - * consulted under the {@link Strategy#CUSTOM} strategy. + * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. - * Only consulted under the {@link Strategy#CUSTOM} strategy. * @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 delegates the CAGRA build parameters to cuVS (derived from the HNSW-equivalent maxConn and beamWidth) or CUSTOM that uses the parameters passed through this class. * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. - * Only consulted under the {@link Strategy#CUSTOM} strategy. * @param hnswHeuristicType the heuristic cuVS applies when deriving the CAGRA build parameters from maxConn and beamWidth under the HEURISTIC strategy. */ private AcceleratedHNSWParams( @@ -157,7 +152,7 @@ public int getWriterThreads() { } /** - * Get the intermediate graph degree. Only consulted under the {@link Strategy#CUSTOM} strategy. + * Get the intermediate graph degree * * @return the graph degree parameter */ @@ -166,7 +161,7 @@ public int getIntermediateGraphDegree() { } /** - * Get the graph degree. Only consulted under the {@link Strategy#CUSTOM} strategy. + * Get the graph degree * * @return the graph degree parameter */ @@ -202,8 +197,7 @@ public int getBeamWidth() { } /** - * Get the CAGRA graph build algorithm. Only consulted under the {@link Strategy#CUSTOM} - * strategy; under {@link Strategy#HEURISTIC} the algorithm is chosen by cuVS. + * Get the CAGRA graph build algorithm * * @return the CAGRA graph build algorithm */ @@ -212,8 +206,7 @@ public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { } /** - * Get the instance of {@link CuVSIvfPqParams}. Only consulted under the {@link - * Strategy#CUSTOM} strategy. + * Get the instance of {@link CuVSIvfPqParams} * * @return the instance of {@link CuVSIvfPqParams} */ @@ -261,8 +254,7 @@ public CuvsDistanceType getCuvsDistanceType() { } /** - * get the number of Iterations to run if building with NN_DESCENT. Only consulted under the - * {@link Strategy#CUSTOM} strategy. + * get the number of Iterations to run if building with NN_DESCENT * * @return the number of iterations for NN_DESCENT */ @@ -347,8 +339,7 @@ public Builder withWriterThreads(int writerThreads) { } /** - * Set the intermediate graph degree to use while building CAGRA index. Only consulted under - * the {@link Strategy#CUSTOM} strategy. + * 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} * @@ -361,8 +352,7 @@ public Builder withIntermediateGraphDegree(int intermediateGraphDegree) { } /** - * Set the graph degree to use while building CAGRA index. Only consulted under the {@link - * Strategy#CUSTOM} strategy. + * 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} * @@ -414,9 +404,8 @@ public Builder withBeamWidth(int beamWidth) { } /** - * Set the CAGRA graph build algorithm to use. Only consulted under the {@link - * Strategy#CUSTOM} strategy; under {@link Strategy#HEURISTIC} the algorithm is chosen by cuVS. - * Default value - AUTO_SELECT + * Set the CAGRA graph build algorithm to use + * Default value - NN_DESCENT * * @param cagraGraphBuildAlgo * @return instance of {@link Builder} @@ -427,8 +416,7 @@ public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo) } /** - * Set the instance of {@link CuVSIvfPqParams}. Only consulted under the {@link - * Strategy#CUSTOM} strategy. + * Set the instance of {@link CuVSIvfPqParams} * * @param cuVSIvfPqParams * @return instance of {@link Builder} @@ -491,8 +479,7 @@ public Builder withCuvsDistanceType(CuvsDistanceType cuvsDistanceType) { } /** - * Set the number of Iterations to run if building with NN_DESCENT. Only consulted under the - * {@link Strategy#CUSTOM} strategy. + * 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} 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 index fdd7dd441f..9c49c07fe0 100644 --- 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 @@ -19,14 +19,8 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; -import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.TaskExecutor; -import org.apache.lucene.store.ByteBuffersDataOutput; -import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -75,7 +69,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens layerAdjacencies.add(adjacencyMatrix); // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } /** @@ -84,25 +78,18 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * (its column count). Ceil is used to accommodate odd graph degrees. * Each layer contains 1/M nodes from the previous layer * Creates layers until the highest layer has ≤ M nodes - *

- * Vectors for higher-layer subsets are read directly from the native matrix - * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])}, - * avoiding any additional heap allocation of the full dataset. Used by both - * the flush and merge paths; the caller provides the vectors as a - * {@link CuVSMatrix}. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, + int size, int dimensions, CuVSMatrix adjacencyListMatrix, - CuVSMatrix vectorDataset, + List vectors, int hnswLayers, CagraIndexParams params, - QuantizationType quantization, - int numThreads) + QuantizationType quantization) throws Throwable { - int size = (int) vectorDataset.size(); int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); // Store all layers data @@ -132,7 +119,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( // Select from previous layer nodes int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); + int idx = random.nextInt(prevLayerNodes.length); + selectedNodesSet.add(prevLayerNodes[idx]); } } @@ -143,22 +131,24 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( layerNodes.add(selectedNodes); if (quantization == QuantizationType.NONE) { - // Read only the sampled rows from the native matrix — no full-dataset heap copy - float[][] selectedVectors = new float[nextLayerSize][dimensions]; + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; for (int i = 0; i < nextLayerSize; i++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); } // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + } else { - // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim. - int bytesPerVector = (int) vectorDataset.columns(); - byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; + + // Extract vectors for selected nodes + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][]; for (int i = 0; i < nextLayerSize; i++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); } // Build CAGRA graph for this layer @@ -176,7 +166,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } /** @@ -194,9 +184,11 @@ private static CuVSMatrix buildCagraGraphForSubset( CuVSMatrix subsetDataset; if (quantization == QuantizationType.BINARY) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); } else { subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } @@ -243,154 +235,56 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) + 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 numLevels = graph.numLevels(); - int[][] offsets = new int[numLevels][]; - - // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is - // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in - // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk - // bytes - // are identical to the fully-serial path (blocks in node order, offsets = per-node byte - // lengths). - // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here - // rather than per level/per task below. - int maxConn = graph.maxConn(); - - int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); - offsets[0] = new int[level0Nodes.length]; - if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { - writeLevel0Parallel( - graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, numThreads); - } else { - writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); - } - - for (int level = 1; level < numLevels; level++) { + 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]; - writeLevelSerial( - graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0, maxConn); - } - // Return offsets (information written while writing the meta info) - return offsets; - } - - /** Node count below which parallel level-0 serialization is not worth the overhead. */ - private static final int PARALLEL_MIN_NODES = 1 << 16; - - /** Nodes per wave — bounds the transient encode buffer regardless of dataset size. */ - private static final int WAVE_NODES = 1 << 20; - - /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ - private static void writeLevelSerial( - GPUBuiltHnswGraph graph, - IndexOutput out, - int level, - int[] sortedNodes, - int[] offsets, - int countOnLevel0, - int maxConn) - throws IOException { - int[] scratch = new int[maxConn * 2]; - int idx = 0; - for (int node : sortedNodes) { - long start = out.getFilePointer(); - encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); - offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); - } - } - - /** - * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node - * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order - * (identical layout to the serial path). - */ - private static void writeLevel0Parallel( - GPUBuiltHnswGraph graph, - IndexOutput out, - int[] nodes, - int[] offsets, - int countOnLevel0, - int maxConn, - int numThreads) - throws IOException { - // invokeAll joins every task before it returns, including tasks still running when a sibling - // fails, so `buffers` is never concatenated while a worker might still be writing into it. - // It also runs one share of each wave on the calling thread instead of parking it, so the - // pool only has to cover the other ranges. - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - TaskExecutor executor = new TaskExecutor(pool); - int n = nodes.length; - for (int waveStart = 0; waveStart < n; waveStart += WAVE_NODES) { - int waveEnd = Math.min(waveStart + WAVE_NODES, n); - int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; - - ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int subStart = waveStart + t * perThread; - final int subEnd = Math.min(subStart + perThread, waveEnd); - final int slot = t; - if (subStart >= subEnd) { - continue; - } - tasks.add( - () -> { - ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); - int[] scratch = new int[maxConn * 2]; - for (int i = subStart; i < subEnd; i++) { - long before = buffer.size(); - encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); - offsets[i] = Math.toIntExact(buffer.size() - before); - } - buffers[slot] = buffer; - return null; - }); + 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; } - executor.invokeAll(tasks); - // Concatenate in thread order (== node order), preserving the serial byte layout. - for (ByteBuffersDataOutput buffer : buffers) { - if (buffer != null) { - buffer.copyTo(out); + // 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]; } - } - } finally { - pool.shutdown(); - } - } - - /** - * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt - * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. - */ - private static void encodeNode( - NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) - throws IOException { - int size = neighbors == null ? 0 : neighbors.size(); - int actualSize = 0; - if (size > 0) { - int[] nnodes = neighbors.nodes(); - Arrays.sort(nnodes, 0, size); - scratch[0] = nnodes[0]; - actualSize = 1; - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - if (nnodes[i - 1] == nnodes[i]) { - continue; + // 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]); } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + offsets[level][nodeOffsetId++] = + Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } } - out.writeVInt(actualSize); - for (int i = 0; i < actualSize; i++) { - out.writeVInt(scratch[i]); - } + // Return offsets (information written while writing the meta info) + return offsets; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java deleted file mode 100644 index e5317f4a01..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java +++ /dev/null @@ -1,214 +0,0 @@ -/* - * 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.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.getCuVSResourcesInstance; - -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.Closeable; -import java.io.IOException; -import java.util.List; -import org.apache.lucene.codecs.CodecUtil; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.IndexFileNames; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.util.IOUtils; - -/** - * Owns the {@code .vem}/{@code .vex} outputs and builds/writes the CAGRA-derived HNSW graph for a - * field. Shared by both {@link Lucene99AcceleratedHNSWVectorsWriter} (heap-buffered flat path) - * and {@link NativeFlatBufferedHNSWVectorsWriter} (native flat-buffered path) — the only - * difference between those two writers is how the flat {@code .vec}/{@code .vemf} files and the - * dataset handed to {@link #writeField(FieldInfo, CuVSMatrix)} are produced. - */ -final class AcceleratedHnswGraphOutput implements Closeable { - - private static final LuceneProvider LUCENE_PROVIDER; - private static final Integer VERSION_CURRENT; - - static { - try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); - VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); - } catch (Exception e) { - throw new ExceptionInInitializerError(e.getMessage()); - } - } - - private final AcceleratedHNSWParams acceleratedHNSWParams; - private final IndexOutput hnswMeta; - private final IndexOutput hnswVectorIndex; - private boolean finished; - - AcceleratedHnswGraphOutput(SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams) - throws IOException { - this.acceleratedHNSWParams = acceleratedHNSWParams; - String vemFileName = - IndexFileNames.segmentFileName( - state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); - String vexFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); - IndexOutput meta = null; - IndexOutput vectorIndex = null; - boolean success = false; - try { - meta = state.directory.createOutput(vemFileName, state.context); - vectorIndex = state.directory.createOutput(vexFileName, state.context); - CodecUtil.writeIndexHeader( - meta, - HNSW_META_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - CodecUtil.writeIndexHeader( - vectorIndex, - HNSW_INDEX_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - success = true; - } finally { - this.hnswMeta = meta; - this.hnswVectorIndex = vectorIndex; - if (success == false) { - IOUtils.closeWhileHandlingException(this); - } - } - } - - /** - * Flush/sorting path: builds a host matrix from the heap vectors, then delegates to {@link - * #writeField(FieldInfo, CuVSMatrix)}. - */ - void writeField(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); - return; - } - if (vectors.size() < 2) { - writeSingleVectorGraph(fieldInfo, vectors); - return; - } - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeField(fieldInfo, dataset); - } - - /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index. Single - * implementation used by both the flush and merge paths of both writers. The dataset is a - * {@link CuVSMatrix} (host-backed on the merge/native paths) so the full set of vectors is - * never double-materialised on the Java heap. - */ - void writeField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try (dataset) { - int size = (int) dataset.size(); - if (size == 0) { - writeEmpty(fieldInfo, hnswMeta); - return; - } - if (size < 2) { - float[] buf = new float[fieldInfo.getVectorDimension()]; - dataset.getRow(0).toArray(buf); - writeSingleVectorGraph(fieldInfo, List.of(buf)); - return; - } - try { - CagraIndexParams params = - CagraIndexParamsFactory.create( - acceleratedHNSWParams, dataset.size(), dataset.columns()); - try (CagraIndex cagraIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(dataset) - .withIndexParams(params) - .build()) { - CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int dimensions = fieldInfo.getVectorDimension(); - GPUBuiltHnswGraph hnswGraph = - createMultiLayerHnswGraph( - fieldInfo, - dimensions, - adjacencyListMatrix, - dataset, - acceleratedHNSWParams.getHnswLayers(), - params, - QuantizationType.NONE, - acceleratedHNSWParams.getWriterThreads()); - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets); - } - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - } - - 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, acceleratedHNSWParams.getWriterThreads()); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets); - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - - void finish() throws IOException { - if (finished) { - throw new IllegalStateException("already finished"); - } - finished = true; - if (hnswMeta != null) { - // write end of fields marker - hnswMeta.writeInt(-1); - CodecUtil.writeFooter(hnswMeta); - } - if (hnswVectorIndex != null) { - CodecUtil.writeFooter(hnswVectorIndex); - } - } - - @Override - public void close() throws IOException { - IOUtils.close(hnswMeta, hnswVectorIndex); - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java deleted file mode 100644 index c1cd5e71b4..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java +++ /dev/null @@ -1,660 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import java.io.Closeable; -import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.ArrayList; -import java.util.Comparator; -import java.util.List; -import java.util.Objects; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.Semaphore; -import java.util.stream.Stream; -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.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.IndexableField; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.misc.store.HardlinkCopyDirectoryWrapper; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; - -/** - * Builds a GPU-accelerated CAGRA/HNSW Lucene index from data the caller already has ready on - * local disk, as few segments as possible, without going through a generically-configurable - * {@link org.apache.lucene.codecs.Codec}. - * - *

This is not a general-purpose Lucene extension point: an instance owns a single - * {@link IndexWriter} and its {@link IndexWriterConfig} so that the invariants the underlying GPU - * writer requires (single unmerged segment, no index sort, exact vector count) are guaranteed by - * this class rather than left to a caller to get right. If you need a general Lucene codec that - * any indexing pipeline (including ones that don't control their own {@code IndexWriter} - * lifecycle, e.g. Solr or Elasticsearch) can register, use {@link Lucene101AcceleratedHNSWCodec} - * directly instead. - * - *

Two ways to use this class: - * - *

    - *
  • Manual, single segment: construct an instance directly, call {@link #addDocument} - * per document exactly like a plain {@link IndexWriter}, then {@link #close}. You own the - * loop and the {@link Document} you build (any fields, not just the vector). - *
  • One-shot, one or many segments: {@link #indexFbin} / {@link #build(VectorSource, - * Config)} own the loop for you — they read vectors from a {@code .fbin} file or {@link - * VectorSource}, optionally split into {@code numSegments} partitions (sequential or - * overlapped), and combine the result. Since they build each row's {@link Document} - * internally, an optional {@link FieldCallback} lets you add extra fields to it. - *
- * - *

Scope: CAGRA_HNSW (GPU build, CPU search) only. This class builds indexes for {@link - * Lucene101AcceleratedHNSWCodec}. It does not support the GPU-search codec ({@code - * CuVS2510GPUSearchCodec}) — that writer does not (yet) have the native flat-buffering - * optimization this class relies on. Reading an existing index for CAGRA_SEARCH-style GPU search - * is unaffected by this class either way; only bulk-building one is out of scope for now. - */ -public final class CagraHnswBulkIndexWriter implements Closeable { - - private static final int DEFAULT_CHUNK_SIZE_MB = 32; - - private final IndexWriter writer; - private final int exactVectorCount; - private int documentsAdded; - private boolean closed; - - /** - * Opens a single-segment, native-flat-buffered writer. {@code exactVectorCount} must equal the - * number of {@link #addDocument} calls that will follow — the native buffer is pre-sized to it, - * so {@link #close} rejects a mismatched count rather than let the underlying writer produce a - * corrupt or incomplete segment. - * - *

{@code conf}'s {@code Analyzer}, {@code Similarity}, {@code InfoStream}, and {@code - * OpenMode} are honored; an explicit {@code IndexSort} is rejected ({@link - * IllegalArgumentException}), since native flat buffering does not support index-sorted - * segments. Codec, merge policy, and flush thresholds are always owned by this class regardless - * of what {@code conf} contains — {@link IndexWriterConfig} does not expose whether the caller - * explicitly set those or left them at Lucene's defaults, so there is no reliable way to - * validate-and-reject a caller-supplied value for them the way {@code IndexSort} can be; this - * class simply never reads them from {@code conf}. - * - *

{@code config.targetDirectory()}, {@code config.numSegments()}, {@code - * config.overlapped()}, and {@code config.pipelineDepth()} are not consulted here — {@code - * directory} is passed explicitly, and this constructor always builds exactly one segment. Those - * fields only matter to {@link #indexFbin} / {@link #build(VectorSource, Config)}. - */ - public CagraHnswBulkIndexWriter( - Directory directory, IndexWriterConfig conf, Config config, int exactVectorCount) - throws Exception { - Objects.requireNonNull(directory, "directory"); - Objects.requireNonNull(conf, "conf"); - Objects.requireNonNull(config, "config"); - if (exactVectorCount <= 0) { - throw new IllegalArgumentException("exactVectorCount must be > 0, got " + exactVectorCount); - } - if (conf.getIndexSort() != null) { - throw new IllegalArgumentException( - "CagraHnswBulkIndexWriter does not support an index-sorted segment (native flat" - + " buffering requires an unsorted single-segment build); leave" - + " IndexWriterConfig.indexSort unset"); - } - this.exactVectorCount = exactVectorCount; - - Codec codec = new Lucene101AcceleratedHNSWCodec(config.graphBuildParams(), exactVectorCount); - IndexWriterConfig ownedConf = - new IndexWriterConfig(conf.getAnalyzer()) - .setSimilarity(conf.getSimilarity()) - .setInfoStream(conf.getInfoStream()) - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(Math.max(2, exactVectorCount + 1)) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE) - .setOpenMode(conf.getOpenMode()); - this.writer = new IndexWriter(directory, ownedConf); - } - - /** - * Adds one document, exactly like {@link IndexWriter#addDocument}. {@code doc} may contain any - * fields — the vector field (matching {@link Config#fieldName()}) is routed into the native - * flat buffer automatically by the underlying codec, the same way any {@link - * KnnFloatVectorField} is for any Lucene codec; every other field is indexed normally. - */ - public long addDocument(Iterable doc) throws IOException { - if (closed) { - throw new IllegalStateException("addDocument called after close()"); - } - if (documentsAdded >= exactVectorCount) { - throw new IllegalStateException( - "addDocument called more than exactVectorCount (" + exactVectorCount + ") times"); - } - long seqNo = writer.addDocument(doc); - documentsAdded++; - return seqNo; - } - - /** - * Runs the single native-buffered flush (this is where the GPU CAGRA build happens) and closes - * the underlying writer. There is no separate {@code commit()}: unlike a plain {@link - * IndexWriter}, only one flush is ever valid for this instance, so folding it into {@code - * close()} removes any way to trigger it early with fewer than {@code exactVectorCount} vectors - * added. Throws {@link IllegalStateException} if {@link #addDocument} was called fewer times - * than {@code exactVectorCount} promised, discarding the buffered documents rather than - * flushing an under-filled segment. - */ - @Override - public void close() throws IOException { - if (closed) { - return; - } - closed = true; - if (documentsAdded != exactVectorCount) { - IllegalStateException mismatch = - new IllegalStateException( - "expected " + exactVectorCount + " documents, got " + documentsAdded); - // The native host matrix is preallocated for exactly exactVectorCount rows, so an - // under-filled buffer can only be discarded, never flushed. rollback() discards it and - // closes the IndexWriter, and is also what frees that preallocated memory: it aborts the - // indexing chain, which closes the vectors writer. Leaving without it leaks the buffer, - // which lives in a shared Arena and is never reclaimed by the GC. - try { - writer.rollback(); - } catch (Throwable t) { - mismatch.addSuppressed(t); - } - throw mismatch; - } - try (writer) { - writer.commit(); - } - } - - /** - * Discards everything added so far without producing a segment: no GPU build runs, nothing is - * committed, and the preallocated native buffer is released (see {@link #close()} for why that - * release is what matters). Unlike {@link #close()} this does not care how many documents were - * added, so it is the cleanup to use on a failure path -- a partially filled buffer is not - * worth building, and reporting its count would only bury the failure that caused it. - * Idempotent, and a no-op once {@link #close()} has run. - */ - void abort() throws IOException { - if (closed) { - return; - } - closed = true; - writer.rollback(); - } - - /** - * Callback invoked once per row by {@link #indexFbin} / {@link #build(VectorSource, Config, - * FieldCallback)}, right after the id and vector fields have been added to {@code document} and - * right before it is added to the index — lets the caller attach additional fields (metadata) - * per vector. Not used by the manual, direct-instance API, where the caller already builds the - * whole {@link Document} themselves. - */ - @FunctionalInterface - public interface FieldCallback { - void addFields(Document document, int id) throws IOException; - } - - /** - * Builds an index from {@code source} into {@code config.targetDirectory()}, partitioned into - * {@code config.numSegments()} contiguous slices as {@code source} is consumed front-to-back. - * - *

{@code config.overlapped()} is not supported here: a single {@link VectorSource} is - * forward-only/single-consumer (see its contract) and so cannot be read concurrently by - * multiple slice builders. Use {@link #indexFbin} for the overlapped pipeline, which knows how - * to open independent, per-slice sources over the same underlying file. - */ - public static void build(VectorSource source, Config config) throws Exception { - build(source, config, null); - } - - /** As {@link #build(VectorSource, Config)}, with a {@link FieldCallback} for extra fields. */ - public static void build(VectorSource source, Config config, FieldCallback callback) - throws Exception { - Objects.requireNonNull(source, "source"); - Objects.requireNonNull(config, "config"); - if (config.overlapped()) { - throw new IllegalArgumentException( - "Config.overlapped() is not supported by build(VectorSource, Config): a VectorSource is" - + " forward-only and cannot be read by multiple concurrent slice builders. Use" - + " indexFbin(...) for the overlapped multi-segment build."); - } - if (source.dimensions() != config.dimensions()) { - throw new IllegalArgumentException( - "source.dimensions() (" - + source.dimensions() - + ") does not match Config.dimensions() (" - + config.dimensions() - + ")"); - } - List slices = sliceEvenly(source.size(), config.numSegments()); - buildSequential(source, config, callback, slices); - } - - /** - * Convenience entry point: builds an index directly from a {@code .fbin} file ({@code - * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows), using - * {@link FbinVectorSource} with {@link #DEFAULT_CHUNK_SIZE_MB}-sized prefetched chunks. - */ - public static void indexFbin(Path fbinPath, Config config) throws Exception { - indexFbin(fbinPath, config, null, DEFAULT_CHUNK_SIZE_MB); - } - - /** As {@link #indexFbin(Path, Config)}, with a {@link FieldCallback} for extra fields. */ - public static void indexFbin(Path fbinPath, Config config, FieldCallback callback) - throws Exception { - indexFbin(fbinPath, config, callback, DEFAULT_CHUNK_SIZE_MB); - } - - /** As {@link #indexFbin(Path, Config)}, with an explicit prefetch chunk size. */ - public static void indexFbin(Path fbinPath, Config config, int chunkSizeMB) throws Exception { - indexFbin(fbinPath, config, null, chunkSizeMB); - } - - /** - * As {@link #indexFbin(Path, Config)}, with an explicit prefetch chunk size and {@link - * FieldCallback}. When {@code config.numSegments() > 1} and {@code config.overlapped()}, builds - * up to {@code config.pipelineDepth()} segments concurrently, each over its own slice of {@code - * fbinPath}, then combines them by hardlink; otherwise builds sequentially over one shared - * reader. - */ - public static void indexFbin( - Path fbinPath, Config config, FieldCallback callback, int chunkSizeMB) throws Exception { - Objects.requireNonNull(fbinPath, "fbinPath"); - Objects.requireNonNull(config, "config"); - int total; - int dim; - try (FbinVectorSource probe = new FbinVectorSource(fbinPath, 1)) { - total = probe.size(); - dim = probe.dimensions(); - } - if (dim != config.dimensions()) { - throw new IllegalArgumentException( - "fbinPath dimension (" - + dim - + ") does not match Config.dimensions() (" - + config.dimensions() - + ")"); - } - List slices = sliceEvenly(total, config.numSegments()); - if (config.overlapped() && slices.size() > 1) { - buildOverlapped(fbinPath, config, callback, chunkSizeMB, slices, dim); - } else { - try (FbinVectorSource source = new FbinVectorSource(fbinPath, chunkSizeMB)) { - buildSequential(source, config, callback, slices); - } - } - } - - /** - * Sequential partitioned build: {@code source} is streamed front-to-back across all slices, - * each slice built as a single native-flat segment appended to the same directory (first slice - * {@code CREATE}, later slices {@code APPEND}). Peak host memory is one slice's native buffer. - */ - private static void buildSequential( - VectorSource source, Config config, FieldCallback callback, List slices) - throws Exception { - float[] scratch = new float[config.dimensions()]; - try (Directory dir = FSDirectory.open(config.targetDirectory())) { - for (int p = 0; p < slices.size(); p++) { - int[] slice = slices.get(p); - buildSegment( - dir, source, scratch, config, callback, slice[0], slice[0], slice[1], p == 0, null); - } - } - } - - /** - * Overlapped partitioned build: a bounded pool builds up to {@code config.pipelineDepth()} - * segments at once, each with its OWN {@link FbinVectorSource} over just its slice, so a - * segment's ingest overlaps a prior segment's GPU commit. The GPU build itself is serialized on - * a single permit. The finished per-segment indexes are combined into {@code - * config.targetDirectory()} by hardlinking their files (no bulk copy of the vector data). - */ - private static void buildOverlapped( - Path fbinPath, - Config config, - FieldCallback callback, - int chunkSizeMB, - List slices, - int dim) - throws Exception { - Path targetDir = config.targetDirectory(); - int depth = Math.min(slices.size(), config.pipelineDepth()); - List segDirs = new ArrayList<>(); - for (int p = 0; p < slices.size(); p++) { - segDirs.add(targetDir.resolveSibling(targetDir.getFileName() + "_p" + p)); - } - for (Path segDir : segDirs) { - deleteRecursivelyQuietly(segDir); - } - try { - Semaphore gpuPermit = new Semaphore(1); // serialize the GPU CAGRA build across segments - ExecutorService pool = Executors.newFixedThreadPool(depth); - List> futures = new ArrayList<>(); - for (int p = 0; p < slices.size(); p++) { - int[] slice = slices.get(p); - Path segDir = segDirs.get(p); - futures.add( - pool.submit( - () -> { - float[] scratch = new float[dim]; - try (FbinVectorSource source = - new FbinVectorSource(fbinPath, slice[0], slice[1], chunkSizeMB); - Directory d = FSDirectory.open(segDir)) { - // createNew=true: each segment is a fresh single-segment index in its own - // dir; sourceStart=0 since this source is already windowed to the slice. - buildSegment( - d, source, scratch, config, callback, 0, slice[0], slice[1], true, - gpuPermit); - } - return null; - })); - } - pool.shutdown(); - try { - for (Future f : futures) { - f.get(); // propagate any build failure - } - } catch (Exception e) { - throw new IOException("Overlapped bulk index build failed", e); - } finally { - pool.shutdownNow(); - } - combineByHardlink(targetDir, segDirs); - } finally { - for (Path segDir : segDirs) { - deleteRecursivelyQuietly(segDir); - } - } - } - - /** - * Builds one segment from {@code size} vectors: {@code source.get(sourceStart + i, ...)} for - * {@code i} in {@code [0, size)}, labelled with ids {@code idStart + i} (the vectors' absolute - * position in the overall build, regardless of whether {@code source} itself is windowed). Opens - * a {@link CagraHnswBulkIndexWriter} sized to {@code size} and drives it exactly like the manual - * API. When {@code gpuPermit} is non-null the close (which runs the GPU CAGRA build) is - * serialized on it while other segments' host-side ingest may proceed. - */ - private static void buildSegment( - Directory dir, - VectorSource source, - float[] scratch, - Config config, - FieldCallback callback, - int sourceStart, - int idStart, - int size, - boolean createNew, - Semaphore gpuPermit) - throws Exception { - IndexWriterConfig conf = - new IndexWriterConfig() - .setOpenMode( - createNew ? IndexWriterConfig.OpenMode.CREATE : IndexWriterConfig.OpenMode.APPEND); - CagraHnswBulkIndexWriter writer = new CagraHnswBulkIndexWriter(dir, conf, config, size); - try { - for (int i = 0; i < size; i++) { - source.get(sourceStart + i, scratch); - int id = idStart + i; - Document doc = new Document(); - if (config.idFieldName() != null) { - doc.add(new StringField(config.idFieldName(), Integer.toString(id), Field.Store.YES)); - } - doc.add(new KnnFloatVectorField(config.fieldName(), scratch, config.similarity())); - if (callback != null) { - callback.addFields(doc, id); - } - writer.addDocument(doc); // copies the vector -> 'scratch' is safe to reuse next iteration - } - // The single flush inside close() is where the GPU CAGRA build runs; serialize it if asked. - if (gpuPermit != null) { - gpuPermit.acquire(); - try { - writer.close(); - } finally { - gpuPermit.release(); - } - } else { - writer.close(); - } - } catch (Throwable t) { - // Cleanup must not become the reported failure. abort() discards the partially filled - // buffer rather than trying to build it, and anything it throws on the way out is attached - // to the original exception instead of replacing it. If the close() above is what failed, - // this is a no-op -- that writer has already released itself. - try { - writer.abort(); - } catch (Throwable cleanupFailure) { - t.addSuppressed(cleanupFailure); - } - throw t; - } - } - - /** - * Combines the per-segment indexes into {@code targetDir} by hardlinking their files (same - * filesystem) rather than copying the vector data. {@link HardlinkCopyDirectoryWrapper} falls - * back to a byte copy automatically if the segment dirs and the final dir are on different - * filesystems. - */ - private static void combineByHardlink(Path targetDir, List segDirs) throws IOException { - Directory[] sources = new Directory[segDirs.size()]; - try { - for (int i = 0; i < segDirs.size(); i++) { - sources[i] = FSDirectory.open(segDirs.get(i)); - } - IndexWriterConfig iwc = - new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE); // keep segments separate - try (Directory target = new HardlinkCopyDirectoryWrapper(FSDirectory.open(targetDir)); - IndexWriter combiner = new IndexWriter(target, iwc)) { - combiner.addIndexes(sources); - } - } finally { - for (Directory s : sources) { - if (s != null) { - s.close(); - } - } - } - } - - /** Splits {@code total} into {@code k} contiguous [start, size] slices, spreading the remainder. */ - private static List sliceEvenly(int total, int k) { - List slices = new ArrayList<>(); - int base = total / k; - int rem = total % k; - int start = 0; - for (int p = 0; p < k; p++) { - int size = base + (p < rem ? 1 : 0); // spread the remainder over the first slices - if (size <= 0) { - continue; - } - slices.add(new int[] {start, size}); - start += size; - } - return slices; - } - - private static void deleteRecursivelyQuietly(Path path) { - if (!Files.exists(path)) { - return; - } - try (Stream walk = Files.walk(path)) { - walk.sorted(Comparator.reverseOrder()) - .forEach( - p -> { - try { - Files.deleteIfExists(p); - } catch (IOException ignored) { - // best-effort cleanup of a temp per-segment dir - } - }); - } catch (IOException ignored) { - // best-effort cleanup of a temp per-segment dir - } - } - - /** Immutable configuration for {@link CagraHnswBulkIndexWriter}. */ - public static final class Config { - private final String fieldName; - private final int dimensions; - private final VectorSimilarityFunction similarity; - private final String idFieldName; - private final AcceleratedHNSWParams graphBuildParams; - private final Path targetDirectory; - private final int numSegments; - private final boolean overlapped; - private final int pipelineDepth; - - private Config(Builder b) { - this.fieldName = b.fieldName; - this.dimensions = b.dimensions; - this.similarity = b.similarity; - this.idFieldName = b.idFieldName; - this.graphBuildParams = b.graphBuildParams; - this.targetDirectory = b.targetDirectory; - this.numSegments = b.numSegments; - this.overlapped = b.overlapped; - this.pipelineDepth = b.pipelineDepth; - } - - public String fieldName() { - return fieldName; - } - - public int dimensions() { - return dimensions; - } - - public VectorSimilarityFunction similarity() { - return similarity; - } - - public String idFieldName() { - return idFieldName; - } - - public AcceleratedHNSWParams graphBuildParams() { - return graphBuildParams; - } - - /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ - public Path targetDirectory() { - return targetDirectory; - } - - /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ - public int numSegments() { - return numSegments; - } - - /** Only consulted by {@link #indexFbin}. */ - public boolean overlapped() { - return overlapped; - } - - /** Only consulted by {@link #indexFbin}. */ - public int pipelineDepth() { - return pipelineDepth; - } - - public static Builder builder() { - return new Builder(); - } - - /** Builder for {@link Config}. */ - public static final class Builder { - private String fieldName = "vector"; - private int dimensions = -1; - private VectorSimilarityFunction similarity = VectorSimilarityFunction.EUCLIDEAN; - private String idFieldName = "id"; - private AcceleratedHNSWParams graphBuildParams; - private Path targetDirectory; - private int numSegments = 1; - private boolean overlapped = false; - private int pipelineDepth = 2; - - /** Sets the vector field name and dimensionality; required. */ - public Builder field(String fieldName, int dimensions, VectorSimilarityFunction similarity) { - this.fieldName = Objects.requireNonNull(fieldName, "fieldName"); - this.dimensions = dimensions; - this.similarity = Objects.requireNonNull(similarity, "similarity"); - return this; - } - - /** - * Sets the stored id field name (holding each document's absolute position in the build, as - * a string), added automatically by {@link #indexFbin} / {@link #build(VectorSource, - * Config)} before the {@link FieldCallback} runs. Defaults to {@code "id"}; pass {@code - * null} to disable. Not used by the manual, direct-instance API. - */ - public Builder idField(String idFieldName) { - this.idFieldName = idFieldName; - return this; - } - - /** Sets the CAGRA/HNSW graph-build parameters; required. */ - public Builder graphBuild(AcceleratedHNSWParams graphBuildParams) { - this.graphBuildParams = Objects.requireNonNull(graphBuildParams, "graphBuildParams"); - return this; - } - - /** Sets the directory the final index is written to; required for {@link #indexFbin}/{@link #build}. */ - public Builder targetDirectory(Path targetDirectory) { - this.targetDirectory = Objects.requireNonNull(targetDirectory, "targetDirectory"); - return this; - } - - /** - * Splits the build into {@code numSegments} contiguous slices, each a single native-flat - * segment. Peak host memory scales as {@code 1/numSegments}; the GPU build itself is always - * serialized across slices regardless of this setting. Default 1 (single segment). - * - * @param overlapped when {@code numSegments > 1} and building via {@link #indexFbin}, - * builds up to {@link #pipelineDepth} slices concurrently (ingest of one overlapping the - * GPU commit of another) instead of strictly sequentially. Ignored by {@link - * #build(VectorSource, Config)}, which always builds sequentially — see that method's - * Javadoc. - */ - public Builder segments(int numSegments, boolean overlapped) { - if (numSegments < 1) { - throw new IllegalArgumentException("numSegments must be >= 1, got " + numSegments); - } - this.numSegments = numSegments; - this.overlapped = overlapped; - return this; - } - - /** Max segments built concurrently in overlap mode; peak host memory is this many slice buffers. */ - public Builder pipelineDepth(int pipelineDepth) { - if (pipelineDepth < 1) { - throw new IllegalArgumentException("pipelineDepth must be >= 1, got " + pipelineDepth); - } - this.pipelineDepth = pipelineDepth; - return this; - } - - public Config build() { - if (dimensions <= 0) { - throw new IllegalStateException( - "field(...) must be called with a positive dimension count"); - } - Objects.requireNonNull(graphBuildParams, "graphBuild(...) must be called"); - return new Config(this); - } - } - } -} 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 index ef0b82b3cf..8bb46faf61 100644 --- 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 @@ -75,8 +75,7 @@ public static CagraIndexParams create( AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { // Delegate the derivation of the graph degrees, build algorithm and its parameters to cuVS, - // expressed in terms of the HNSW-equivalent maxConn/beamWidth. cagraGraphBuildAlgo is not - // consulted here; it only applies under CUSTOM. + // expressed in terms of the HNSW-equivalent maxConn/beamWidth. CagraIndexParams derived = CagraIndexParams.fromHnswParams( rows, @@ -85,9 +84,9 @@ public static CagraIndexParams create( acceleratedHNSWParams.getBeamWidth(), acceleratedHNSWParams.getHnswHeuristicType(), acceleratedHNSWParams.getCuvsDistanceType()); - // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS - // default (not a heuristic value). We can rebuild the CagraIndexParams with the - // caller-supplied writerThreads for now but should fix this in cuVS in the future. + // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS default + // (not a heuristic value). We can rebuild the CagraIndexParams with the caller-supplied + // writerThreads for now but should fix this in cuVS in the future. return new CagraIndexParams.Builder() .withGraphDegree(derived.getGraphDegree()) .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) @@ -98,7 +97,6 @@ public static CagraIndexParams create( .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) .build(); } - // CUSTOM: forward the caller's algorithm and the parameters it consumes. return new CagraIndexParams.Builder() .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) 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 index 3a7eb7852e..5feb211afe 100644 --- 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 @@ -202,7 +202,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { CuVSMatrix cagraDataset = - Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. @@ -223,7 +224,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isBruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = - Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java deleted file mode 100644 index 548874d9d6..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java +++ /dev/null @@ -1,213 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Path; -import java.nio.file.StandardOpenOption; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; - -/** - * Prefetching, double-buffered {@link VectorSource} over an uncompressed {@code .fbin} file - * ({@code [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). - * - *

Opens the file ONCE. A background thread reads the (optionally sliced) range front-to-back - * into two reusable direct buffers: while the caller consumes the current chunk, the reader fills - * the next one, so the disk read overlaps with the caller's per-vector work. {@link #get(int, - * float[])} unpacks directly into a caller-supplied array (no per-vector allocation). - * - *

Forward-only, single-consumer. {@code get} must be called with strictly increasing - * indices from a single thread, relative to the window this instance was constructed over — index 0 is - * the first vector in that window, not necessarily the first vector in the file. To read only a - * slice of a larger file (e.g. one segment of a partitioned build), construct with a {@code - * [firstVector, count)} range so each instance streams just its own portion of the file. - */ -public final class FbinVectorSource implements VectorSource { - - private static final long HEADER_BYTES = 8; - - private static final class Chunk { - final ByteBuffer buf; - final long start; - final int len; - - Chunk(ByteBuffer buf, long start, int len) { - this.buf = buf; - this.start = start; - this.len = len; - } - } - - /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ - private static final Chunk POISON = new Chunk(null, -1, 0); - - private final FileChannel channel; - private final int dimension; - private final int firstVector; // absolute index (in the file) of the first vector served - private final int windowSize; // number of vectors this instance serves - private final int vectorBytes; - private final int chunkVectors; - - private final BlockingQueue free = new ArrayBlockingQueue<>(2); - private final BlockingQueue ready = new ArrayBlockingQueue<>(2); - private final Thread reader; - private volatile IOException readerError; - - private Chunk current; // consumer-owned; the chunk currently being served - private long nextExpectedRelative; // forward-only guard, relative index - - /** Reads the whole file from index 0. */ - public FbinVectorSource(Path path, int chunkSizeMB) throws IOException { - this(path, 0, -1, chunkSizeMB); - } - - /** - * Reads the contiguous range {@code [firstVector, firstVector + count)} of {@code path}, or to - * end of file if {@code count <= 0}. {@link #get} then serves indices relative to {@code - * firstVector} (0 = {@code firstVector}). - */ - public FbinVectorSource(Path path, int firstVector, int count, int chunkSizeMB) - throws IOException { - this.channel = FileChannel.open(path, StandardOpenOption.READ); - ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); - readFully(header, 0); - header.flip(); - int numVectors = header.getInt(); - this.dimension = header.getInt(); - this.vectorBytes = dimension * Float.BYTES; - this.firstVector = firstVector; - int endVector = count > 0 ? (int) Math.min((long) firstVector + count, numVectors) : numVectors; - this.windowSize = Math.max(0, endVector - firstVector); - - long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; - int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int - this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); - - // Two reusable direct buffers: the reader fills one while the consumer drains the other. - for (int i = 0; i < 2; i++) { - free.add( - ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); - } - - this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); - this.reader.setDaemon(true); - this.reader.start(); - } - - @Override - public int dimensions() { - return dimension; - } - - @Override - public int size() { - return windowSize; - } - - /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ - private void readLoop() { - long next = firstVector; - long endVector = firstVector + windowSize; - try { - while (next < endVector) { - ByteBuffer buf = free.take(); - int toRead = (int) Math.min(chunkVectors, endVector - next); - buf.clear(); - buf.limit(toRead * vectorBytes); - readFully(buf, HEADER_BYTES + next * (long) vectorBytes); - ready.put(new Chunk(buf, next, toRead)); - next += toRead; - } - ready.put(POISON); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); // close() requested; stop quietly - } catch (IOException e) { - readerError = e; - try { - ready.put(POISON); // unblock the consumer so it can observe the error - } catch (InterruptedException ignored) { - Thread.currentThread().interrupt(); - } - } - } - - private void advance() throws IOException { - if (current != null) { - try { - free.put(current.buf); // hand the drained buffer back to the reader - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted returning chunk buffer", e); - } - current = null; - } - Chunk next; - try { - next = ready.take(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted awaiting next chunk", e); - } - if (next == POISON) { - if (readerError != null) { - throw new IOException("Prefetch reader failed", readerError); - } - throw new IOException("No more chunks available (unexpected EOF in prefetch)"); - } - current = next; - } - - @Override - public void get(int index, float[] dst) throws IOException { - if (index < 0 || index >= windowSize) { - throw new IndexOutOfBoundsException( - "Index " + index + " out of bounds [0, " + windowSize + ")"); - } - if (index < nextExpectedRelative) { - throw new UnsupportedOperationException( - "FbinVectorSource requires forward-only sequential access; got index " - + index - + " before previously served index " - + (nextExpectedRelative - 1)); - } - nextExpectedRelative = index + 1; - long absolute = firstVector + index; - while (current == null || absolute >= current.start + current.len) { - advance(); - } - int base = (int) (absolute - current.start) * vectorBytes; - for (int i = 0; i < dimension; i++) { - dst[i] = current.buf.getFloat(base + i * Float.BYTES); - } - } - - /** Convenience allocating variant (e.g. for a one-off query vector). */ - public float[] get(int index) throws IOException { - float[] dst = new float[dimension]; - get(index, dst); - return dst; - } - - private void readFully(ByteBuffer buf, long position) throws IOException { - long pos = position; - while (buf.hasRemaining()) { - int n = channel.read(buf, pos); - if (n < 0) { - throw new IOException("Unexpected EOF reading at position " + pos); - } - pos += n; - } - } - - @Override - public void close() throws IOException { - reader.interrupt(); - channel.close(); - } -} 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 index 1898f5b985..8e2d9a70e2 100644 --- 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 @@ -23,7 +23,6 @@ public class FieldWriter extends KnnFieldVectorsWriter { RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); private final FieldInfo fieldInfo; - private final int dimension; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; private QuantizationType quantizationType; @@ -35,7 +34,6 @@ public FieldWriter( FlatFieldVectorsWriter flatFieldVectorsWriter) { this.quantizationType = quantizationType; this.fieldInfo = fieldInfo; - this.dimension = fieldInfo.getVectorDimension(); this.flatFieldVectorsWriter = (FlatFieldVectorsWriter) flatFieldVectorsWriter; } @@ -73,10 +71,6 @@ DocsWithFieldSet getDocsWithFieldSet() { return flatFieldVectorsWriter.getDocsWithFieldSet(); } - int dimension() { - return dimension; - } - @Override public Object copyValue(Object vectorValue) { throw new UnsupportedOperationException(); 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 index 2007bc5c30..7e9f888e32 100644 --- 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 @@ -6,17 +6,10 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; -import java.io.IOException; import java.util.ArrayList; import java.util.List; -import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -45,15 +38,9 @@ public class GPUBuiltHnswGraph extends HnswGraph { * @param dimensions the vector dimension * @param layerNodes the nodes on the layer * @param layerAdjacencies adjacency list - * @param numThreads threads to use for materializing the adjacency (1 = serial) */ public GPUBuiltHnswGraph( - int size, - int dimensions, - List layerNodes, - List layerAdjacencies, - int numThreads) - throws IOException { + int size, int dimensions, List layerNodes, List layerAdjacencies) { this.size = size; this.dimensions = dimensions; @@ -63,102 +50,38 @@ public GPUBuiltHnswGraph( // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size, numThreads); + 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, numThreads)); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); } } - /** Node count below which parallel materialization is not worth the thread overhead. */ - private static final int PARALLEL_MIN_NODES = 1 << 16; - /** - * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. - * - *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe - * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose - * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so - * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node - * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are - * read directly in both paths. + * Fills the neighbor array using the adjacency matrix. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes - * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) * @return the NeighborArray */ - private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) - throws IOException { + private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { NeighborArray[] neighbors = new NeighborArray[size]; - if (numThreads <= 1 || size < PARALLEL_MIN_NODES) { - fillNeighborRange(adjacency, neighbors, 0, size); - return neighbors; - } - CuVSMatrix source = adjacency; - CuVSHostMatrix hostCopy = null; - if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { - hostCopy = deviceAdjacency.toHost(); - source = hostCopy; - } - try { - fillNeighborArrayParallel(source, neighbors, size, numThreads); - return neighbors; - } finally { - if (hostCopy != null) { - hostCopy.close(); - } - } - } - - /** - * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code - * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code - * source} must be a host matrix (stateless {@code getRow}). - */ - private static void fillNeighborArrayParallel( - CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) throws IOException { - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - int perThread = (size + numThreads - 1) / numThreads; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int start = t * perThread; - final int end = Math.min(start + perThread, size); - if (start >= end) { - break; - } - tasks.add( - () -> { - fillNeighborRange(source, neighbors, start, end); - return null; - }); - } - new TaskExecutor(pool).invokeAll(tasks); - } finally { - pool.shutdown(); - } - } - - /** Fills {@code neighbors[start, end)} from the adjacency rows. */ - private static void fillNeighborRange( - CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { - for (int i = start; i < end; i++) { - RowView rv = source.getRow(i); + for (int i = 0; i < size; i++) { + RowView rv = adjacency.getRow(i); if (rv != null && rv.size() > 0) { - NeighborArray na = new NeighborArray((int) rv.size(), true); + neighbors[i] = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { - na.addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); } - neighbors[i] = na; } else { neighbors[i] = new NeighborArray(0, true); } } + return neighbors; } /** 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 index f95487e953..e9f4d6fead 100644 --- 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 @@ -53,25 +53,7 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(acceleratedHNSWParams, 0); - } - - /** - * Constructor for {@link Lucene101AcceleratedHNSWCodec} used for the native flat-buffered build - * path (see {@link NativeFlatBufferedHNSWVectorsWriter}). Package-private: only {@link - * CagraHnswBulkIndexWriter} can guarantee the invariants (single unmerged segment, no index - * sort, exact vector count) that native flat buffering requires, so this overload is not part - * of the public API. - * - * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} - * @param numInputVectors the exact number of vectors to be indexed, used to pre-size the native - * flat buffer - * @throws Exception exception - */ - Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) - throws Exception { - this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(acceleratedHNSWParams, numInputVectors); + initializeFormat(acceleratedHNSWParams); } /** @@ -79,19 +61,17 @@ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams * with an instance of {@link AcceleratedHNSWParams} with default parameter values. */ private void initializeFormatDefaultValues() { - initializeFormat(new AcceleratedHNSWParams.Builder().build(), 0); + initializeFormat(new AcceleratedHNSWParams.Builder().build()); } /** * Initialize an instance of {@link Lucene99AcceleratedHNSWVectorsFormat}. * * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} to use - * @param numInputVectors the exact number of vectors to be indexed, used to pre-size the native - * flat buffer (0 = disabled, the default heap-buffered path) */ - private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) { + private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { try { - format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams, numInputVectors); + format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams); setKnnFormat(format); } catch (LibraryException ex) { log.log( 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 index 59e012427d..1b1bc3cbd2 100644 --- 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 @@ -31,7 +31,6 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final int numInputVectors; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; @@ -66,25 +65,8 @@ public Lucene99AcceleratedHNSWVectorsFormat() { * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} */ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNSWParams) { - this(acceleratedHNSWParams, 0); - } - - /** - * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} for the native flat-buffered build - * path (see {@link NativeFlatBufferedHNSWVectorsWriter}). Package-private: only {@link - * CagraHnswBulkIndexWriter} (via {@link Lucene101AcceleratedHNSWCodec}) can guarantee the - * invariants that native flat buffering requires, so this overload is not part of the public - * API. - * - * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} - * @param numInputVectors the exact number of vectors to be indexed, used to pre-size the native - * flat buffer (0 = disabled, the default heap-buffered path) - */ - Lucene99AcceleratedHNSWVectorsFormat( - AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) { super("Lucene99AcceleratedHNSWVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; - this.numInputVectors = numInputVectors; } /** @@ -92,23 +74,14 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - boolean nativeMode = isSupported() && numInputVectors > 0; + var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); if (isSupported()) { - if (nativeMode) { - log.log(Level.FINE, "cuVS is supported so using the NativeFlatBufferedHNSWVectorsWriter"); - // In hint mode the accelerated writer owns the flat .vec/.vemf files, so the Lucene flat - // writer must not be created (it would open the same outputs). - return new NativeFlatBufferedHNSWVectorsWriter( - state, acceleratedHNSWParams, numInputVectors); - } log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); } else { log.log( Level.WARNING, "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); try { return LUCENE_PROVIDER.getLuceneHnswVectorsWriterInstance( state, 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 index bc27af0b06..13edc64975 100644 --- 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 @@ -4,31 +4,42 @@ */ 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.CuVSHostMatrix; +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.KnnVectorValues; +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.search.DocIdSetIterator; -import org.apache.lucene.util.Bits; +import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; @@ -43,13 +54,28 @@ 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 AcceleratedHnswGraphOutput graphOutput; + 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} * @@ -64,11 +90,30 @@ public Lucene99AcceleratedHNSWVectorsWriter( FlatVectorsWriter flatVectorsWriter) throws IOException { super(); - this.flatVectorsWriter = Objects.requireNonNull(flatVectorsWriter); + 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 { - graphOutput = new AcceleratedHnswGraphOutput(state, acceleratedHNSWParams); + 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 { @@ -93,6 +138,66 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException 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(), + 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); + cagraIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + /** * Build the indexes and writes it to the disk. */ @@ -115,7 +220,7 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * @throws IOException */ private void writeField(FieldWriter fieldData) throws IOException { - graphOutput.writeField(fieldData.fieldInfo(), fieldData.getFloatVectors()); + writeFieldInternal(fieldData.fieldInfo(), fieldData.getFloatVectors()); } /** @@ -134,68 +239,48 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr for (int i = 0; i < floatVectors.size(); i++) { sortedVectors.add(floatVectors.get(new2OldOrd[i])); } - graphOutput.writeField(fieldData.fieldInfo(), sortedVectors); + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); } /** - * Streams merged vectors directly into a native host-memory matrix (CuVSHostMatrix) - * without materialising a List on the Java heap, then calls graphOutput.writeField. - * This avoids the double-copy OOM (heap list + native matrix simultaneously) that - * occurs when force-merging large segments. + * 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 vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { try { - // FloatVectorValues#size() on the merged view is the raw sum of every source segment's - // on-disk vector count (MergedVectorValues.MergedFloat32VectorValues computes it once at - // construction from each sub-reader's unfiltered size) -- NOT the number of live - // (non-deleted) vectors the iterator below will actually yield, which is what - // CuVSMatrix.hostBuilder needs since it preallocates a fixed-size native buffer. Using - // size() here under-fills that buffer whenever the merge drops deleted docs, leaving the - // graph built over more rows than were actually populated. - // - // size() IS trustworthy when no segment being merged has any deletions: per-segment vector - // counts already exclude docs without a value for this field (sparse fields are handled at - // the single-segment level, independent of deletions), so the raw sum equals the live count - // in that case and the extra counting pass below can be skipped. - boolean anySegmentHasDeletions = false; - for (Bits liveDocs : mergeState.liveDocs) { - if (liveDocs != null) { - anySegmentHasDeletions = true; - break; - } - } - - int size; - if (anySegmentHasDeletions) { - // Count the live vectors via a throwaway iteration first (mergeFloatVectorValues - // constructs a fresh, independent view each call, so this doesn't disturb the real build - // pass below). - size = 0; - FloatVectorValues counting = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - KnnVectorValues.DocIndexIterator countingIt = counting.iterator(); - for (int doc = countingIt.nextDoc(); - doc != DocIdSetIterator.NO_MORE_DOCS; - doc = countingIt.nextDoc()) { - size++; - } - } else { - size = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) - .size(); - } + 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); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } - FloatVectorValues mergedVectors = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - int dims = fieldInfo.getVectorDimension(); - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); - KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - builder.addVector(mergedVectors.vectorValue(it.index())); - } - CuVSHostMatrix dataset = builder.build(); - graphOutput.writeField(fieldInfo, dataset); + /** + * 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); } @@ -220,7 +305,14 @@ public void finish() throws IOException { } finished = true; flatVectorsWriter.finish(); - graphOutput.finish(); + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } } /** @@ -229,7 +321,7 @@ public void finish() throws IOException { @Override public void close() throws IOException { printInfoStream(infoStream, COMPONENT, "Closing resources"); - IOUtils.close(graphOutput, flatVectorsWriter); + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); closeCuVSResourcesInstance(); } 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 index cf6d19f7ca..87907d2cbb 100644 --- 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 @@ -155,7 +155,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; - CuVSMatrix dataset = Utils.createByteMatrix(vectors, bytesPerVector); + CuVSMatrix dataset = + Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance()); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, vectors); @@ -178,18 +179,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, + size, dimensions, adjacencyListMatrix, - dataset, + vectors, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.BINARY, - acceleratedHNSWParams.getWriterThreads()); + QuantizationType.BINARY); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -277,8 +277,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; 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 index aebfbea7c6..7141af56ee 100644 --- 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 @@ -181,7 +181,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE } // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = Utils.createByteMatrix(unsignedVectors, dimensions); + CuVSMatrix dataset = + Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, unsignedVectors); @@ -203,19 +204,18 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, + size, dimensions, adjacencyListMatrix, - dataset, + unsignedVectors, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.SCALAR, - acceleratedHNSWParams.getWriterThreads()); + QuantizationType.SCALAR); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -302,8 +302,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java deleted file mode 100644 index a246efb69d..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java +++ /dev/null @@ -1,124 +0,0 @@ -/* - * 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.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import java.io.IOException; -import org.apache.lucene.codecs.KnnFieldVectorsWriter; -import org.apache.lucene.index.DocsWithFieldSet; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.util.RamUsageEstimator; - -/** - * Native-flat-buffered field writer, used exclusively by {@link - * NativeFlatBufferedHNSWVectorsWriter}: streams incoming vectors directly into a native host - * matrix (see {@link CuVSMatrix#hostBuilder}) instead of accumulating a {@code List} on - * the Java heap, so the full dataset is never held twice (heap list + native copy). The matrix is - * preallocated for exactly {@code numInputVectors} rows, so that count must equal the number of - * vectors actually added (validated by {@link NativeFlatBufferedHNSWVectorsWriter}). - * - *

Package-private and float32-only — quantization is not supported on this path. Distinct from - * the heap-buffered {@link FieldWriter} used by the other accelerated writers; the two share no - * state or behavior beyond both extending {@link KnnFieldVectorsWriter}. - */ -final class NativeFieldWriter extends KnnFieldVectorsWriter { - - private static final long SHALLOW_SIZE = - RamUsageEstimator.shallowSizeOfInstance(NativeFieldWriter.class); - - private final FieldInfo fieldInfo; - private final int numInputVectors; - private final CuVSMatrix.Builder hostMatrixBuilder; - private final DocsWithFieldSet docsWithField = new DocsWithFieldSet(); - private int lastDocID = -1; - private int count; - private CuVSHostMatrix builtMatrix; - - NativeFieldWriter(FieldInfo fieldInfo, int numInputVectors) { - this.fieldInfo = fieldInfo; - this.numInputVectors = numInputVectors; - // Preallocates one contiguous native region of numInputVectors * dimension * 4 bytes. - this.hostMatrixBuilder = - CuVSMatrix.hostBuilder( - numInputVectors, fieldInfo.getVectorDimension(), CuVSMatrix.DataType.FLOAT); - } - - @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)"); - } - if (count >= numInputVectors) { - throw new IllegalStateException( - "Buffered vectors (" - + (count + 1) - + ") exceed numInputVectors (" - + numInputVectors - + ") for field \"" - + fieldInfo.name - + "\". This usually means more vectors arrived than the numInputVectors hint" - + " promised (e.g. a merge or a later flush cycle reused this config); see" - + " CagraHnswBulkIndexWriter."); - } - // hostMatrixBuilder.addVector validates the dimension and performs the native row copy. - hostMatrixBuilder.addVector((float[]) vectorValue); - docsWithField.add(docID); - count++; - lastDocID = docID; - } - - FieldInfo fieldInfo() { - return fieldInfo; - } - - DocsWithFieldSet getDocsWithFieldSet() { - return docsWithField; - } - - /** - * The native host matrix holding the buffered vectors. Built once and cached; the caller owns - * closing it via {@link #releaseNativeBuffer()} once the CAGRA build has consumed it. - */ - CuVSHostMatrix getHostMatrix() { - if (builtMatrix == null) { - builtMatrix = hostMatrixBuilder.build(); - } - return builtMatrix; - } - - /** Number of vectors buffered so far. */ - int getNativeVectorCount() { - return count; - } - - /** - * Closes the native host matrix. Safe to call multiple times, including on a field whose - * buffer was never fully populated (e.g. after a count-mismatch or other flush failure): {@code - * build()} just returns the pre-allocated matrix regardless of how many rows were written, and - * {@code close()} on that matrix is itself idempotent. - */ - void releaseNativeBuffer() { - getHostMatrix().close(); - builtMatrix = null; - } - - @Override - public Object copyValue(Object vectorValue) { - throw new UnsupportedOperationException(); - } - - @Override - public long ramBytesUsed() { - // The native host matrix is off-heap and intentionally excluded from Lucene's heap RAM - // accounting. - return SHALLOW_SIZE; - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java deleted file mode 100644 index 4f52d882d7..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java +++ /dev/null @@ -1,214 +0,0 @@ -/* - * 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.printInfoStream; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; -import static org.apache.lucene.index.VectorEncoding.FLOAT32; -import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; - -import com.nvidia.cuvs.CuVSHostMatrix; -import java.io.IOException; -import java.util.ArrayList; -import java.util.List; -import org.apache.lucene.codecs.KnnFieldVectorsWriter; -import org.apache.lucene.codecs.KnnVectorsWriter; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.MergeState; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.index.Sorter.DocMap; -import org.apache.lucene.util.IOUtils; -import org.apache.lucene.util.InfoStream; - -/** - * Native-flat-buffered {@link KnnVectorsWriter}, streaming vectors directly into a native host - * matrix (see {@link NativeFieldWriter}) rather than a heap {@code List}, and writing the - * flat {@code .vec}/{@code .vemf} files itself via {@link NativeFlatVectorsWriter} instead of - * delegating to Lucene's {@link org.apache.lucene.codecs.hnsw.FlatVectorsWriter}. - * - *

This writer supports only the unsorted, single-segment flush path: merges and index-sorted - * flushes are rejected. It is package-private and constructed only by {@link - * Lucene99AcceleratedHNSWVectorsFormat#fieldsWriter} when {@code numInputVectors > 0}, which is - * reachable only via the package-private constructors on {@link Lucene99AcceleratedHNSWVectorsFormat} - * and {@link Lucene101AcceleratedHNSWCodec} used exclusively by {@link CagraHnswBulkIndexWriter} — - * the only caller that owns its {@link org.apache.lucene.index.IndexWriter} and can guarantee - * those invariants. - */ -final class NativeFlatBufferedHNSWVectorsWriter extends KnnVectorsWriter { - - private static final long SHALLOW_RAM_BYTES_USED = - shallowSizeOfInstance(NativeFlatBufferedHNSWVectorsWriter.class); - private static final String COMPONENT = "NativeFlatBufferedHNSWVectorsWriter"; - - private final int numInputVectors; - private final List fields = new ArrayList<>(); - private final InfoStream infoStream; - private NativeFlatVectorsWriter nativeFlat; - private AcceleratedHnswGraphOutput graphOutput; - private boolean finished; - - NativeFlatBufferedHNSWVectorsWriter( - SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) - throws IOException { - super(); - if (numInputVectors <= 0) { - throw new IllegalArgumentException("numInputVectors must be > 0, got " + numInputVectors); - } - if (state.segmentInfo.getIndexSort() != null) { - throw new IllegalArgumentException( - "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" - + " index-sorted segments; unset it (0) to use the heap-buffered path"); - } - this.infoStream = state.infoStream; - this.numInputVectors = numInputVectors; - boolean success = false; - try { - // In hint mode we own the flat files; the Lucene flat writer must be absent to avoid opening - // the same .vec/.vemf outputs. - nativeFlat = new NativeFlatVectorsWriter(state); - graphOutput = new AcceleratedHnswGraphOutput(state, acceleratedHNSWParams); - success = true; - printInfoStream(infoStream, COMPONENT, "NativeFlatBufferedHNSWVectorsWriter 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); - } - // Buffer directly into a native host matrix; return the field writer itself so Lucene routes - // addValue() here rather than to a (nonexistent) Lucene flat field writer. - var cuvsFieldWriter = new NativeFieldWriter(fieldInfo, numInputVectors); - fields.add(cuvsFieldWriter); - return cuvsFieldWriter; - } - - /** - * Writes the flat {@code .vec}/{@code .vemf} from each field's native host matrix, builds the - * CAGRA/HNSW graph from the same matrix, then releases the matrix. - */ - @Override - public void flush(int maxDoc, DocMap sortMap) throws IOException { - if (sortMap != null) { - throw new UnsupportedOperationException( - "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" - + " index-sorted segments; unset it (0) to enable the sorted flush path"); - } - for (var field : fields) { - writeFieldNative(field, maxDoc); - } - } - - private void writeFieldNative(NativeFieldWriter fieldData, int maxDoc) throws IOException { - int count = fieldData.getNativeVectorCount(); - if (count != numInputVectors) { - throw new IllegalStateException( - "numInputVectors (" - + numInputVectors - + ") must equal the number of vectors added (" - + count - + ") for field \"" - + fieldData.fieldInfo().name - + "\"; the native host matrix is sized for the hint exactly. This usually means" - + " IndexWriterConfig's auto-flush wasn't disabled (setMaxBufferedDocs /" - + " setRAMBufferSizeMB(DISABLE_AUTO_FLUSH)), so a flush landed before exactly" - + " numInputVectors vectors were added; see CagraHnswBulkIndexWriter."); - } - FieldInfo fieldInfo = fieldData.fieldInfo(); - try { - CuVSHostMatrix dataset = fieldData.getHostMatrix(); - nativeFlat.writeField(fieldInfo, dataset, maxDoc, fieldData.getDocsWithFieldSet()); - graphOutput.writeField(fieldInfo, dataset); - } finally { - fieldData.releaseNativeBuffer(); - } - } - - /** - * Native flat buffering supports only the unsorted single-segment flush path; merges are - * rejected. - */ - @Override - public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { - throw new UnsupportedOperationException( - "AcceleratedHNSWParams.numInputVectors (native flat buffering) supports only the" - + " unsorted single-segment flush path; unset it (0) to enable merges"); - } - - /** - * Called once at the end before close. - */ - @Override - public void finish() throws IOException { - if (finished) { - throw new IllegalStateException("already finished"); - } - finished = true; - nativeFlat.finish(); - graphOutput.finish(); - } - - /** - * Closes the resources. - */ - @Override - public void close() throws IOException { - printInfoStream(infoStream, COMPONENT, "Closing resources"); - try { - releaseAllNativeBuffers(); - } finally { - IOUtils.close(graphOutput, nativeFlat); - closeCuVSResourcesInstance(); - } - } - - /** - * Releases every field's native host matrix. {@link #writeFieldNative} already releases a - * field's buffer once its write succeeds, but that is not a reliable place to guarantee cleanup: - * a count mismatch throws before that field's buffer is ever touched, and any field failing - * aborts {@link #flush}'s loop before later fields are even reached. This is the guaranteed - * backstop, run unconditionally on close regardless of how flush exited. - * {@link NativeFieldWriter#releaseNativeBuffer} is idempotent, so re-releasing an - * already-released field here is a safe no-op. - */ - private void releaseAllNativeBuffers() { - RuntimeException firstFailure = null; - for (var field : fields) { - try { - field.releaseNativeBuffer(); - } catch (RuntimeException e) { - if (firstFailure == null) { - firstFailure = e; - } else { - firstFailure.addSuppressed(e); - } - } - } - if (firstFailure != null) { - throw firstFailure; - } - } - - /** - * 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/NativeFlatVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java deleted file mode 100644 index 4b37686590..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java +++ /dev/null @@ -1,199 +0,0 @@ -/* - * 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.CuVSHostMatrix; -import java.io.Closeable; -import java.io.IOException; -import java.lang.foreign.MemorySegment; -import java.lang.foreign.ValueLayout; -import java.nio.ByteOrder; -import org.apache.lucene.codecs.CodecUtil; -import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; -import org.apache.lucene.index.DocsWithFieldSet; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.IndexFileNames; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.util.IOUtils; - -/** - * Writes the flat vector files ({@code .vec} data + {@code .vemf} meta) directly from a native host - * matrix, byte-for-byte compatible with Lucene's {@code Lucene99FlatVectorsWriter} so the stock - * {@code Lucene99FlatVectorsReader} can read them. - * - *

This is the hint-path counterpart to delegating to Lucene's {@code FlatVectorsWriter}: the - * accelerated writer streams vectors into a {@link CuVSHostMatrix} during indexing (see - * {@link NativeFieldWriter}) and never materialises the full dataset as a {@code List} - * on the Java heap, so the flat file is written here from that native matrix instead. - * - *

Ported code — pinned to the {@code Lucene99} format. The dense float32 layout (format - * constants, header, meta field order, and footer) is transcribed from {@code - * Lucene99FlatVectorsFormat}/{@code Lucene99FlatVectorsWriter} (tag {@code - * releases/lucene/10.2.0}): - * - *

    - *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat} — format constants: - * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsFormat.java - *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter} — write sequence: - * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsWriter.java - *
- * - *

{@code Lucene99} is a frozen, versioned codec name: its constants and byte layout stay fixed - * for the life of the codec, and files written to it stay readable by the stock {@code - * Lucene99FlatVectorsReader} — via {@code lucene-backward-codecs} once a newer default codec ships - * — for the rest of the Lucene 10.x major version. The constants below (codec names, extensions, - * {@code VERSION_CURRENT}) are the fixed handshake the reader string-matches ({@code - * CodecUtil.checkIndexHeader}) and file-extension-matches against, so they stay as transcribed here - * for the life of the {@code Lucene99} codec. - * - *

On a {@code lucene-core} version bump: run {@code TestNativeFlatVectorsWriterRoundTrip} - * to confirm the stock reader on the new classpath still accepts what this class writes; it builds - * a small index and asserts every vector round-trips byte-exact through {@code - * Lucene99FlatVectorsReader}. - * - *

Moving to a newer flat-vector format (e.g. once {@code Lucene99} support is dropped - * ahead of an 11.x major bump): diff the new format's writer against the one linked above and - * re-derive the {@code writeField}/{@code writeMeta} byte sequence (header, meta field order, - * footer) here. This class writes directly from native memory to avoid the per-vector {@code - * FloatVectorValues} indirection Lucene's own writer requires — that's the reason to keep - * hand-porting the format rather than delegating to it. Then update the codec name/extension/version - * constants below and the Lucene version and links in this javadoc. - */ -final class NativeFlatVectorsWriter implements Closeable { - - // Mirrors org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat (10.2.0) so the standard - // Lucene99FlatVectorsReader accepts the header/codec of the files written here. - private static final String META_CODEC_NAME = "Lucene99FlatVectorsFormatMeta"; - private static final String VECTOR_DATA_CODEC_NAME = "Lucene99FlatVectorsFormatData"; - private static final String META_EXTENSION = "vemf"; - private static final String VECTOR_DATA_EXTENSION = "vec"; - private static final int VERSION_CURRENT = 0; - private static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; - - // Little-endian float layout matching Lucene's on-disk .vec byte order. Must be UNALIGNED: the - // destination is a heap byte[]-backed MemorySegment whose max alignment is 1 byte, so a 4-byte - // aligned JAVA_FLOAT layout is rejected with "incompatible with alignment constraints". - private static final ValueLayout.OfFloat LE_FLOAT = - ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN); - - // Byte granularity for a single writeBytes call; bounds the transient encode buffer. - private static final int CHUNK_BYTES = 1 << 18; // 256 KiB - - private final IndexOutput meta; - private final IndexOutput vectorData; - private boolean finished; - - NativeFlatVectorsWriter(SegmentWriteState state) throws IOException { - String metaFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); - String vectorDataFileName = - IndexFileNames.segmentFileName( - state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); - boolean success = false; - try { - meta = state.directory.createOutput(metaFileName, state.context); - vectorData = state.directory.createOutput(vectorDataFileName, state.context); - CodecUtil.writeIndexHeader( - meta, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); - CodecUtil.writeIndexHeader( - vectorData, - VECTOR_DATA_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - success = true; - } finally { - if (success == false) { - IOUtils.closeWhileHandlingException(this); - } - } - } - - /** - * Writes one dense float32 field: the raw vectors to {@code .vec} and the field metadata (plus the - * ordinal-to-doc mapping) to {@code .vemf}. Vectors are read from {@code matrix} in ordinal order, - * which matches the ascending-docID order in which {@code docsWithField} was populated. - * - * @param field the field being written - * @param matrix the native host matrix holding {@code docsWithField.cardinality()} rows of {@code - * field.getVectorDimension()} floats each - * @param maxDoc the segment's maxDoc, used to build the ordinal-to-doc mapping - * @param docsWithField the set of docs that have a value for this field - */ - void writeField( - FieldInfo field, CuVSHostMatrix matrix, int maxDoc, DocsWithFieldSet docsWithField) - throws IOException { - // Mirrors Lucene99FlatVectorsWriter#writeField (see class-level version pin). - int count = docsWithField.cardinality(); - int dim = field.getVectorDimension(); - long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES); - writeFloat32Vectors(matrix, count, dim); - long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset; - writeMeta(field, maxDoc, count, vectorDataOffset, vectorDataLength, docsWithField); - } - - private void writeFloat32Vectors(CuVSHostMatrix matrix, int count, int dim) throws IOException { - int rowBytes = dim * Float.BYTES; - int chunkRows = Math.max(1, CHUNK_BYTES / rowBytes); - byte[] chunk = new byte[chunkRows * rowBytes]; - MemorySegment chunkSeg = MemorySegment.ofArray(chunk); - float[] rowBuf = new float[dim]; - int r = 0; - for (int ord = 0; ord < count; ord++) { - matrix.getRow(ord).toArray(rowBuf); // native -> heap float[] (bulk) - MemorySegment.copy(rowBuf, 0, chunkSeg, LE_FLOAT, (long) r * rowBytes, dim); // -> LE bytes - if (++r == chunkRows) { - vectorData.writeBytes(chunk, r * rowBytes); - r = 0; - } - } - if (r > 0) { - vectorData.writeBytes(chunk, r * rowBytes); - } - } - - private void writeMeta( - FieldInfo field, - int maxDoc, - int count, - long vectorDataOffset, - long vectorDataLength, - DocsWithFieldSet docsWithField) - throws IOException { - // Mirrors Lucene99FlatVectorsWriter#writeMeta (see class-level version pin); field order is - // load-bearing and must match Lucene's reader. - meta.writeInt(field.number); - meta.writeInt(field.getVectorEncoding().ordinal()); - meta.writeInt(field.getVectorSimilarityFunction().ordinal()); - meta.writeVLong(vectorDataOffset); - meta.writeVLong(vectorDataLength); - meta.writeVInt(field.getVectorDimension()); - meta.writeInt(count); - OrdToDocDISIReaderConfiguration.writeStoredMeta( - DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, count, maxDoc, docsWithField); - } - - /** Writes the end-of-fields marker and footers. Mirrors {@code Lucene99FlatVectorsWriter.finish}. */ - void finish() throws IOException { - if (finished) { - throw new IllegalStateException("already finished"); - } - finished = true; - if (meta != null) { - meta.writeInt(-1); - CodecUtil.writeFooter(meta); - } - if (vectorData != null) { - CodecUtil.writeFooter(vectorData); - } - } - - @Override - public void close() throws IOException { - IOUtils.close(meta, vectorData); - } -} 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 index 9b8028543c..034f8aa5d8 100644 --- 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 @@ -51,50 +51,82 @@ static RuntimeException handleThrowable(Throwable t) throws IOException { } /** - * Builds a host-memory CuVSMatrix from a list of float vectors. + * A method to build a CuVSMatrix from a list of float vectors. * - *

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, - * without creating an intermediate {@code float[][]} on the heap. + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. * * @param data The float vectors - * @param dimensions The number of float elements in each vector - * @return a host-memory CuVSMatrix + * @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) { + 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.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + 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(); } /** - * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). + * 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 - * @return a host-memory CuVSMatrix with BYTE data type + * @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) { + 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.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); + 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(); } /** - * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). + * 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 - * @return a host-memory CuVSMatrix with BYTE data type + * @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) { + static CuVSMatrix createByteMatrixFromArray( + byte[][] data, int bytesPerVector, CuVSResources resources) { CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); + 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); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java deleted file mode 100644 index eeab85cf1b..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java +++ /dev/null @@ -1,35 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import java.io.Closeable; -import java.io.IOException; - -/** - * A forward-only, single-consumer source of float vectors for {@link CagraHnswBulkIndexWriter}. - * - *

Implementations own how the underlying data is fetched (a local file, a database cursor, an - * object-store stream, ...); {@link CagraHnswBulkIndexWriter} only depends on this contract, not on - * any particular storage. - * - *

{@link #get(int, float[])} must be called with strictly increasing indices from a single - * thread -- repeating the previous call's index is not permitted. This mirrors how a bulk indexer - * consumes its input (front-to-back, one pass, exactly once per vector) and lets implementations - * use a simple prefetch/streaming strategy instead of arbitrary random access. - */ -public interface VectorSource extends Closeable { - - /** Number of dimensions of every vector returned by this source. */ - int dimensions(); - - /** Number of vectors available, i.e. the exclusive upper bound on indices passed to {@link #get}. */ - int size(); - - /** - * Fills {@code dst} with the vector at {@code index} (no allocation). {@code index} must be - * strictly greater than the index passed to the previous call, if any. - */ - void get(int index, float[] dst) throws IOException; -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java deleted file mode 100644 index ac027e2701..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static org.junit.Assert.assertFalse; - -import java.lang.reflect.Constructor; -import java.lang.reflect.Modifier; -import org.junit.Test; - -/** - * Regression guard for the native-flat-buffering surface narrowing: the {@code numInputVectors} - * constructor overloads on {@link Lucene101AcceleratedHNSWCodec} and {@link - * Lucene99AcceleratedHNSWVectorsFormat} must stay package-private, reachable only from within - * {@code com.nvidia.cuvs.lucene} (i.e. only by {@link CagraHnswBulkIndexWriter}), not from a - * generic external Lucene codec user. A future change that accidentally re-widens either - * constructor to {@code public} would reopen exactly the footgun this package's design is meant - * to close, without necessarily being caught by any functional test -- this test exists to catch - * that specific mistake directly. - */ -public class TestAcceleratedHNSWParamsSurface { - - @Test - public void testCodecNumInputVectorsConstructorIsNotPublic() throws NoSuchMethodException { - Constructor ctor = - Lucene101AcceleratedHNSWCodec.class.getDeclaredConstructor( - AcceleratedHNSWParams.class, int.class); - assertFalse( - "Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams, int) must not be public; it should" - + " only be reachable from within com.nvidia.cuvs.lucene (CagraHnswBulkIndexWriter)", - Modifier.isPublic(ctor.getModifiers())); - } - - @Test - public void testFormatNumInputVectorsConstructorIsNotPublic() throws NoSuchMethodException { - Constructor ctor = - Lucene99AcceleratedHNSWVectorsFormat.class.getDeclaredConstructor( - AcceleratedHNSWParams.class, int.class); - assertFalse( - "Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams, int) must not be public; it" - + " should only be reachable from within com.nvidia.cuvs.lucene" - + " (CagraHnswBulkIndexWriter)", - Modifier.isPublic(ctor.getModifiers())); - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java deleted file mode 100644 index 2a299a6440..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java +++ /dev/null @@ -1,446 +0,0 @@ -/* - * 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 com.nvidia.cuvs.spi.CuVSProvider; -import java.io.File; -import java.io.IOException; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -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.IndexWriterConfig; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.SortField; -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.BeforeClass; -import org.junit.Test; - -/** - * Functional coverage for {@link CagraHnswBulkIndexWriter}: the manual/direct-instance {@link - * CagraHnswBulkIndexWriter#addDocument} API (including its safety checks) and the one-shot {@link - * CagraHnswBulkIndexWriter#indexFbin}/{@link CagraHnswBulkIndexWriter#build(VectorSource, Config)} - * convenience entry points (single-segment, K-segment sequential, K-segment overlapped, - * {@link CagraHnswBulkIndexWriter.FieldCallback} metadata, rejection of {@code overlapped} for - * {@code build}). Guard-rail behavior of the underlying native-buffered writer itself is already - * covered by {@link TestNativeFlatBufferingGuardRails}; this class covers this class's own - * orchestration and its own safety checks layered on top. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestCagraHnswBulkIndexWriter extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - private static final String CATEGORY_FIELD = "category"; - - private Random random; - private Path indexDirPath; - private Path fbinPath; - - @BeforeClass - public static void beforeClass() { - // It is recommended to enable RMM allocation mode at application start, before constructing - // any CagraHnswBulkIndexWriter, to avoid device-wide sync from the default allocator. - try { - CuVSProvider.provider().enableRMMAsyncMemory(); - } catch (UnsupportedOperationException unsupported) { - assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); - } - } - - @Before - public void beforeTest() { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); - } - - @After - public void afterTest() throws IOException { - if (indexDirPath != null) { - File dir = indexDirPath.toFile(); - if (dir.exists() && dir.isDirectory()) { - FileUtils.deleteDirectory(dir); - } - } - if (fbinPath != null) { - new File(fbinPath.toString()).delete(); - } - } - - @Test - public void testSingleSegmentBuildIsSearchable() throws Exception { - int numDocs = 300; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, 1, false)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); - } - - @Test - public void testPartitionedSequentialBuildProducesKSegments() throws Exception { - int numDocs = 400; - int dimension = 24; - int k = 4; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, false)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ k); - } - - @Test - public void testOverlappedBuildProducesKSegments() throws Exception { - int numDocs = 400; - int dimension = 24; - int k = 4; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, true)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ k); - } - - @Test - public void testGenericBuildViaVectorSource() throws Exception { - int numDocs = 200; - int dimension = 16; - float[][] dataset = generateDataset(random, numDocs, dimension); - - CagraHnswBulkIndexWriter.build( - new InMemoryVectorSource(dataset), configFor(dimension, 1, false)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); - } - - @Test - public void testGenericBuildRejectsOverlapped() throws Exception { - int dimension = 8; - float[][] dataset = generateDataset(random, 50, dimension); - CagraHnswBulkIndexWriter.Config config = configFor(dimension, 2, true); - try { - CagraHnswBulkIndexWriter.build(new InMemoryVectorSource(dataset), config); - fail( - "expected IllegalArgumentException: overlapped is not supported by build(VectorSource," - + " Config)"); - } catch (IllegalArgumentException expected) { - // expected - } - } - - @Test - public void testIdFieldCanBeDisabled() throws Exception { - int numDocs = 50; - int dimension = 8; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .idField(null) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .segments(1, false) - .targetDirectory(indexDirPath) - .build(); - CagraHnswBulkIndexWriter.indexFbin(fbinPath, config); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); - assertEquals(5, results.scoreDocs.length); - assertNull( - "id field should not be stored when idField(null) is used", - searcher.storedFields().document(results.scoreDocs[0].doc).get(ID_FIELD)); - } - } - - /** {@link CagraHnswBulkIndexWriter.FieldCallback} lets the one-shot API attach metadata per row. */ - @Test - public void testFieldCallbackAttachesMetadata() throws Exception { - int numDocs = 60; - int dimension = 8; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin( - fbinPath, - configFor(dimension, 1, false), - (doc, id) -> - doc.add( - new StringField(CATEGORY_FIELD, id % 2 == 0 ? "even" : "odd", Field.Store.YES))); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 1), 1); - String category = - searcher.storedFields().document(results.scoreDocs[0].doc).get(CATEGORY_FIELD); - assertEquals("even", category); // id=0 is even - } - } - - /** - * The manual/direct-instance API: caller builds the {@link Document} (including arbitrary extra - * fields) and drives {@link CagraHnswBulkIndexWriter#addDocument} directly, same shape as a - * plain {@link org.apache.lucene.index.IndexWriter}. - */ - @Test - public void testManualAddDocumentWithMetadata() throws Exception { - int numDocs = 40; - int dimension = 12; - float[][] dataset = generateDataset(random, numDocs, dimension); - - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - try (Directory dir = FSDirectory.open(indexDirPath); - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, numDocs)) { - for (int i = 0; i < numDocs; i++) { - Document doc = new Document(); - doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - doc.add(new StringField(CATEGORY_FIELD, i % 2 == 0 ? "even" : "odd", Field.Store.YES)); - doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(doc); - } - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals(1, reader.leaves().size()); - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 1), 1); - assertEquals( - "even", searcher.storedFields().document(results.scoreDocs[0].doc).get(CATEGORY_FIELD)); - } - } - - @Test - public void testCloseWithTooFewDocumentsThrows() throws Exception { - int dimension = 8; - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - try (Directory dir = FSDirectory.open(indexDirPath)) { - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, 10); - Document doc = new Document(); - doc.add( - new KnnFloatVectorField( - VECTOR_FIELD, generateDataset(random, 1, dimension)[0], EUCLIDEAN)); - writer.addDocument(doc); // only 1 of the promised 10 - - try { - writer.close(); - fail("expected IllegalStateException: fewer documents added than exactVectorCount"); - } catch (IllegalStateException expected) { - // Specifically this class's mismatch error, not the underlying native writer's own count - // check: close() must roll the buffered documents back, not flush them. - assertEquals("expected 10 documents, got 1", expected.getMessage()); - } - // The rollback left nothing committed, so there is no segment to open. - assertFalse(DirectoryReader.indexExists(dir)); - } - } - - @Test - public void testSourceFailureMidBuildIsNotMaskedByCleanup() throws Exception { - int dimension = 8; - float[][] dataset = generateDataset(random, 200, dimension); - // Fails halfway through, so the writer is cleaned up holding 100 of the 200 promised vectors. - VectorSource failing = new FailingVectorSource(dataset, 100); - - try { - CagraHnswBulkIndexWriter.build(failing, configFor(dimension, 1, false)); - fail("expected the source's IOException to propagate"); - } catch (IOException expected) { - // The cleanup path must not report the count mismatch it necessarily sees instead of the - // failure that caused it. - assertEquals("source failed at 100", expected.getMessage()); - } - } - - @Test - public void testAddDocumentBeyondExactCountThrows() throws Exception { - int dimension = 8; - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - try (Directory dir = FSDirectory.open(indexDirPath)) { - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, 1); - float[][] vectors = generateDataset(random, 2, dimension); - Document doc1 = new Document(); - doc1.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[0], EUCLIDEAN)); - writer.addDocument(doc1); - - Document doc2 = new Document(); - doc2.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[1], EUCLIDEAN)); - try { - writer.addDocument(doc2); // exactVectorCount was 1 - fail("expected IllegalStateException: more documents added than exactVectorCount"); - } catch (IllegalStateException expected) { - // expected -- the rejected call never reached the underlying writer or incremented the - // count, so it's still exactly 1 and close() below completes normally. - } finally { - writer.close(); - } - } - } - - @Test - public void testConstructorRejectsIndexSort() throws Exception { - int dimension = 8; - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - IndexWriterConfig conf = - new IndexWriterConfig() - .setIndexSort(new Sort(new SortField(ID_FIELD, SortField.Type.STRING))); - - try (Directory dir = FSDirectory.open(indexDirPath)) { - try { - new CagraHnswBulkIndexWriter(dir, conf, config, 10); - fail("expected IllegalArgumentException: index-sorted segments are not supported"); - } catch (IllegalArgumentException expected) { - // expected - } - } - } - - private CagraHnswBulkIndexWriter.Config configFor( - int dimension, int numSegments, boolean overlapped) { - return CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .segments(numSegments, overlapped) - .targetDirectory(indexDirPath) - .build(); - } - - private void assertSearchable(int numDocs, float[][] dataset, int expectedSegments) - throws Exception { - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("unexpected segment count", expectedSegments, reader.leaves().size()); - - IndexSearcher searcher = new IndexSearcher(reader); - int topK = 10; - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], topK), topK); - assertEquals(topK, results.scoreDocs.length); - boolean sawQueryVectorItself = false; - for (var scoreDoc : results.scoreDocs) { - String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); - int idValue = Integer.parseInt(id); - assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); - sawQueryVectorItself |= idValue == 0; - } - // Querying with dataset[0] itself (an exact, zero-distance match) should reliably surface - // id=0 within topK for a graph this small, across all segments the query fans out over. - assertTrue( - "expected id=0 (the query vector itself) within topK results", sawQueryVectorItself); - } - } - - /** An {@link InMemoryVectorSource} that throws once a given row is reached. */ - private static final class FailingVectorSource implements VectorSource { - private final InMemoryVectorSource delegate; - private final int failAt; - - FailingVectorSource(float[][] dataset, int failAt) { - this.delegate = new InMemoryVectorSource(dataset); - this.failAt = failAt; - } - - @Override - public int dimensions() { - return delegate.dimensions(); - } - - @Override - public int size() { - return delegate.size(); - } - - @Override - public void get(int index, float[] dst) throws IOException { - if (index >= failAt) { - throw new IOException("source failed at " + failAt); - } - delegate.get(index, dst); - } - - @Override - public void close() { - delegate.close(); - } - } - - private static final class InMemoryVectorSource implements VectorSource { - private final float[][] dataset; - - InMemoryVectorSource(float[][] dataset) { - this.dataset = dataset; - } - - @Override - public int dimensions() { - return dataset.length == 0 ? 0 : dataset[0].length; - } - - @Override - public int size() { - return dataset.length; - } - - @Override - public void get(int index, float[] dst) { - System.arraycopy(dataset[index], 0, dst, 0, dst.length); - } - - @Override - public void close() { - // nothing to release - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java index 302629d6c1..f559ea3808 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java @@ -195,43 +195,6 @@ public void testHnswHeuristicDelegatesToCuVS() { assertEquals(7, cagraParams.getNumWriterThreads()); } - /** - * {@code cagraGraphBuildAlgo} is only consulted under {@code CUSTOM}; an explicit override left - * on the builder while the strategy is {@code HEURISTIC} must not change the delegated-to-cuVS - * result. Requires the native cuVS library. - */ - @Test - public void testHnswHeuristicIgnoresCagraGraphBuildAlgoOverride() { - assumeTrue("cuVS not supported", isSupported()); - - AcceleratedHNSWParams withoutOverride = - new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) - .withMaxConn(16) - .withBeamWidth(100) - .build(); - AcceleratedHNSWParams withOverride = - new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withMaxConn(16) - .withBeamWidth(100) - .build(); - - CagraIndexParams derivedWithoutOverride = - CagraIndexParamsFactory.create(withoutOverride, 10_000, 128); - CagraIndexParams derivedWithOverride = - CagraIndexParamsFactory.create(withOverride, 10_000, 128); - - assertEquals( - derivedWithoutOverride.getCagraGraphBuildAlgo(), - derivedWithOverride.getCagraGraphBuildAlgo()); - assertEquals(derivedWithoutOverride.getGraphDegree(), derivedWithOverride.getGraphDegree()); - assertEquals( - derivedWithoutOverride.getIntermediateGraphDegree(), - derivedWithOverride.getIntermediateGraphDegree()); - } - /** * The heuristic type must reach cuVS rather than being pinned to the default: under * SIMILAR_SEARCH_PERFORMANCE cuVS derives {@code graph_degree = 2 + maxConn * 2 / 3} instead of @@ -291,12 +254,10 @@ public void testStrategySpecificParamsRemainAccepted() { .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) .withGraphDegree(96) .withIntermediateGraphDegree(192) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) .build(); // Retained verbatim on the instance; CagraIndexParamsFactory is what declines to apply them. assertEquals(96, hnswParams.getGraphdegree()); assertEquals(192, hnswParams.getIntermediateGraphDegree()); - assertEquals(CagraGraphBuildAlgo.NN_DESCENT, hnswParams.getCagraGraphBuildAlgo()); GPUSearchParams gpuParams = new GPUSearchParams.Builder() 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 index 2b960bf3cd..6db380a3ce 100644 --- 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 @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -14,7 +15,6 @@ 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; @@ -44,12 +44,12 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); w.addDocument(doc2); w.flush(); @@ -69,11 +69,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // 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); + assertVectorsKeepTheirDocuments(r, "f", f); } } } @@ -85,12 +81,12 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); w.addDocument(doc2); @@ -98,17 +94,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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); + assertVectorsKeepTheirDocuments(r, "f1", f1); + assertVectorsKeepTheirDocuments(r, "f2", f2); // opportunistically check boundary condition - search with a 0 topK var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java deleted file mode 100644 index 41e7d4a081..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java +++ /dev/null @@ -1,120 +0,0 @@ -/* - * 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.TestUtils.writeFbin; - -import java.io.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; - -/** - * Correctness coverage for {@link FbinVectorSource}: header parsing, whole-file vs. windowed - * (sliced) reads, and the forward-only/single-consumer contract it shares with {@link - * VectorSource}. Does not require GPU support -- this is pure file I/O. - */ -public class TestFbinVectorSource extends LuceneTestCase { - - private Path fbinPath; - - @Before - public void beforeTest() { - fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); - } - - @After - public void afterTest() { - if (fbinPath != null) { - new File(fbinPath.toString()).delete(); - } - } - - @Test - public void testWholeFileReadMatchesDataset() throws Exception { - Random random = new Random(1); - int numVectors = 250; - int dimension = 17; - float[][] dataset = generateDataset(random, numVectors, dimension); - writeFbin(fbinPath, dataset); - - // A small chunk size forces multiple prefetch chunks so this also exercises the chunk - // boundary/advance() path, not just a single-chunk read. - try (FbinVectorSource source = new FbinVectorSource(fbinPath, /* chunkSizeMB= */ 1)) { - assertEquals(dimension, source.dimensions()); - assertEquals(numVectors, source.size()); - - float[] scratch = new float[dimension]; - for (int i = 0; i < numVectors; i++) { - source.get(i, scratch); - assertArrayEquals("vector " + i + " mismatch", dataset[i], scratch, 0f); - } - } - } - - @Test - public void testWindowedReadServesOnlyItsSlice() throws Exception { - Random random = new Random(2); - int numVectors = 100; - int dimension = 8; - int sliceStart = 30; - int sliceSize = 25; - float[][] dataset = generateDataset(random, numVectors, dimension); - writeFbin(fbinPath, dataset); - - try (FbinVectorSource source = new FbinVectorSource(fbinPath, sliceStart, sliceSize, 1)) { - assertEquals(dimension, source.dimensions()); - assertEquals(sliceSize, source.size()); - - for (int i = 0; i < sliceSize; i++) { - float[] got = source.get(i); - assertArrayEquals( - "relative index " + i + " should map to absolute " + (sliceStart + i), - dataset[sliceStart + i], - got, - 0f); - } - } - } - - @Test - public void testOutOfOrderAccessIsRejected() throws Exception { - Random random = new Random(3); - float[][] dataset = generateDataset(random, 20, 4); - writeFbin(fbinPath, dataset); - - try (FbinVectorSource source = new FbinVectorSource(fbinPath, 1)) { - source.get(5); - try { - source.get(2); // going backwards must be rejected: forward-only contract - fail("expected UnsupportedOperationException for out-of-order access"); - } catch (UnsupportedOperationException expected) { - // expected - } - } - } - - @Test - public void testOutOfBoundsIndexIsRejected() throws Exception { - Random random = new Random(4); - float[][] dataset = generateDataset(random, 10, 4); - writeFbin(fbinPath, dataset); - - try (FbinVectorSource source = new FbinVectorSource(fbinPath, 1)) { - try { - source.get(10); // window is [0, 10) - fail("expected IndexOutOfBoundsException"); - } catch (IndexOutOfBoundsException expected) { - // expected - } - } - } -} 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 index 9bb140228b..707ee35709 100644 --- 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 @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -14,7 +15,6 @@ 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; @@ -44,12 +44,12 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); w.addDocument(doc2); w.flush(); @@ -69,11 +69,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // 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); + assertVectorsKeepTheirDocuments(r, "f", f); } } } @@ -85,12 +81,12 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); w.addDocument(doc2); @@ -98,17 +94,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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); + assertVectorsKeepTheirDocuments(r, "f1", f1); + assertVectorsKeepTheirDocuments(r, "f2", f2); // opportunistically check boundary condition - search with a 0 topK var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java deleted file mode 100644 index 12b00c6e12..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java +++ /dev/null @@ -1,198 +0,0 @@ -/* - * 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 static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import java.util.Random; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.hnsw.HnswGraphProvider; -import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; -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.CodecReader; -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.NoMergePolicy; -import org.apache.lucene.index.Term; -import org.apache.lucene.index.TieredMergePolicy; -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.tests.util.TestUtil; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.junit.Test; - -/** - * Repro for the CI-observed {@code EOFException} in {@code OffHeapFloatVectorValues} during - * concurrent KNN search over an accelerated-HNSW index built via {@link - * TestAcceleratedHNSWDeletedDocuments}/{@link TestCuVSAcceleratedHNSWDeletedDocuments} (both: - * deletions + a real merge, heap-buffered path, no {@code numInputVectors}). - * - *

Rather than relying on a random concurrent search happening to traverse a bad graph node - * (which only reproduced intermittently, on one CI node), this walks the entire merged - * HNSW graph directly and asserts every neighbor ordinal is within the merged segment's actual - * flat-vector count. This targets the suspected root cause: {@code - * Lucene99AcceleratedHNSWVectorsWriter#mergeOneField} derives the merged vector set twice, - * independently -- once via the real {@code flatVectorsWriter.mergeOneField} (the authoritative - * flat {@code .vec} file) and again via {@code vectorBasedMerge}'s own call to {@code - * KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues} to build the CAGRA/HNSW graph. If - * those two independently-derived views of "the merged, post-deletion vector set" ever disagree in - * count or ordinal order, the graph ends up referencing ordinals the flat file doesn't actually - * have, which is exactly what an out-of-bounds read (EOFException) during traversal would look - * like. - * - *

This test is deterministic: it fails on any disagreement between the graph and the flat file, - * rather than depending on a search happening to reach the bad node. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestMergedGraphOrdinalBounds extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String FIELD = "vector"; - - @Test - public void testMergedGraphOrdinalsStayWithinFlatVectorBounds() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - - Random random = new Random(1234); - int segmentSize = 300; - int dimension = 32; - // Interspersed deletions on both segments, so the merge must drop a scattered subset of - // ordinals from each -- not just a contiguous prefix/suffix -- when it re-derives the merged - // vector set. - int deleteEveryNth = 4; - - Codec codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - IndexWriterConfig config = - new IndexWriterConfig().setCodec(codec).setMergePolicy(NoMergePolicy.INSTANCE); - - int expectedLiveVectors; - try (Directory dir = newDirectory(); - IndexWriter writer = new IndexWriter(dir, config)) { - int deletedFromSegment1 = - addSegmentWithInterspersedDeletions( - writer, 0, segmentSize, dimension, deleteEveryNth, random); - writer.commit(); // segment 1, alone - int deletedFromSegment2 = - addSegmentWithInterspersedDeletions( - writer, segmentSize, segmentSize, dimension, deleteEveryNth, random); - writer.commit(); // segment 2, alone - - expectedLiveVectors = 2 * segmentSize - deletedFromSegment1 - deletedFromSegment2; - - // NoMergePolicy blocks forced merges too, so swap it out now that the two segments (each - // with their own interspersed deletions already committed) are set up. - writer.getConfig().setMergePolicy(new TieredMergePolicy()); - writer.forceMerge(1); - writer.commit(); - - try (DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals( - "expected the forced merge to produce a single segment", 1, reader.leaves().size()); - LeafReader leaf = reader.leaves().get(0).reader(); - - FloatVectorValues flatValues = leaf.getFloatVectorValues(FIELD); - assertEquals( - "merged flat vector count should equal (added - deleted)", - expectedLiveVectors, - flatValues.size()); - - HnswGraph graph = graphOf(leaf); - int level0NodeCount = graph.getNodesOnLevel(0).size(); - assertEquals( - "HNSW graph's level-0 node count disagrees with the merged flat vector file's actual" - + " count -- the graph and the flat file were derived independently by" - + " vectorBasedMerge and flatVectorsWriter.mergeOneField and disagree", - flatValues.size(), - level0NodeCount); - - assertAllNeighborOrdinalsInBounds(graph, flatValues.size()); - } - } - } - - /** - * Every neighbor referenced anywhere in the graph, at every level, must be a valid ordinal into - * the merged segment's actual flat vector data -- otherwise a reader resolving that neighbor's - * vector (e.g. mid-search, to score it) reads past the end of the flat file. - */ - private static void assertAllNeighborOrdinalsInBounds(HnswGraph graph, int liveVectorCount) - throws Exception { - for (int level = 0; level < graph.numLevels(); level++) { - HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(level); - while (nodes.hasNext()) { - int node = nodes.nextInt(); - assertTrue( - "node " - + node - + " at level " - + level - + " is itself out of bounds (live vectors: " - + liveVectorCount - + ")", - node >= 0 && node < liveVectorCount); - graph.seek(level, node); - for (int neighbor = graph.nextNeighbor(); - neighbor != NO_MORE_DOCS; - neighbor = graph.nextNeighbor()) { - assertTrue( - "node " - + node - + " at level " - + level - + " has a neighbor ordinal " - + neighbor - + " out of bounds for the merged segment's " - + liveVectorCount - + " live vectors", - neighbor >= 0 && neighbor < liveVectorCount); - } - } - } - } - - private static HnswGraph graphOf(LeafReader leaf) throws Exception { - KnnVectorsReader knnReader = ((CodecReader) leaf).getVectorReader(); - if (knnReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { - knnReader = fieldsReader.getFieldReader(FIELD); - } - return ((HnswGraphProvider) knnReader).getGraph(FIELD); - } - - /** - * Adds {@code count} documents (global ids {@code [startId, startId + count)}), then deletes - * every {@code deleteEveryNth}-th one by id, scattering the deletions across the segment rather - * than leaving a contiguous surviving range. - * - * @return the number of documents deleted from this segment - */ - private static int addSegmentWithInterspersedDeletions( - IndexWriter writer, int startId, int count, int dimension, int deleteEveryNth, Random random) - throws Exception { - float[][] dataset = generateDataset(random, count, dimension); - for (int i = 0; i < count; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); - document.add(new KnnFloatVectorField(FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - int deleted = 0; - for (int i = 0; i < count; i += deleteEveryNth) { - writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(startId + i))); - deleted++; - } - return deleted; - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java deleted file mode 100644 index e116e9f02a..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java +++ /dev/null @@ -1,269 +0,0 @@ -/* - * 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.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -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.NumericDocValuesField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.SerialMergeScheduler; -import org.apache.lucene.index.TieredMergePolicy; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.SortField; -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; - -/** - * Negative-path coverage for the guard rails in {@code NativeFlatBufferedHNSWVectorsWriter} around - * {@code numInputVectors} (native flat buffering): a count mismatch, an index-sorted segment, a - * merge attempt, and that a count mismatch on one field doesn't leave another field's native - * buffer unreleased on close. - * - *

Positive-path coverage (does a natively-buffered index actually search correctly, tolerate - * deletions, etc.) lives separately in {@link TestNativeFlatBufferingIndexAndSearch}. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestNativeFlatBufferingGuardRails extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - private Random random; - private Path indexDirPath; - - @Before - public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @After - public void afterTest() throws Exception { - if (indexDirPath == null) { - return; - } - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } - - /** - * The count-mismatch guard exists for a caller-side bookkeeping error: declaring {@code - * numInputVectors} against the pre-filter document count instead of the number of vectors that - * actually reach {@code addValue} (e.g. an ingest-time filter skips some documents' vector - * field). It is unrelated to, and not triggered by, Lucene-level deletion -- see the javadoc on - * {@link TestNativeFlatBufferingIndexAndSearch#testDeletedDocsAfterNativeFlatBufferedFlush}. - */ - @Test - public void testCountMismatchFromIngestTimeFilterIsRejected() throws Exception { - int declaredNumInputVectors = 100; - int actuallyIndexed = declaredNumInputVectors - 1; // one doc "filtered out" before addValue - int dimension = 32; - float[][] dataset = generateDataset(random, actuallyIndexed, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, declaredNumInputVectors); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(declaredNumInputVectors + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < actuallyIndexed; 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)); - writer.addDocument(document); - } - IllegalStateException thrown = expectThrows(IllegalStateException.class, writer::commit); - assertTrue( - "unexpected message: " + thrown.getMessage(), - thrown.getMessage().contains("numInputVectors")); - } - } - - /** - * When one field's native buffer fails the count-mismatch guard, {@code flush}'s per-field loop - * throws immediately -- a later field, even if it was populated correctly, is never reached by - * {@code writeFieldNative} and so never gets its native buffer released there. Confirms {@code - * NativeFlatBufferedHNSWVectorsWriter#close} still releases every field's buffer as a backstop - * regardless of what flush reached, by asserting close does not throw even though the second - * field's buffer was never touched during the failed flush. - */ - @Test - public void testCloseAfterCountMismatchReleasesEveryFieldsBuffer() throws Exception { - int declaredNumInputVectors = 50; - int dimension = 16; - // fieldA is added to documents first, so it is flush()'s first (and only-attempted) field; - // fieldB is added second and, unlike fieldA, is fully and correctly filled on every document. - String fieldA = "vector_field_a"; - String fieldB = "vector_field_b"; - float[][] datasetA = generateDataset(random, declaredNumInputVectors - 1, dimension); - float[][] datasetB = generateDataset(random, declaredNumInputVectors, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, declaredNumInputVectors); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(declaredNumInputVectors + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < declaredNumInputVectors; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - // Skip fieldA on the last doc so it ends up one short of declaredNumInputVectors. - if (i < declaredNumInputVectors - 1) { - document.add(new KnnFloatVectorField(fieldA, datasetA[i], EUCLIDEAN)); - } - document.add(new KnnFloatVectorField(fieldB, datasetB[i], EUCLIDEAN)); - writer.addDocument(document); - } - IllegalStateException thrown = expectThrows(IllegalStateException.class, writer::commit); - assertTrue( - "unexpected message: " + thrown.getMessage(), - thrown.getMessage().contains("numInputVectors")); - // The try-with-resources close() below must not throw even though fieldB's fully-populated - // native buffer was never reached by the aborted flush() loop. - } - } - - /** Native flat buffering pre-sizes a single flush's buffer and cannot support sorted flushes. */ - @Test - public void testIndexSortedSegmentIsRejected() throws Exception { - int numDocs = 50; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - Sort indexSort = new Sort(new SortField("sort_key", SortField.Type.LONG)); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setIndexSort(indexSort) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - // KnnVectorsFormat#fieldsWriter (and so the writer's index-sort check in its constructor) - // is invoked on the first addDocument() for the segment, not at commit() -- so the guard - // must be expected around the whole indexing loop, not just the flush. - IllegalArgumentException thrown = - expectThrows( - IllegalArgumentException.class, - () -> { - 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 NumericDocValuesField("sort_key", numDocs - i)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - }); - assertTrue( - "unexpected message: " + thrown.getMessage(), - thrown.getMessage().contains("index-sorted")); - } - } - - /** - * Native flat buffering supports only the unsorted single-segment flush path; merging two - * natively-buffered segments must be rejected rather than silently mis-sizing the native buffer. - */ - @Test - public void testMergeIsRejected() throws Exception { - int segmentSize = 40; - int dimension = 32; - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, segmentSize); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(segmentSize + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - // Keep the two flushes below as separate segments; NoMergePolicy also blocks forced - // merges, so it is swapped out before forceMerge() is called. - .setMergePolicy(NoMergePolicy.INSTANCE) - // Force merges to run synchronously on the calling thread, so the guard's exception - // (or whatever IndexWriter/SegmentMerger wraps it as) surfaces directly from - // forceMerge() instead of on a background merge thread. - .setMergeScheduler(new SerialMergeScheduler()); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - addSegment(writer, 0, segmentSize, dimension); - writer.commit(); // segment 1: exactly segmentSize vectors, matching numInputVectors - addSegment(writer, segmentSize, segmentSize, dimension); - writer.commit(); // segment 2: exactly segmentSize vectors, matching numInputVectors - - writer.getConfig().setMergePolicy(new TieredMergePolicy()); - Throwable thrown = expectThrows(Throwable.class, () -> writer.forceMerge(1)); - assertTrue( - "expected UnsupportedOperationException somewhere in the cause chain of: " + thrown, - causedBy(thrown, UnsupportedOperationException.class)); - } - } - - private void addSegment(IndexWriter writer, int startId, int count, int dimension) - throws Exception { - float[][] dataset = generateDataset(random, count, dimension); - for (int i = 0; i < count; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - } - - private static boolean causedBy(Throwable t, Class type) { - for (Throwable cur = t; cur != null; cur = cur.getCause()) { - if (type.isInstance(cur)) { - return true; - } - for (Throwable suppressed : cur.getSuppressed()) { - if (causedBy(suppressed, type)) { - return true; - } - } - } - return false; - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java deleted file mode 100644 index 0ba476c652..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java +++ /dev/null @@ -1,220 +0,0 @@ -/* - * 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 com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import java.io.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -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.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.Term; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -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; - -/** - * Positive-path functional coverage for native flat buffering ({@code - * AcceleratedHNSWParams.numInputVectors}) beyond {@code TestNativeFlatVectorsWriterRoundTrip}, - * which only checks that the flat {@code .vec} file round-trips -- not that the resulting index is - * actually searchable, tolerates deletions, or composes correctly with the odd-graph-degree fix. - * - *

The negative/guard-rail paths (count mismatch, index-sorted segments, merges) are covered - * separately in {@link TestNativeFlatBufferingGuardRails}. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestNativeFlatBufferingIndexAndSearch extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - private Random random; - private Path indexDirPath; - - @Before - public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @After - public void afterTest() throws Exception { - if (indexDirPath == null) { - return; - } - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } - - /** A natively-buffered index must still be searchable through the normal Lucene KNN query API. */ - @Test - public void testIndexAndSearch() throws Exception { - int numDocs = 500; - int dimension = 32; - int topK = 10; - float[][] dataset = generateDataset(random, numDocs, dimension); - - buildNativeFlatBufferedIndex(numDocs, dimension, dataset); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - - IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateDataset(random, 1, dimension)[0]; - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); - - assertEquals("expected topK results", topK, results.scoreDocs.length); - for (var scoreDoc : results.scoreDocs) { - String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); - int idValue = Integer.parseInt(id); - assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); - } - } - } - - /** - * Deletion is orthogonal to native flat buffering: {@code IndexWriter.deleteDocuments} only - * updates Lucene's liveDocs bitset at search time -- it never touches {@code NativeFieldWriter} - * or the native host matrix, so it cannot trip (and isn't meant to be caught by) the count-mismatch - * guard rail tested in {@link TestNativeFlatBufferingGuardRails}. This test instead confirms that - * deletions applied after a natively-buffered flush are still honored correctly at search time. - */ - @Test - public void testDeletedDocsAfterNativeFlatBufferedFlush() throws Exception { - int numDocs = 300; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, 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)); - writer.addDocument(document); - } - writer.commit(); // single natively-buffered flush: NativeFieldWriter's count matches numDocs - - // Delete every 3rd doc. No new vectors are added, so this does not trigger another flush of - // the vector field and cannot interact with the numInputVectors hint. - for (int i = 0; i < numDocs; i += 3) { - writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(i))); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals( - "expected the deletions to land in the same single segment", 1, reader.leaves().size()); - assertTrue("expected some deleted docs", reader.numDeletedDocs() > 0); - - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], numDocs), numDocs); - for (var scoreDoc : results.scoreDocs) { - String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); - assertNotEquals( - "deleted doc id=" + id + " was still returned by search", 0, Integer.parseInt(id) % 3); - } - } - } - - /** - * The M = ceil(cagraGraphDegree / 2) fix ({@link TestAcceleratedHNSWOddGraphDegree}) must also - * hold on the native-flat-buffered write path ({@code writeFieldNative}), which is a distinct - * call path from the heap-buffered one that test exercises. - */ - @Test - public void testOddGraphDegreeWithNativeFlatBuffering() throws Exception { - int numDocs = 200; - int dimension = 32; - int oddGraphDegree = 63; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = - new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withIntermediateGraphDegree(128) - .withGraphDegree(oddGraphDegree) - .build(); - - buildNativeFlatBufferedIndex(numDocs, dimension, dataset, params); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); - assertEquals(5, results.scoreDocs.length); - } - } - - private void buildNativeFlatBufferedIndex(int numDocs, int dimension, float[][] dataset) - throws Exception { - buildNativeFlatBufferedIndex( - numDocs, dimension, dataset, new AcceleratedHNSWParams.Builder().build()); - } - - private void buildNativeFlatBufferedIndex( - int numDocs, int dimension, float[][] dataset, AcceleratedHNSWParams params) - throws Exception { - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, 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)); - writer.addDocument(document); - } - writer.commit(); - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java deleted file mode 100644 index 027a086f88..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java +++ /dev/null @@ -1,279 +0,0 @@ -/* - * 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.File; -import java.io.IOException; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.Set; -import java.util.UUID; -import java.util.function.IntFunction; -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.KnnVectorValues; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.search.DocIdSetIterator; -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; - -/** - * Positive round-trip check for {@link NativeFlatVectorsWriter}: builds a single-segment index - * with {@code AcceleratedHNSWParams.numInputVectors} set (native flat buffering) and reads the - * {@code .vec}/{@code .vemf} files back through the stock {@code Lucene99FlatVectorsReader}, - * asserting every vector round-trips byte-exact. - * - *

This is the check called for by {@link NativeFlatVectorsWriter}'s "on a Lucene upgrade" class - * javadoc: it confirms the hand-transcribed format is still readable by Lucene's real reader, and - * should pass on every {@code lucene-core} version bump. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestNativeFlatVectorsWriterRoundTrip extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - private Random random; - private Path indexDirPath; - - @Before - public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @After - public void afterTest() throws Exception { - if (indexDirPath == null) { - return; - } - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } - - @Test - public void vectorsRoundTripThroughStockLucene99FlatVectorsReader() throws Exception { - int numDocs = 500; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - - // Force everything into a single unsorted, unmerged flush: native flat buffering requires - // numInputVectors to equal the exact number of vectors landing in that one flush. - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, 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)); - writer.addDocument(document); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - LeafReader leafReader = reader.leaves().get(0).reader(); - assertFieldRoundTrips(leafReader, VECTOR_FIELD, dimension, numDocs, docId -> dataset[docId]); - } - } - - /** - * Cardinality (8000) exceeds {@code IndexedDISI.MAX_ARRAY_LENGTH} (4095), so {@code - * OrdToDocDISIReaderConfiguration} picks the DENSE (bitset) {@code docsWithField} encoding for - * this field. - */ - @Test - public void vectorsRoundTripWithDenseDocsWithFieldEncoding() throws Exception { - roundTripPartialField(10000, 8000, 32); - } - - /** - * Cardinality (800) stays under {@code IndexedDISI.MAX_ARRAY_LENGTH} (4095), so {@code - * OrdToDocDISIReaderConfiguration} picks the SPARSE (array) {@code docsWithField} encoding - * instead. - */ - @Test - public void vectorsRoundTripWithSparseDocsWithFieldEncoding() throws Exception { - roundTripPartialField(1000, 800, 32); - } - - /** - * Two vector fields written to the same segment, exercising the boundary between consecutive - * per-field records in {@code .vemf}: the second field's header must be found where the first - * field's record actually ends. - */ - @Test - public void vectorsRoundTripAcrossMultipleFieldsInSameSegment() throws Exception { - int numDocs = 500; - int dimensionA = 32; - int dimensionB = 16; - String fieldB = "vector_field_two"; - float[][] datasetA = generateDataset(random, numDocs, dimensionA); - float[][] datasetB = generateDataset(random, numDocs, dimensionB); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, 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, datasetA[i], EUCLIDEAN)); - document.add(new KnnFloatVectorField(fieldB, datasetB[i], EUCLIDEAN)); - writer.addDocument(document); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - LeafReader leafReader = reader.leaves().get(0).reader(); - assertFieldRoundTrips( - leafReader, VECTOR_FIELD, dimensionA, numDocs, docId -> datasetA[docId]); - assertFieldRoundTrips(leafReader, fieldB, dimensionB, numDocs, docId -> datasetB[docId]); - } - } - - /** - * Builds a segment of {@code numDocs} documents where only a random {@code numDocsWithVector} - * of them carry {@code VECTOR_FIELD}, then verifies the round trip. - */ - private void roundTripPartialField(int numDocs, int numDocsWithVector, int dimension) - throws Exception { - Set docsWithVector = randomDocSubset(numDocs, numDocsWithVector); - float[][] dataset = generateDataset(random, numDocsWithVector, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocsWithVector); - - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - Map docIdToVectorIndex = new HashMap<>(); - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - int vectorIndex = 0; - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - if (docsWithVector.contains(i)) { - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[vectorIndex], EUCLIDEAN)); - docIdToVectorIndex.put(i, vectorIndex); - vectorIndex++; - } - writer.addDocument(document); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - LeafReader leafReader = reader.leaves().get(0).reader(); - assertFieldRoundTrips( - leafReader, - VECTOR_FIELD, - dimension, - numDocsWithVector, - docId -> dataset[docIdToVectorIndex.get(docId)]); - } - } - - /** Picks {@code count} distinct doc ids out of {@code [0, numDocs)}. */ - private Set randomDocSubset(int numDocs, int count) { - List allDocs = new ArrayList<>(numDocs); - for (int i = 0; i < numDocs; i++) { - allDocs.add(i); - } - Collections.shuffle(allDocs, random); - return new HashSet<>(allDocs.subList(0, count)); - } - - private void assertFieldRoundTrips( - LeafReader leafReader, - String fieldName, - int dimension, - int expectedCount, - IntFunction expectedVectorForDocId) - throws IOException { - FloatVectorValues values = leafReader.getFloatVectorValues(fieldName); - assertNotNull(values); - assertEquals(expectedCount, values.size()); - assertEquals(dimension, values.dimension()); - - int seen = 0; - KnnVectorValues.DocIndexIterator it = values.iterator(); - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - String id = leafReader.storedFields().document(doc).get(ID_FIELD); - float[] roundTripped = values.vectorValue(it.index()); - assertArrayEquals( - "vector for field=" - + fieldName - + " id=" - + id - + " did not round-trip byte-exact through the stock Lucene99FlatVectorsReader", - expectedVectorForDocId.apply(Integer.parseInt(id)), - roundTripped, - 0f); - seen++; - } - assertEquals("did not visit every vector for field=" + fieldName, expectedCount, seen); - } -} 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 index 5140782fc7..3b05beaacf 100644 --- 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 @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.COSINE; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -77,7 +78,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { 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 StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(F, f[i], EUCLIDEAN)); w.addDocument(doc); w.commit(); @@ -95,12 +96,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { 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); - } + assertVectorsKeepTheirDocuments(r, F, f); } } } @@ -121,7 +117,7 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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 StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(F1, f1[i], EUCLIDEAN)); doc.add(new KnnFloatVectorField(F2, f2[i], EUCLIDEAN)); w.addDocument(doc); @@ -130,19 +126,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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); - } + assertVectorsKeepTheirDocuments(r, F1, f1); + assertVectorsKeepTheirDocuments(r, F2, f2); } } } 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 index e19e5ce64f..22cd509e7b 100644 --- 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 @@ -4,41 +4,63 @@ */ package com.nvidia.cuvs.lucene; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Path; -import java.nio.file.StandardOpenOption; +import java.util.HashSet; import java.util.Random; +import java.util.Set; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.StoredFields; public class TestUtils { - /** Writes {@code dataset} as an uncompressed {@code .fbin} file (little-endian float32 rows). */ - public static void writeFbin(Path path, float[][] dataset) throws IOException { - int numVectors = dataset.length; - int dimension = numVectors == 0 ? 0 : dataset[0].length; - ByteBuffer buf = - ByteBuffer.allocate(8 + numVectors * dimension * Float.BYTES) - .order(ByteOrder.LITTLE_ENDIAN); - buf.putInt(numVectors); - buf.putInt(dimension); - for (float[] vector : dataset) { - for (float v : vector) { - buf.putFloat(v); - } - } - buf.flip(); - try (FileChannel ch = - FileChannel.open( - path, - StandardOpenOption.CREATE, - StandardOpenOption.WRITE, - StandardOpenOption.TRUNCATE_EXISTING)) { - while (buf.hasRemaining()) { - ch.write(buf); - } + /** + * Asserts that every vector in {@code field} is still paired with the document that indexed it, + * matching on the document's stored {@code id} rather than on ordinal position. + * + *

Nothing fixes the order documents land in after a merge, and the randomized test + * framework's {@code MockRandomMergePolicy} disturbs it two independent ways: {@code + * findForcedMerges} shuffles the segments it is about to merge, and {@code + * MockRandomOneMerge.reorder} reverses doc IDs outright. The second applies even when there is + * only one segment, so a test that never commits between documents is no safer than one that + * does. Asserting {@code vectorValue(i) == expectedById[i]} therefore fails on some seeds with + * nothing wrong. What has to hold is that no document loses its own vector, which is what this + * checks. + * + * @param expectedById the vector indexed for each document id, indexed by that id + */ + public static void assertVectorsKeepTheirDocuments( + LeafReader reader, String field, float[][] expectedById) throws IOException { + FloatVectorValues values = reader.getFloatVectorValues(field); + assertNotNull("no vector values for field " + field, values); + assertEquals(expectedById.length, values.size()); + StoredFields storedFields = reader.storedFields(); + Set seen = new HashSet<>(); + int previousDoc = -1; + for (int ord = 0; ord < values.size(); ord++) { + int doc = values.ordToDoc(ord); + // Ordinals are assigned in docID order. Where every document has a vector -- every caller + // today -- ordToDoc is the hardcoded identity and this cannot fire. It earns its keep only + // if a caller passes a field that some documents lack, where the id lookup below would + // otherwise be free to agree with a garbled mapping. + assertTrue("ordToDoc went backwards at ordinal " + ord, doc > previousDoc); + previousDoc = doc; + String storedId = storedFields.document(doc).get("id"); + assertNotNull("document at ordinal " + ord + " has no stored id", storedId); + int id = Integer.parseInt(storedId); + assertTrue("document id " + id + " appeared twice", seen.add(id)); + assertArrayEquals( + "vector for document id " + id + " (doc " + doc + ", ordinal " + ord + ")", + expectedById[id], + values.vectorValue(ord), + 0.0f); } + assertEquals("not every document was found", expectedById.length, seen.size()); } public static float[][] generateDataset(Random random, int size, int dimensions) { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java deleted file mode 100644 index d95c60ae6b..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java +++ /dev/null @@ -1,156 +0,0 @@ -/* - * 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 static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import com.nvidia.cuvs.CuVSMatrix; -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Random; -import org.apache.lucene.store.ByteBuffersDirectory; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; -import org.junit.Test; - -/** - * Verifies that {@code writerThreads > 1} produces the same result as the serial path for the - * two parallelizations gated on {@code AcceleratedHNSWUtils}/{@code GPUBuiltHnswGraph}'s {@code - * PARALLEL_MIN_NODES} threshold: materializing the CAGRA adjacency into {@code NeighborArray}s - * (the {@code GPUBuiltHnswGraph} constructor), and encoding level 0 to disk ({@code - * AcceleratedHNSWUtils#writeGraph}). - * - *

Both tests use a synthetic adjacency ({@link CuVSMatrix#ofArray(int[][])}, the same host-matrix - * construction the higher-layer subset builder already uses) rather than a real CAGRA build, so - * that the comparison isolates these two parallelizations from CAGRA's own build-to-build - * variance -- a real GPU build is not guaranteed to produce the identical graph twice even with the - * same input and thread count, which would make an end-to-end build comparison unreliable for this - * purpose. This needs cuVS/GPU only to allocate the host matrix, not to run a build. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestWriterThreadsGraphEquivalence extends LuceneTestCase { - - // Must be >= PARALLEL_MIN_NODES (1 << 16) in both AcceleratedHNSWUtils and - // GPUBuiltHnswGraph, or the "parallel" runs below silently fall through to the serial branch and - // the test would pass without exercising anything. - private static final int NUM_NODES = (1 << 16) + 1000; - private static final int DEGREE = 12; - private static final int NUM_THREADS = 4; - - @Test - public void fillNeighborArrayParallelMatchesSerial() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); - - try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency)) { - GPUBuiltHnswGraph serial = newSingleLayerGraph(matrix, 1); - GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); - assertGraphsEqual(serial, parallel); - } - } - - @Test - public void writeGraphParallelMatchesSerial() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); - - try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency); - Directory dir = new ByteBuffersDirectory()) { - // Materialize once, serially, so any difference found below is attributable only to - // writeGraph's own parallelization, not to fillNeighborArray's. - GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, 1); - - int[][] serialOffsets; - try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { - serialOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, 1); - } - int[][] parallelOffsets; - try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { - parallelOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, NUM_THREADS); - } - - assertEquals(serialOffsets.length, parallelOffsets.length); - for (int level = 0; level < serialOffsets.length; level++) { - assertArrayEquals( - "per-node byte-length offsets differ for level " + level, - serialOffsets[level], - parallelOffsets[level]); - } - - assertArrayEquals( - "writeGraph's parallel level-0 encoding produced different bytes than the serial path", - readAllBytes(dir, "serial"), - readAllBytes(dir, "parallel")); - } - } - - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) - throws IOException { - // A single layer (layer 0 only): the constructor never consults layerNodes in that case, so - // the placeholder null entry mirrors the convention used elsewhere for "layer 0 needs no node - // list" without actually being read. - return new GPUBuiltHnswGraph( - NUM_NODES, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(layer0Adjacency), - numThreads); - } - - /** Every node/level's in-order arc list must match exactly between the two graphs. */ - private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { - assertEquals(a.numLevels(), b.numLevels()); - for (int level = 0; level < a.numLevels(); level++) { - int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); - for (int node : nodes) { - assertArrayEquals( - "node " + node + " at level " + level + " has different neighbors", - arcsOf(a, level, node), - arcsOf(b, level, node)); - } - } - } - - private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { - graph.seek(level, node); - List arcs = new ArrayList<>(); - for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { - arcs.add(n); - } - return arcs.stream().mapToInt(Integer::intValue).toArray(); - } - - private static byte[] readAllBytes(Directory dir, String name) throws Exception { - try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { - byte[] bytes = new byte[(int) in.length()]; - in.readBytes(bytes, 0, bytes.length); - return bytes; - } - } - - /** - * A deterministic, seeded pseudo-adjacency. It doesn't need to be a real CAGRA graph -- only a - * realistic shape (fixed degree, valid node ids) -- since {@code GPUBuiltHnswGraph} and {@code - * AcceleratedHNSWUtils#writeGraph} don't interpret the neighbor ids semantically. - */ - private static int[][] randomAdjacency(int numNodes, int degree, Random random) { - int[][] adjacency = new int[numNodes][degree]; - for (int[] row : adjacency) { - for (int j = 0; j < degree; j++) { - row[j] = random.nextInt(numNodes); - } - } - return adjacency; - } -} diff --git a/java/panama-bindings/generate-bindings.sh b/java/panama-bindings/generate-bindings.sh index cdefd1ae1c..315b58f89f 100755 --- a/java/panama-bindings/generate-bindings.sh +++ b/java/panama-bindings/generate-bindings.sh @@ -10,7 +10,16 @@ REPODIR=$(cd "$(dirname "$0")"; cd ../../ ; pwd) CURDIR=$(cd "$(dirname "$0")"; pwd) TARGET_PACKAGE="com.nvidia.cuvs.internal.panama" -TARGET_DIR="targets/x86_64-linux/include" +ARCH=$(uname -m) +case "${ARCH}" in + x86_64) CUDA_TARGET_ARCH="x86_64-linux" ;; + aarch64) CUDA_TARGET_ARCH="sbsa-linux" ;; + *) + echo "Unsupported architecture for CUDA include directory lookup: ${ARCH}" + exit 1 + ;; +esac +TARGET_DIR="targets/${CUDA_TARGET_ARCH}/include" if [ -n "${CONDA_PREFIX:-}" ] && [ -d "${CONDA_PREFIX}/${TARGET_DIR}" ]; then CUDA_INCLUDE_DIR="${CONDA_PREFIX}/${TARGET_DIR}" elif [ -d "/usr/local/cuda/${TARGET_DIR}" ]; then @@ -25,7 +34,15 @@ export PATH if [[ $(command -v jextract) == "" ]]; then - JEXTRACT_FILENAME="openjdk-22-jextract+6-47_linux-x64_bin.tar.gz" + case "${ARCH}" in + x86_64) JEXTRACT_ARCH="linux-x64" ;; + aarch64) JEXTRACT_ARCH="linux-aarch64" ;; + *) + echo "Unsupported architecture for jextract download: ${ARCH}" + exit 1 + ;; + esac + JEXTRACT_FILENAME="openjdk-22-jextract+6-47_${JEXTRACT_ARCH}_bin.tar.gz" JEXTRACT_DOWNLOAD_URL="https://download.java.net/java/early_access/jextract/22/6/${JEXTRACT_FILENAME}" echo "jextract doesn't exist. Downloading it from $JEXTRACT_DOWNLOAD_URL."; wget -c $JEXTRACT_DOWNLOAD_URL diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 42b5e8e213..37d4533ef3 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -84,7 +84,6 @@ build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", "cuda-tile", - "cuda-toolkit[tileiras]==13.*", "libkvikio==26.10.*,>=0.0.0a0", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0",