From b7ecf252d0bd00c618bffe4cf559724f321d963f Mon Sep 17 00:00:00 2001 From: Shaunak Kapur Date: Fri, 4 Sep 2026 19:16:27 +0000 Subject: [PATCH] Expose Vamana build and serialize through a public cuvs-java API cuVS provides a GPU Vamana builder and DiskANN-compatible serialization, and the generated Panama bindings for it already exist in cuvs-java, but there is no public API, so nothing can call Vamana from Java. This adds VamanaIndex and VamanaIndexParams beside CagraIndex. Scope is build and serialize only, matching the native surface. cuVS exposes no Vamana search entry point, so neither does this. serialize takes a Path prefix rather than an OutputStream, because one native call writes two files, the graph at the prefix and the dataset at prefix + ".data". VamanaIndexParams.Builder mirrors the native RAFT_EXPECTS checks so an invalid configuration fails in Java with a readable message instead of inside a GPU kernel. The native index may retain a non-owning device view of the dataset, so the index holds a reference to keep it alive and closes it only when it created the matrix from a float[][]. newVamanaIndexBuilder is added to CuVSProvider as a default method rather than an abstract one so that providers written against an earlier version of the interface keep compiling. --- fern/pages/neighbors/vamana.md | 43 +- .../java/com/nvidia/cuvs/VamanaIndex.java | 126 ++++++ .../com/nvidia/cuvs/VamanaIndexParams.java | 369 +++++++++++++++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 12 + .../nvidia/cuvs/spi/UnsupportedProvider.java | 5 + .../cuvs/internal/CuVSParamsHelper.java | 19 + .../nvidia/cuvs/internal/VamanaIndexImpl.java | 292 ++++++++++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 5 + .../cuvs/VamanaBuildAndSerializeIT.java | 426 ++++++++++++++++++ 9 files changed, 1296 insertions(+), 1 deletion(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/VamanaIndexImpl.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/VamanaBuildAndSerializeIT.java diff --git a/fern/pages/neighbors/vamana.md b/fern/pages/neighbors/vamana.md index 6a8a3b25a5..439870981b 100644 --- a/fern/pages/neighbors/vamana.md +++ b/fern/pages/neighbors/vamana.md @@ -10,7 +10,7 @@ Vamana works well when you want to build large DiskANN-compatible graph indexes [C API](/api-reference/c-api-neighbors-vamana) | [C++ API](/api-reference/cpp-api-neighbors-vamana) | [Python API](/api-reference/python-api-neighbors-vamana) | [Rust API](/api-reference/rust-api-cuvs-neighbors-vamana) -Vamana currently supports build and serialize operations in NVIDIA cuVS. Search is performed by loading the serialized index with DiskANN. Java and Go do not currently expose standalone Vamana bindings. +Vamana currently supports build and serialize operations in NVIDIA cuVS. Search is performed by loading the serialized index with DiskANN. Java exposes build and serialize through `VamanaIndex`; Go does not currently expose standalone Vamana bindings. ### Building an index @@ -82,6 +82,30 @@ index_params = vamana.IndexParams( index = vamana.build(index_params, dataset) ``` + + + +```java +import com.nvidia.cuvs.*; + +float[][] dataset = loadData(); + +try (CuVSResources resources = CuVSResources.create()) { + VamanaIndexParams indexParams = new VamanaIndexParams.Builder() + .withGraphDegree(64) + .withVisitedSize(128) + .withQueueSize(255) + .build(); + + try (VamanaIndex index = VamanaIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build()) { + // ... + } +} +``` + @@ -164,6 +188,23 @@ index = vamana.build(vamana.IndexParams(), dataset) vamana.save("/tmp/cuvs-vamana/index", index, include_dataset=True) ``` + + + +```java +import com.nvidia.cuvs.*; +import java.nio.file.Path; + +try (CuVSResources resources = CuVSResources.create(); + VamanaIndex index = VamanaIndex.newBuilder(resources) + .withDataset(loadData()) + .build()) { + + // Writes DiskANN-compatible files using this path prefix. + index.serialize(Path.of("/tmp/cuvs-vamana/index"), true); +} +``` + diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java new file mode 100644 index 0000000000..2deaf24f1e --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import com.nvidia.cuvs.spi.CuVSProvider; +import java.nio.file.Path; +import java.util.Objects; + +/** + * {@link VamanaIndex} encapsulates a Vamana index, along with methods to build + * it on the GPU and serialize it in the DiskANN file format. + *

+ * Vamana is the graph construction algorithm behind DiskANN. cuVS currently + * provides build and serialize only. There is no Vamana search API, so a + * serialized index is searched by loading it with DiskANN. + * + * @since 26.10 + */ +public interface VamanaIndex extends AutoCloseable { + + @Override + void close() throws Exception; + + /** + * Gets the dimensionality of the vectors in this index. + * + * @return the number of dimensions + */ + int getDimensions() throws Throwable; + + /** + * Serializes the index in the DiskANN file format, including the dataset. + *

+ * This writes two files, {@code filePrefix} holding the graph and + * {@code filePrefix + ".data"} holding the dataset. + * + * @param filePrefix the prefix that output file names are derived from + */ + default void serialize(Path filePrefix) throws Throwable { + serialize(filePrefix, true); + } + + /** + * Serializes the index in the DiskANN file format. + *

+ * When {@code includeDataset} is true this writes {@code filePrefix} holding + * the graph and {@code filePrefix + ".data"} holding the dataset. When it is + * false only {@code filePrefix} is written. + *

+ * The argument is a prefix and not a complete file name, matching the native + * {@code file_prefix} parameter. + * + * @param filePrefix the prefix that output file names are derived from + * @param includeDataset whether to write the dataset alongside the graph + */ + void serialize(Path filePrefix, boolean includeDataset) throws Throwable; + + /** + * Gets an instance of {@link CuVSResources} + * + * @return an instance of {@link CuVSResources} + */ + CuVSResources getCuVSResources(); + + /** + * Creates a new Builder with an instance of {@link CuVSResources}. + * + * @param cuvsResources an instance of {@link CuVSResources} + * @throws UnsupportedOperationException if the provider does not support cuvs + */ + static Builder newBuilder(CuVSResources cuvsResources) { + Objects.requireNonNull(cuvsResources); + return CuVSProvider.provider().newVamanaIndexBuilder(cuvsResources); + } + + /** + * Builder helps configure and create an instance of {@link VamanaIndex}. + */ + interface Builder { + + /** + * Sets the dataset for building the {@link VamanaIndex}. + * + * @param vectors a two-dimensional float array + * @return an instance of this Builder + */ + Builder withDataset(float[][] vectors); + + /** + * Sets the dataset for building the {@link VamanaIndex}. + *

+ * The native builder accepts {@code float}, {@code half}, {@code uint8}, + * and {@code int8} datasets. Of those, {@link CuVSMatrix.DataType#FLOAT}, + * {@link CuVSMatrix.DataType#HALF}, and {@link CuVSMatrix.DataType#BYTE} + * are reachable from Java today, where {@code BYTE} is unsigned. + * {@code int8} has no corresponding {@code DataType}. + *

+ * The native index may retain a non-owning device view of the dataset + * rather than copying it, so the caller must keep this matrix open for at + * least as long as the index and close it afterwards. A dataset supplied as + * a {@code float[][]} is created and closed by the index instead. + * + * @param dataset a {@link CuVSMatrix} object containing the vectors + * @return an instance of this Builder + */ + Builder withDataset(CuVSMatrix dataset); + + /** + * Registers an instance of configured {@link VamanaIndexParams} with this + * Builder. + * + * @param vamanaIndexParameters An instance of VamanaIndexParams + * @return An instance of this Builder + */ + Builder withIndexParams(VamanaIndexParams vamanaIndexParameters); + + /** + * Builds and returns an instance of {@link VamanaIndex}. + * + * @return an instance of {@link VamanaIndex} + */ + VamanaIndex build() throws Throwable; + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java new file mode 100644 index 0000000000..d44d443ba3 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java @@ -0,0 +1,369 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import java.util.Arrays; +import java.util.Objects; + +/** + * Supplemental parameters to build a Vamana index. + *

+ * The defaults match the native {@code cuvs::neighbors::vamana::index_params} + * defaults. + * + * @since 26.10 + */ +public class VamanaIndexParams { + + /** + * The graph degrees the native builder supports, matching {@code DEGREE_SIZES} + * in the cuVS Vamana implementation. Sorted, so it can be searched. + */ + private static final int[] SUPPORTED_GRAPH_DEGREES = {32, 64, 128, 256}; + + /** + * Returns the graph degrees the native Vamana builder supports. + * + * @return a copy of the supported graph degrees, in ascending order + */ + public static int[] supportedGraphDegrees() { + return SUPPORTED_GRAPH_DEGREES.clone(); + } + + /** + * Distance metric types supported by the Vamana builder. + *

+ * The native build kernel accepts these two only. Other metrics fail inside + * the kernel rather than at parameter validation time, so they are not + * exposed here. + */ + public enum CuvsDistanceType { + /** + * Squared L2. + */ + L2Expanded(0), + + /** + * Euclidean, the square root of {@link #L2Expanded}. + */ + L2SqrtExpanded(1); + + /** + * The value for the enum choice. + */ + public final int value; + + private CuvsDistanceType(int value) { + this.value = value; + } + } + + private final int graphDegree; + private final int visitedSize; + private final float vamanaIters; + private final float alpha; + private final float maxFraction; + private final float batchBase; + private final int queueSize; + private final int reverseBatchSize; + private final CuvsDistanceType metric; + + private VamanaIndexParams( + int graphDegree, + int visitedSize, + float vamanaIters, + float alpha, + float maxFraction, + float batchBase, + int queueSize, + int reverseBatchSize, + CuvsDistanceType metric) { + this.graphDegree = graphDegree; + this.visitedSize = visitedSize; + this.vamanaIters = vamanaIters; + this.alpha = alpha; + this.maxFraction = maxFraction; + this.batchBase = batchBase; + this.queueSize = queueSize; + this.reverseBatchSize = reverseBatchSize; + this.metric = metric; + } + + /** + * Gets the maximum degree of the output graph, the R parameter in the Vamana + * literature. + */ + public int getGraphDegree() { + return graphDegree; + } + + /** + * Gets the maximum number of visited nodes per search, the L parameter in the + * Vamana literature. + */ + public int getVisitedSize() { + return visitedSize; + } + + /** + * Gets the number of Vamana vector insertion iterations. + */ + public float getVamanaIters() { + return vamanaIters; + } + + /** + * Gets the alpha pruning parameter. + */ + public float getAlpha() { + return alpha; + } + + /** + * Gets the maximum fraction of the dataset inserted per batch. + */ + public float getMaxFraction() { + return maxFraction; + } + + /** + * Gets the growth rate base for batch sizes. + */ + public float getBatchBase() { + return batchBase; + } + + /** + * Gets the candidate queue size. + */ + public int getQueueSize() { + return queueSize; + } + + /** + * Gets the maximum batch size of reverse edge processing. + */ + public int getReverseBatchSize() { + return reverseBatchSize; + } + + /** + * Gets the distance metric. + */ + public CuvsDistanceType getMetric() { + return metric; + } + + @Override + public String toString() { + return "VamanaIndexParams [graphDegree=" + + graphDegree + + ", visitedSize=" + + visitedSize + + ", vamanaIters=" + + vamanaIters + + ", alpha=" + + alpha + + ", maxFraction=" + + maxFraction + + ", batchBase=" + + batchBase + + ", queueSize=" + + queueSize + + ", reverseBatchSize=" + + reverseBatchSize + + ", metric=" + + metric + + "]"; + } + + /** + * Builder configures and creates an instance of {@link VamanaIndexParams}. + */ + public static class Builder { + + private int graphDegree = 32; + private int visitedSize = 64; + private float vamanaIters = 1.0f; + private float alpha = 1.2f; + private float maxFraction = 0.06f; + private float batchBase = 2.0f; + private int queueSize = 127; + private int reverseBatchSize = 1000000; + private CuvsDistanceType metric = CuvsDistanceType.L2Expanded; + + public Builder() {} + + /** + * Sets the maximum degree of the output graph. + * + * @param graphDegree the graph degree, one of + * {@link VamanaIndexParams#supportedGraphDegrees()} + * @return an instance of this Builder + */ + public Builder withGraphDegree(int graphDegree) { + this.graphDegree = graphDegree; + return this; + } + + /** + * Sets the maximum number of visited nodes per search. + *

+ * The native builder requires this to be greater than the graph degree. + * + * @param visitedSize the visited size + * @return an instance of this Builder + */ + public Builder withVisitedSize(int visitedSize) { + this.visitedSize = visitedSize; + return this; + } + + /** + * Sets the number of Vamana vector insertion iterations. + * + * @param vamanaIters the iteration count + * @return an instance of this Builder + */ + public Builder withVamanaIters(float vamanaIters) { + this.vamanaIters = vamanaIters; + return this; + } + + /** + * Sets the alpha pruning parameter. + * + * @param alpha the alpha value + * @return an instance of this Builder + */ + public Builder withAlpha(float alpha) { + this.alpha = alpha; + return this; + } + + /** + * Sets the maximum fraction of the dataset inserted per batch. A larger + * batch decreases graph quality but improves build speed. + * + * @param maxFraction the maximum fraction + * @return an instance of this Builder + */ + public Builder withMaxFraction(float maxFraction) { + this.maxFraction = maxFraction; + return this; + } + + /** + * Sets the growth rate base for batch sizes. + * + * @param batchBase the batch base + * @return an instance of this Builder + */ + public Builder withBatchBase(float batchBase) { + this.batchBase = batchBase; + return this; + } + + /** + * Sets the candidate queue size. The native builder expects a value of the + * form {@code (2^x) - 1}. + * + * @param queueSize the queue size + * @return an instance of this Builder + */ + public Builder withQueueSize(int queueSize) { + this.queueSize = queueSize; + return this; + } + + /** + * Sets the maximum batch size of reverse edge processing, which bounds the + * memory footprint of that stage. + * + * @param reverseBatchSize the reverse batch size + * @return an instance of this Builder + */ + public Builder withReverseBatchSize(int reverseBatchSize) { + this.reverseBatchSize = reverseBatchSize; + return this; + } + + /** + * Sets the distance metric. + * + * @param metric the distance metric + * @return an instance of this Builder + */ + public Builder withMetric(CuvsDistanceType metric) { + this.metric = metric; + return this; + } + + /** + * Builds an instance of {@link VamanaIndexParams}. + * + * @return an instance of {@link VamanaIndexParams} + */ + public VamanaIndexParams build() { + validate(); + return new VamanaIndexParams( + graphDegree, + visitedSize, + vamanaIters, + alpha, + maxFraction, + batchBase, + queueSize, + reverseBatchSize, + metric); + } + + /** + * Mirrors the checks the native builder performs, so that an invalid + * configuration fails here with a readable message rather than inside a + * GPU kernel. + */ + private void validate() { + if (Arrays.binarySearch(SUPPORTED_GRAPH_DEGREES, graphDegree) < 0) { + throw new IllegalArgumentException( + "graphDegree must be one of " + + Arrays.toString(SUPPORTED_GRAPH_DEGREES) + + ", was " + + graphDegree); + } + if (visitedSize <= graphDegree) { + throw new IllegalArgumentException( + "visitedSize must be greater than graphDegree, was " + + visitedSize + + " with graphDegree " + + graphDegree); + } + if (!Float.isFinite(vamanaIters) || vamanaIters < 1.0f) { + throw new IllegalArgumentException( + "vamanaIters must be finite and at least 1.0, was " + vamanaIters); + } + if (!Float.isFinite(alpha) || alpha <= 0.0f) { + throw new IllegalArgumentException("alpha must be finite and positive, was " + alpha); + } + if (!Float.isFinite(maxFraction) || maxFraction <= 0.0f || maxFraction > 1.0f) { + throw new IllegalArgumentException( + "maxFraction must be finite and in (0, 1], was " + maxFraction); + } + if (!Float.isFinite(batchBase) || batchBase <= 1.0f) { + throw new IllegalArgumentException( + "batchBase must be finite and greater than 1.0, was " + batchBase); + } + if (queueSize <= 0 || Integer.bitCount(queueSize + 1) != 1) { + throw new IllegalArgumentException( + "queueSize must be positive and of the form (2^x) - 1, was " + queueSize); + } + if (reverseBatchSize <= 0) { + throw new IllegalArgumentException( + "reverseBatchSize must be positive, was " + reverseBatchSize); + } + Objects.requireNonNull(metric, "metric must not be null"); + } + } +} 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 44706d0b6a..4df6b19477 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 @@ -136,6 +136,18 @@ BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; + /** + * Creates a new VamanaIndex Builder. + *

+ * This is a {@code default} method rather than an abstract one so that + * providers written against an earlier version of this interface keep + * compiling. + */ + default VamanaIndex.Builder newVamanaIndexBuilder(CuVSResources cuVSResources) + throws UnsupportedOperationException { + throw new UnsupportedOperationException("This provider does not support Vamana indexes"); + } + /** Creates a new HnswIndex Builder. */ HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index fd1cf7746c..4244b91676 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -64,6 +64,11 @@ public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) { throw new UnsupportedOperationException(reasons); } + @Override + public VamanaIndex.Builder newVamanaIndexBuilder(CuVSResources cuVSResources) { + throw new UnsupportedOperationException(reasons); + } + @Override public HnswIndex hnswIndexFromCagra(HnswIndexParams hnswParams, CagraIndex cagraIndex) throws Throwable { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java index 5e3088cdb2..dd117515f3 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java @@ -92,6 +92,25 @@ public void close() { } } + public static CloseableHandle createVamanaIndexParams() { + try (var localArena = Arena.ofConfined()) { + var paramsPtrPtr = localArena.allocate(cuvsVamanaIndexParams_t); + checkCuVSError(cuvsVamanaIndexParamsCreate(paramsPtrPtr), "cuvsVamanaIndexParamsCreate"); + var paramsPtr = paramsPtrPtr.get(cuvsVamanaIndexParams_t, 0L); + return new CloseableHandle() { + @Override + public MemorySegment handle() { + return paramsPtr; + } + + @Override + public void close() { + checkCuVSError(cuvsVamanaIndexParamsDestroy(paramsPtr), "cuvsVamanaIndexParamsDestroy"); + } + }; + } + } + public static CloseableHandle createIvfPqIndexParams() { try (var localArena = Arena.ofConfined()) { var paramsPtrPtr = localArena.allocate(cuvsIvfPqIndexParams_t); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/VamanaIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/VamanaIndexImpl.java new file mode 100644 index 0000000000..08c250e49e --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/VamanaIndexImpl.java @@ -0,0 +1,292 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.CuVSParamsHelper.createVamanaIndexParams; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; +import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; +import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; +import static com.nvidia.cuvs.internal.panama.headers_h.*; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.VamanaIndex; +import com.nvidia.cuvs.VamanaIndexParams; +import com.nvidia.cuvs.internal.common.CloseableHandle; +import com.nvidia.cuvs.internal.panama.cuvsVamanaIndexParams; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.nio.file.Path; +import java.util.Objects; + +/** + * {@link VamanaIndex} encapsulates a Vamana index, along with methods to build + * it on the GPU and serialize it in the DiskANN file format. + *

+ * cuVS provides build and serialize for Vamana but no search entry point, so + * this class deliberately exposes no search method. + * + * @since 26.10 + */ +public class VamanaIndexImpl implements VamanaIndex { + + private final CuVSResources resources; + private final MemorySegment vamanaIndexReference; + private final CuVSMatrix dataset; + private final boolean ownsDataset; + private boolean destroyed; + + private VamanaIndexImpl( + VamanaIndexParams indexParameters, + CuVSMatrix dataset, + boolean ownsDataset, + CuVSResources resources) { + Objects.requireNonNull(dataset); + this.resources = resources; + // the native index may retain a non-owning device view of the dataset, so + // we hold a reference to keep it alive for at least as long as the index + this.dataset = dataset; + this.ownsDataset = ownsDataset; + if (!(dataset instanceof CuVSMatrixInternal internalDataset)) { + throw new IllegalArgumentException( + "dataset must be created through CuVSMatrix, was " + dataset.getClass().getName()); + } + checkSupportedDataType(dataset.dataType()); + this.vamanaIndexReference = build(indexParameters, internalDataset); + } + + /** + * The native Vamana builder is instantiated for {@code float}, {@code half}, + * {@code int8}, and {@code uint8} only. Of those, {@code int8} has no + * corresponding {@link CuVSMatrix.DataType}. Reject anything else here rather + * than inside a kernel. + */ + private static void checkSupportedDataType(CuVSMatrix.DataType dataType) { + switch (dataType) { + case FLOAT, HALF, BYTE -> {} + default -> + throw new IllegalArgumentException( + "Vamana supports FLOAT, HALF, and BYTE datasets, was " + dataType); + } + } + + private void checkNotDestroyed() { + if (destroyed) { + throw new IllegalStateException("destroyed"); + } + } + + @Override + public void close() throws Exception { + checkNotDestroyed(); + destroyed = true; + Throwable failure = null; + try { + checkCuVSError(cuvsVamanaIndexDestroy(vamanaIndexReference), "cuvsVamanaIndexDestroy"); + } catch (Throwable t) { + failure = t; + } + if (ownsDataset) { + // attempt this even if the index failed to destroy, so an owned dataset + // is never stranded + try { + dataset.close(); + } catch (Throwable t) { + if (failure == null) { + failure = t; + } else { + failure.addSuppressed(t); + } + } + } + if (failure instanceof Error error) { + throw error; + } + if (failure != null) { + throw (Exception) failure; + } + } + + /** + * Creates the native index handle. The handle is a native heap allocation and + * must be released with {@code cuvsVamanaIndexDestroy}, so it is deliberately + * not tied to an {@link Arena}. + */ + private static MemorySegment createVamanaIndex() { + try (var localArena = Arena.ofConfined()) { + MemorySegment indexPtrPtr = localArena.allocate(cuvsVamanaIndex_t); + checkCuVSError(cuvsVamanaIndexCreate(indexPtrPtr), "cuvsVamanaIndexCreate"); + return indexPtrPtr.get(cuvsVamanaIndex_t, 0); + } + } + + /** + * Populates a native parameter struct from the Java parameters. A null + * argument leaves the native defaults in place. + */ + private static CloseableHandle segmentFromIndexParams(VamanaIndexParams params) { + var handle = createVamanaIndexParams(); + if (params == null) { + return handle; + } + try { + MemorySegment seg = handle.handle(); + cuvsVamanaIndexParams.graph_degree(seg, params.getGraphDegree()); + cuvsVamanaIndexParams.visited_size(seg, params.getVisitedSize()); + cuvsVamanaIndexParams.vamana_iters(seg, params.getVamanaIters()); + cuvsVamanaIndexParams.alpha(seg, params.getAlpha()); + cuvsVamanaIndexParams.max_fraction(seg, params.getMaxFraction()); + cuvsVamanaIndexParams.batch_base(seg, params.getBatchBase()); + cuvsVamanaIndexParams.queue_size(seg, params.getQueueSize()); + cuvsVamanaIndexParams.reverse_batchsize(seg, params.getReverseBatchSize()); + cuvsVamanaIndexParams.metric(seg, params.getMetric().value); + return handle; + } catch (RuntimeException | Error e) { + handle.close(); + throw e; + } + } + + /** + * Invokes the native {@code cuvsVamanaBuild} function to build the + * {@link VamanaIndex}. + * + * @return the handle of the built index + */ + private MemorySegment build(VamanaIndexParams indexParameters, CuVSMatrixInternal dataset) { + try (var indexParams = segmentFromIndexParams(indexParameters); + var localArena = Arena.ofConfined()) { + + var datasetTensor = dataset.toTensor(localArena); + var index = createVamanaIndex(); + try { + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + checkCuVSError( + cuvsVamanaBuild(cuvsRes, indexParams.handle(), datasetTensor, index), + "cuvsVamanaBuild"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + } + } catch (RuntimeException | Error e) { + // the index handle is a native allocation, so release it if the build + // never completed + checkCuVSError(cuvsVamanaIndexDestroy(index), "cuvsVamanaIndexDestroy"); + throw e; + } + return index; + } + } + + @Override + public int getDimensions() { + checkNotDestroyed(); + try (var localArena = Arena.ofConfined()) { + MemorySegment dims = localArena.allocate(C_INT); + checkCuVSError(cuvsVamanaIndexGetDims(vamanaIndexReference, dims), "cuvsVamanaIndexGetDims"); + return dims.get(C_INT, 0); + } + } + + @Override + public void serialize(Path filePrefix, boolean includeDataset) { + checkNotDestroyed(); + Objects.requireNonNull(filePrefix); + try (var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { + MemorySegment prefix = buildMemorySegment(localArena, filePrefix.toAbsolutePath().toString()); + checkCuVSError( + cuvsVamanaSerialize( + resourcesAccessor.handle(), prefix, vamanaIndexReference, includeDataset), + "cuvsVamanaSerialize"); + } + } + + @Override + public CuVSResources getCuVSResources() { + return resources; + } + + public static VamanaIndex.Builder newBuilder(CuVSResources cuvsResources) { + return new Builder(Objects.requireNonNull(cuvsResources)); + } + + /** + * Builder helps configure and create an instance of {@link VamanaIndex}. + */ + public static class Builder implements VamanaIndex.Builder { + + private final CuVSResources cuvsResources; + private CuVSMatrix dataset; + private boolean ownsDataset; + private VamanaIndexParams vamanaIndexParams; + + public Builder(CuVSResources cuvsResources) { + this.cuvsResources = cuvsResources; + } + + @Override + public Builder withDataset(float[][] vectors) { + // build the matrix first, then release any matrix this builder previously + // created, so a second call cannot strand the first one + CuVSMatrix created = CuVSMatrix.ofArray(vectors); + releaseOwnedDataset(); + this.dataset = created; + // we created it, so we close it + this.ownsDataset = true; + return this; + } + + @Override + public Builder withDataset(CuVSMatrix dataset) { + releaseOwnedDataset(); + this.dataset = dataset; + // the caller created it, so the caller closes it + this.ownsDataset = false; + return this; + } + + private void releaseOwnedDataset() { + if (ownsDataset && dataset != null) { + try { + dataset.close(); + } catch (Exception e) { + throw new RuntimeException("Failed to close the previously supplied dataset", e); + } + } + this.dataset = null; + this.ownsDataset = false; + } + + @Override + public Builder withIndexParams(VamanaIndexParams vamanaIndexParameters) { + this.vamanaIndexParams = vamanaIndexParameters; + return this; + } + + @Override + public VamanaIndexImpl build() { + if (dataset == null) { + throw new IllegalArgumentException("dataset must be provided"); + } + boolean transferred = false; + try { + VamanaIndexImpl index = + new VamanaIndexImpl(vamanaIndexParams, dataset, ownsDataset, cuvsResources); + // ownership now belongs to the index + transferred = true; + this.dataset = null; + this.ownsDataset = false; + return index; + } finally { + // a failed construction must not strand a matrix this builder created + if (!transferred) { + releaseOwnedDataset(); + } + } + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 1f16a4e904..bb27f36250 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -264,6 +264,11 @@ public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) { return CagraIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources)); } + @Override + public VamanaIndex.Builder newVamanaIndexBuilder(CuVSResources cuVSResources) { + return VamanaIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources)); + } + @Override public FilterBitsetHandle newFilterBitsetHandle(long[] combinedLongs) { return new FilterBitsetHandleImpl(combinedLongs); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/VamanaBuildAndSerializeIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/VamanaBuildAndSerializeIT.java new file mode 100644 index 0000000000..f347b43a21 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/VamanaBuildAndSerializeIT.java @@ -0,0 +1,426 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.*; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Comparator; +import java.util.Random; +import java.util.stream.Stream; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +/** + * Build and serialize tests for {@link VamanaIndex}. + * + *

cuVS exposes no Vamana search entry point, so these tests cover + * construction, dimensions, the DiskANN file layout produced by serialization, + * parameter validation, and lifecycle. + */ +@RunWith(RandomizedRunner.class) +public class VamanaBuildAndSerializeIT extends CuVSTestCase { + + private static final int ROWS = 1000; + private static final int DIMENSIONS = 32; + private static final int GRAPH_DEGREE = 32; + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + initializeRandom(); + } + + private static float[][] randomFloatDataset() { + Random random = new Random(42); + float[][] dataset = new float[ROWS][DIMENSIONS]; + for (int i = 0; i < ROWS; i++) { + for (int j = 0; j < DIMENSIONS; j++) { + dataset[i][j] = random.nextFloat(); + } + } + return dataset; + } + + private static VamanaIndexParams defaultParams() { + return new VamanaIndexParams.Builder() + .withGraphDegree(GRAPH_DEGREE) + .withVisitedSize(64) + .build(); + } + + /** Runs the body against a fresh output prefix and removes every file afterwards. */ + private static void withPrefix(PrefixConsumer body) throws Throwable { + Path dir = Files.createTempDirectory("cuvs-vamana"); + try { + body.accept(dir.resolve("index")); + } finally { + deleteRecursively(dir); + } + } + + private interface PrefixConsumer { + void accept(Path prefix) throws Throwable; + } + + private static void deleteRecursively(Path root) throws IOException { + if (!Files.exists(root)) { + return; + } + try (Stream paths = Files.walk(root)) { + paths + .sorted(Comparator.reverseOrder()) + .forEach( + path -> { + try { + Files.deleteIfExists(path); + } catch (IOException e) { + throw new UncheckedIOExceptionWrapper(e); + } + }); + } catch (UncheckedIOExceptionWrapper e) { + throw e.cause; + } + } + + private static final class UncheckedIOExceptionWrapper extends RuntimeException { + private final IOException cause; + + UncheckedIOExceptionWrapper(IOException cause) { + super(cause); + this.cause = cause; + } + } + + private static long sizeOf(Path path) throws IOException { + assertTrue(path + " should exist", Files.exists(path)); + long size = Files.size(path); + assertTrue(path + " should not be empty", size > 0); + return size; + } + + private static Path dataFile(Path prefix) { + return prefix.resolveSibling(prefix.getFileName() + ".data"); + } + + private static ByteBuffer readHead(Path path, int bytes) throws IOException { + byte[] head = new byte[bytes]; + try (var in = Files.newInputStream(path)) { + assertEquals(bytes, in.readNBytes(head, 0, bytes)); + } + return ByteBuffer.wrap(head).order(ByteOrder.LITTLE_ENDIAN); + } + + private static CuVSMatrix hostMatrix(CuVSMatrix.DataType dataType) { + return fill(CuVSMatrix.hostBuilder(ROWS, DIMENSIONS, dataType), dataType); + } + + private static CuVSMatrix deviceMatrix(CuVSResources resources, CuVSMatrix.DataType dataType) { + return fill(CuVSMatrix.deviceBuilder(resources, ROWS, DIMENSIONS, dataType), dataType); + } + + private static CuVSMatrix fill( + CuVSMatrix.Builder builder, CuVSMatrix.DataType dataType) { + Random random = new Random(7); + for (int i = 0; i < ROWS; i++) { + switch (dataType) { + case FLOAT -> { + float[] row = new float[DIMENSIONS]; + for (int j = 0; j < DIMENSIONS; j++) { + row[j] = random.nextFloat(); + } + builder.addVector(row); + } + case BYTE -> { + byte[] row = new byte[DIMENSIONS]; + random.nextBytes(row); + builder.addVector(row); + } + case HALF -> { + short[] row = new short[DIMENSIONS]; + for (int j = 0; j < DIMENSIONS; j++) { + row[j] = Float.floatToFloat16(random.nextFloat()); + } + builder.addVector(row); + } + case INT -> { + int[] row = new int[DIMENSIONS]; + for (int j = 0; j < DIMENSIONS; j++) { + row[j] = random.nextInt(100); + } + builder.addVector(row); + } + default -> throw new IllegalArgumentException("unhandled type " + dataType); + } + } + return builder.build(); + } + + @Test + public void testBuildFloatAndGetDimensions() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(randomFloatDataset()) + .withIndexParams(defaultParams()) + .build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + } + } + + @Test + public void testBuildUnsignedByteDataset() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix dataset = hostMatrix(CuVSMatrix.DataType.BYTE); + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(defaultParams()) + .build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + + withPrefix( + prefix -> { + index.serialize(prefix, true); + // one byte per component, plus the two 32-bit header values + assertEquals((long) ROWS * DIMENSIONS + 8, sizeOf(dataFile(prefix))); + }); + } + } + + @Test + public void testSerializeWritesGraphAndDataset() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(randomFloatDataset()) + .withIndexParams(defaultParams()) + .build()) { + + withPrefix( + prefix -> { + index.serialize(prefix, true); + + long graphSize = sizeOf(prefix); + Path data = dataFile(prefix); + + // the DiskANN .data file is two 32-bit header values followed by the + // raw vectors + ByteBuffer dataHead = readHead(data, 8); + assertEquals(ROWS, dataHead.getInt()); + assertEquals(DIMENSIONS, dataHead.getInt()); + assertEquals((long) ROWS * DIMENSIONS * Float.BYTES + 8, sizeOf(data)); + + // the graph file opens with its own length, then the observed + // maximum degree, which the configured graph degree bounds + ByteBuffer graphHead = readHead(prefix, 12); + assertEquals(graphSize, graphHead.getLong()); + int maxDegree = graphHead.getInt(); + assertTrue("max degree should be positive, was " + maxDegree, maxDegree > 0); + assertTrue( + "max degree " + maxDegree + " should not exceed the configured graph degree", + maxDegree <= GRAPH_DEGREE); + }); + } + } + + @Test + public void testSerializeWithoutDataset() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(randomFloatDataset()) + .withIndexParams(defaultParams()) + .build()) { + + withPrefix( + prefix -> { + index.serialize(prefix, false); + sizeOf(prefix); + assertFalse( + "the dataset file should not be written when includeDataset is false", + Files.exists(dataFile(prefix))); + }); + } + } + + @Test + public void testDefaultParametersAreUsedWhenNoneAreGiven() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + VamanaIndex index = + VamanaIndex.newBuilder(resources).withDataset(randomFloatDataset()).build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + } + } + + @Test + public void testUnsupportedGraphDegreeIsRejectedInJava() { + var builder = new VamanaIndexParams.Builder().withGraphDegree(48).withVisitedSize(128); + IllegalArgumentException e = assertThrows(IllegalArgumentException.class, builder::build); + assertTrue(e.getMessage(), e.getMessage().contains("graphDegree")); + } + + @Test + public void testVisitedSizeBelowGraphDegreeIsRejectedInJava() { + var builder = new VamanaIndexParams.Builder().withGraphDegree(64).withVisitedSize(8); + IllegalArgumentException e = assertThrows(IllegalArgumentException.class, builder::build); + assertTrue(e.getMessage(), e.getMessage().contains("visitedSize")); + } + + @Test + public void testInvalidVamanaItersIsRejectedInJava() { + var builder = new VamanaIndexParams.Builder().withVamanaIters(0.5f); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void testInvalidQueueSizeIsRejectedInJava() { + var builder = new VamanaIndexParams.Builder().withQueueSize(100); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void testSupportedGraphDegreesAreDefensivelyCopied() { + int[] degrees = VamanaIndexParams.supportedGraphDegrees(); + assertArrayEquals(new int[] {32, 64, 128, 256}, degrees); + degrees[0] = -1; + assertArrayEquals(new int[] {32, 64, 128, 256}, VamanaIndexParams.supportedGraphDegrees()); + } + + @Test + public void testUnsupportedDataTypeIsRejected() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix dataset = hostMatrix(CuVSMatrix.DataType.INT)) { + var builder = + VamanaIndex.newBuilder(resources).withDataset(dataset).withIndexParams(defaultParams()); + IllegalArgumentException e = assertThrows(IllegalArgumentException.class, builder::build); + assertTrue(e.getMessage(), e.getMessage().contains("FLOAT, HALF, and BYTE")); + } + } + + @Test + public void testMissingDatasetIsRejected() throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + var builder = VamanaIndex.newBuilder(resources).withIndexParams(defaultParams()); + assertThrows(IllegalArgumentException.class, builder::build); + } + } + + @Test + public void testCallerSuppliedDatasetOutlivesTheIndex() throws Throwable { + // the native index may retain a non-owning device view, so a caller + // supplied matrix is the caller's to close + try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix dataset = hostMatrix(CuVSMatrix.DataType.FLOAT)) { + try (VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(defaultParams()) + .build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + } + // still usable after the index is closed + assertEquals(ROWS, dataset.size()); + } + } + + @Test + public void testBuilderReuseDoesNotStrandAnOwnedDataset() throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + var builder = VamanaIndex.newBuilder(resources).withIndexParams(defaultParams()); + // the first matrix is created and then replaced, which must release it + builder.withDataset(randomFloatDataset()); + builder.withDataset(randomFloatDataset()); + try (VamanaIndex index = builder.build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + } + // the builder handed ownership to the index, so it now has no dataset + assertThrows(IllegalArgumentException.class, builder::build); + } + } + + @Test + public void testBuildHalfDatasetOnDevice() throws Throwable { + // covers float16 and the device-backed matrix path together, since the + // native index may retain a non-owning device view of this matrix + try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix dataset = deviceMatrix(resources, CuVSMatrix.DataType.HALF); + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(defaultParams()) + .build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + + withPrefix( + prefix -> { + index.serialize(prefix, true); + // two bytes per component, plus the two 32-bit header values + assertEquals((long) ROWS * DIMENSIONS * 2 + 8, sizeOf(dataFile(prefix))); + }); + } + } + + @Test + public void testBuildFloatDatasetOnDevice() throws Throwable { + try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix dataset = deviceMatrix(resources, CuVSMatrix.DataType.FLOAT)) { + try (VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(defaultParams()) + .build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + } + // the caller's device matrix outlives the index + assertEquals(ROWS, dataset.size()); + } + } + + @Test + public void testSqrtL2MetricIsAccepted() throws Throwable { + VamanaIndexParams params = + new VamanaIndexParams.Builder() + .withGraphDegree(GRAPH_DEGREE) + .withVisitedSize(64) + .withMetric(VamanaIndexParams.CuvsDistanceType.L2SqrtExpanded) + .build(); + try (CuVSResources resources = CuVSResources.create(); + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(randomFloatDataset()) + .withIndexParams(params) + .build()) { + assertEquals(DIMENSIONS, index.getDimensions()); + } + } + + @Test + public void testInfiniteVamanaItersIsRejectedInJava() { + var builder = new VamanaIndexParams.Builder().withVamanaIters(Float.POSITIVE_INFINITY); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void testUseAfterCloseIsRejected() throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + VamanaIndex index = + VamanaIndex.newBuilder(resources) + .withDataset(randomFloatDataset()) + .withIndexParams(defaultParams()) + .build(); + index.close(); + assertThrows(IllegalStateException.class, index::getDimensions); + assertThrows(IllegalStateException.class, index::close); + } + } +}