From b780b7ff90fa33514d4e1b762285379c0f951f3f Mon Sep 17 00:00:00 2001 From: Shaunak Kapur Date: Wed, 9 Sep 2026 17:54:19 +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/docs.yml | 4 + fern/pages/java_api/index.md | 2 + ...va-api-com-nvidia-cuvs-spi-cuvsprovider.md | 46 +- .../java-api-com-nvidia-cuvs-vamanaindex.md | 201 +++++++++ ...a-api-com-nvidia-cuvs-vamanaindexparams.md | 345 ++++++++++++++ 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 ++++++++++++++++++ 14 files changed, 1878 insertions(+), 17 deletions(-) create mode 100644 fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindex.md create mode 100644 fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindexparams.md 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/docs.yml b/fern/docs.yml index 9e47025b33..637f6c5ea8 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -496,6 +496,10 @@ navigation: path: "./pages/java_api/java-api-com-nvidia-cuvs-rowview.md" - page: "SynchronizedCuVSResources" path: "./pages/java_api/java-api-com-nvidia-cuvs-synchronizedcuvsresources.md" + - page: "VamanaIndex" + path: "./pages/java_api/java-api-com-nvidia-cuvs-vamanaindex.md" + - page: "VamanaIndexParams" + path: "./pages/java_api/java-api-com-nvidia-cuvs-vamanaindexparams.md" - page: "BruteForceIndex" path: "./pages/java_api/java-api-com-nvidia-cuvs-bruteforceindex.md" - page: "BruteForceIndexParams" diff --git a/fern/pages/java_api/index.md b/fern/pages/java_api/index.md index d2385b9a1a..753210a613 100644 --- a/fern/pages/java_api/index.md +++ b/fern/pages/java_api/index.md @@ -18,6 +18,8 @@ These pages are generated from the Java source files in `java/cuvs-java/src/main - [LibraryException](/api-reference/java-api-com-nvidia-cuvs-libraryexception) - [RowView](/api-reference/java-api-com-nvidia-cuvs-rowview) - [SynchronizedCuVSResources](/api-reference/java-api-com-nvidia-cuvs-synchronizedcuvsresources) +- [VamanaIndex](/api-reference/java-api-com-nvidia-cuvs-vamanaindex) +- [VamanaIndexParams](/api-reference/java-api-com-nvidia-cuvs-vamanaindexparams) ## Nearest Neighbors diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md index f19b50fb6b..84f130ff31 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md @@ -211,6 +211,20 @@ Creates a new CagraIndex Builder. _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:135`_ +### newVamanaIndexBuilder + +```java +default VamanaIndex.Builder newVamanaIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException +``` + +Creates a new VamanaIndex Builder. + +This is a `default` method rather than an abstract one so that +providers written against an earlier version of this interface keep +compiling. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:145`_ + ### newHnswIndexBuilder ```java @@ -219,7 +233,7 @@ HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) throws Unsupp Creates a new HnswIndex Builder. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:139`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:151`_ ### hnswIndexFromCagra @@ -246,7 +260,7 @@ A new HNSW index | --- | --- | | `Throwable` | if an error occurs during conversion | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:150`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:162`_ ### hnswIndexBuild @@ -274,7 +288,7 @@ A new HNSW index ready for search | --- | --- | | `Throwable` | if an error occurs during building | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:161`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:173`_ ### newTieredIndexBuilder @@ -284,7 +298,7 @@ TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) throws Un Creates a new TieredIndex Builder. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:165`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:177`_ ### isCagraPaddedDataset @@ -317,7 +331,7 @@ true when the rows are already padded the way CAGRA requires | --- | --- | | `UnsupportedOperationException` | if this provider cannot answer | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:182`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:194`_ ### mergeCagraIndexes @@ -347,7 +361,7 @@ A new merged CAGRA index | --- | --- | | `Throwable` | if an error occurs during the merge operation | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:198`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:210`_ ### newFilterBitsetHandle @@ -364,7 +378,7 @@ Per-partition bit offsets are recomputed inside cuVS from the index sizes. | --- | --- | | `combinedLongs` | packed bitset words for a single partition | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:207`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:219`_ ### searchCagraMultiPartition @@ -390,7 +404,7 @@ Searches multiple CAGRA index partitions for the global top-k nearest neighbors | --- | --- | | `Throwable` | if an error occurs during the search | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:220`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:232`_ ### gpuInfoProvider @@ -400,7 +414,7 @@ GPUInfoProvider gpuInfoProvider() Returns a `GPUInfoProvider` to query the system for GPU related information -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:229`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:241`_ ### enableRMMPooledMemory @@ -426,7 +440,7 @@ concurrent construction of that class from multiple threads. | `initialPoolSizePercent` | The initial pool size, in percentage of the total GPU memory | | `maxPoolSizePercent` | The maximum pool size, in percentage of the total GPU memory | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:250`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:262`_ ### enableRMMManagedPooledMemory @@ -447,7 +461,7 @@ Ownership: see `#enableRMMPooledMemory`. | `initialPoolSizePercent` | The initial pool size, in percentage of the total GPU memory | | `maxPoolSizePercent` | The maximum pool size, in percentage of the total GPU memory | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:262`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:274`_ ### enableRMMAsyncMemory @@ -464,7 +478,7 @@ launches. This operation has a global effect and will affect all resources on th Ownership: see `#enableRMMPooledMemory`. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:274`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:286`_ ### resetRMMPooledMemory @@ -476,7 +490,7 @@ Disables pooled memory on the current device, reverting back to the default sett Ownership: see `#enableRMMPooledMemory`. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:281`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:293`_ ### provider @@ -486,7 +500,7 @@ static CuVSProvider provider() Retrieves the system-wide provider. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:284`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:296`_ ### cagraIndexParamsFromHnswParams @@ -516,7 +530,7 @@ may be shifted along the curve right or left. See the heuristics descriptions fo A new CAGRA index parameters object -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:304`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:316`_ ### cagraIndexParamsFromDataset @@ -540,6 +554,6 @@ Create CAGRA index parameters heuristically tuned for a dataset. A new CAGRA index parameters object -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:322`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:334`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:18`_ diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindex.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindex.md new file mode 100644 index 0000000000..22fa850fe4 --- /dev/null +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindex.md @@ -0,0 +1,201 @@ +--- +slug: api-reference/java-api-com-nvidia-cuvs-vamanaindex +--- + +# VamanaIndex + +_Java package: `com.nvidia.cuvs`_ + +```java +public interface VamanaIndex extends AutoCloseable +``` + +`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. + +## Public Members + +### getDimensions + +```java +int getDimensions() throws Throwable +``` + +Gets the dimensionality of the vectors in this index. + +**Returns** + +the number of dimensions + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:30`_ + +### serialize + +```java +default void serialize(Path filePrefix) throws Throwable +``` + +Serializes the index in the DiskANN file format, including the dataset. + +This writes two files, `filePrefix` holding the graph and +`filePrefix + ".data"` holding the dataset. + +**Parameters** + +| Name | Description | +| --- | --- | +| `filePrefix` | the prefix that output file names are derived from | + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:40`_ + +### serialize + +```java +void serialize(Path filePrefix, boolean includeDataset) throws Throwable +``` + +Serializes the index in the DiskANN file format. + +When `includeDataset` is true this writes `filePrefix` holding +the graph and `filePrefix + ".data"` holding the dataset. When it is +false only `filePrefix` is written. + +The argument is a prefix and not a complete file name, matching the native +`file_prefix` parameter. + +**Parameters** + +| Name | Description | +| --- | --- | +| `filePrefix` | the prefix that output file names are derived from | +| `includeDataset` | whether to write the dataset alongside the graph | + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:57`_ + +### getCuVSResources + +```java +CuVSResources getCuVSResources() +``` + +Gets an instance of `CuVSResources` + +**Returns** + +an instance of `CuVSResources` + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:64`_ + +### newBuilder + +```java +static Builder newBuilder(CuVSResources cuvsResources) +``` + +Creates a new Builder with an instance of `CuVSResources`. + +**Parameters** + +| Name | Description | +| --- | --- | +| `cuvsResources` | an instance of `CuVSResources` | + +**Throws** + +| Type | Description | +| --- | --- | +| `UnsupportedOperationException` | if the provider does not support cuvs | + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:72`_ + +### withDataset + +```java +Builder withDataset(float[][] vectors) +``` + +Sets the dataset for building the `VamanaIndex`. + +**Parameters** + +| Name | Description | +| --- | --- | +| `vectors` | a two-dimensional float array | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:88`_ + +### withDataset + +```java +Builder withDataset(CuVSMatrix dataset) +``` + +Sets the dataset for building the `VamanaIndex`. + +The native builder accepts `float`, `half`, `uint8`, +and `int8` datasets. Of those, `CuVSMatrix.DataType#FLOAT`, +`CuVSMatrix.DataType#HALF`, and `CuVSMatrix.DataType#BYTE` +are reachable from Java today, where `BYTE` is unsigned. +`int8` has no corresponding `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 `float[][]` is created and closed by the index instead. + +**Parameters** + +| Name | Description | +| --- | --- | +| `dataset` | a `CuVSMatrix` object containing the vectors | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:107`_ + +### withIndexParams + +```java +Builder withIndexParams(VamanaIndexParams vamanaIndexParameters) +``` + +Registers an instance of configured `VamanaIndexParams` with this +Builder. + +**Parameters** + +| Name | Description | +| --- | --- | +| `vamanaIndexParameters` | An instance of VamanaIndexParams | + +**Returns** + +An instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:116`_ + +### build + +```java +VamanaIndex build() throws Throwable +``` + +Builds and returns an instance of `VamanaIndex`. + +**Returns** + +an instance of `VamanaIndex` + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:123`_ + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndex.java:21`_ diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindexparams.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindexparams.md new file mode 100644 index 0000000000..c3e7d8a9f5 --- /dev/null +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-vamanaindexparams.md @@ -0,0 +1,345 @@ +--- +slug: api-reference/java-api-com-nvidia-cuvs-vamanaindexparams +--- + +# VamanaIndexParams + +_Java package: `com.nvidia.cuvs`_ + +```java +public class VamanaIndexParams +``` + +Supplemental parameters to build a Vamana index. + +The defaults match the native `cuvs::neighbors::vamana::index_params` +defaults. + +## Public Members + +### supportedGraphDegrees + +```java +public static int[] supportedGraphDegrees() +``` + +Returns the graph degrees the native Vamana builder supports. + +**Returns** + +a copy of the supported graph degrees, in ascending order + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:30`_ + +### L2Expanded + +```java +L2Expanded(0), /** * Euclidean, the square root of +``` + +Squared L2. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:45`_ + +### L2SqrtExpanded + +```java +L2SqrtExpanded(1) +``` + +Euclidean, the square root of `#L2Expanded`. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:50`_ + +### getGraphDegree + +```java +public int getGraphDegree() +``` + +Gets the maximum degree of the output graph, the R parameter in the Vamana +literature. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:97`_ + +### getVisitedSize + +```java +public int getVisitedSize() +``` + +Gets the maximum number of visited nodes per search, the L parameter in the +Vamana literature. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:105`_ + +### getVamanaIters + +```java +public float getVamanaIters() +``` + +Gets the number of Vamana vector insertion iterations. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:112`_ + +### getAlpha + +```java +public float getAlpha() +``` + +Gets the alpha pruning parameter. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:119`_ + +### getMaxFraction + +```java +public float getMaxFraction() +``` + +Gets the maximum fraction of the dataset inserted per batch. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:126`_ + +### getBatchBase + +```java +public float getBatchBase() +``` + +Gets the growth rate base for batch sizes. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:133`_ + +### getQueueSize + +```java +public int getQueueSize() +``` + +Gets the candidate queue size. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:140`_ + +### getReverseBatchSize + +```java +public int getReverseBatchSize() +``` + +Gets the maximum batch size of reverse edge processing. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:147`_ + +### getMetric + +```java +public CuvsDistanceType getMetric() +``` + +Gets the distance metric. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:154`_ + +### withGraphDegree + +```java +public Builder withGraphDegree(int graphDegree) +``` + +Sets the maximum degree of the output graph. + +**Parameters** + +| Name | Description | +| --- | --- | +| `graphDegree` | the graph degree, one of `VamanaIndexParams#supportedGraphDegrees()` | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:205`_ + +### withVisitedSize + +```java +public Builder withVisitedSize(int visitedSize) +``` + +Sets the maximum number of visited nodes per search. + +The native builder requires this to be greater than the graph degree. + +**Parameters** + +| Name | Description | +| --- | --- | +| `visitedSize` | the visited size | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:218`_ + +### withVamanaIters + +```java +public Builder withVamanaIters(float vamanaIters) +``` + +Sets the number of Vamana vector insertion iterations. + +**Parameters** + +| Name | Description | +| --- | --- | +| `vamanaIters` | the iteration count | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:229`_ + +### withAlpha + +```java +public Builder withAlpha(float alpha) +``` + +Sets the alpha pruning parameter. + +**Parameters** + +| Name | Description | +| --- | --- | +| `alpha` | the alpha value | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:240`_ + +### withMaxFraction + +```java +public Builder withMaxFraction(float maxFraction) +``` + +Sets the maximum fraction of the dataset inserted per batch. A larger +batch decreases graph quality but improves build speed. + +**Parameters** + +| Name | Description | +| --- | --- | +| `maxFraction` | the maximum fraction | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:252`_ + +### withBatchBase + +```java +public Builder withBatchBase(float batchBase) +``` + +Sets the growth rate base for batch sizes. + +**Parameters** + +| Name | Description | +| --- | --- | +| `batchBase` | the batch base | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:263`_ + +### withQueueSize + +```java +public Builder withQueueSize(int queueSize) +``` + +Sets the candidate queue size. The native builder expects a value of the +form `(2^x) - 1`. + +**Parameters** + +| Name | Description | +| --- | --- | +| `queueSize` | the queue size | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:275`_ + +### withReverseBatchSize + +```java +public Builder withReverseBatchSize(int reverseBatchSize) +``` + +Sets the maximum batch size of reverse edge processing, which bounds the +memory footprint of that stage. + +**Parameters** + +| Name | Description | +| --- | --- | +| `reverseBatchSize` | the reverse batch size | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:287`_ + +### withMetric + +```java +public Builder withMetric(CuvsDistanceType metric) +``` + +Sets the distance metric. + +**Parameters** + +| Name | Description | +| --- | --- | +| `metric` | the distance metric | + +**Returns** + +an instance of this Builder + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:298`_ + +### build + +```java +public VamanaIndexParams build() +``` + +Builds an instance of `VamanaIndexParams`. + +**Returns** + +an instance of `VamanaIndexParams` + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:308`_ + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/VamanaIndexParams.java:18`_ 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..b23cced2d5 --- /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.12 + */ +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..b0e8ded341 --- /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.12 + */ +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 cb44f9b672..12899b6f00 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..13eff52b46 --- /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.12 + */ +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); + } + } +}