+ * 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