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