[Java] Add CAGRA index graph accessor/build from graph (host memory) (#1216)

ldematte · web-flow · commit 34032a813c9c · 2025-08-12T05:50:12.000Z
This PR leverages the functions introduced by #1086 and the data structures introduced by #1111 to access, copy, and re-create an index to/from a CAGRA graph. Supersedes #1105 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: #1216
diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java
@@ -50,6 +50,13 @@ public interface CagraIndex {
    */
   SearchResults search(CagraQuery query) throws Throwable;
 
+  /** Returns the CAGRA graph
+   *
+   * @return a {@link CuVSMatrix} encapsulating the native int (uint32_t) array used to represent
+   * the cagra graph
+   */
+  CuVSMatrix getGraph();
+
   /**
    * A method to persist a CAGRA index using an instance of {@link OutputStream}
    * for writing index bytes.
@@ -208,6 +215,12 @@ interface Builder {
      */
     Builder from(InputStream inputStream);
 
+    /**
+     * Sets a CAGRA graph instance to re-create an index from a
+     * previously built graph.
+     */
+    Builder from(CuVSMatrix graph);
+
     /**
      * Sets the dataset vectors for building the {@link CagraIndex}.
      *
diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java
@@ -30,6 +30,7 @@ public interface CuVSMatrix extends AutoCloseable {
   enum DataType {
     FLOAT,
     INT,
+    UINT,
     BYTE
   }
 
diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java
@@ -45,6 +45,7 @@
 import com.nvidia.cuvs.internal.common.CloseableHandle;
 import com.nvidia.cuvs.internal.common.CloseableRMMAllocation;
 import com.nvidia.cuvs.internal.common.CompositeCloseableHandle;
+import com.nvidia.cuvs.internal.common.Util;
 import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams;
 import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams;
 import com.nvidia.cuvs.internal.panama.cuvsCagraMergeParams;
@@ -122,6 +123,31 @@ private CagraIndexImpl(IndexReference indexReference, CuVSResources resources) {
     this.destroyed = false;
   }
 
+  /**
+   * Constructor for creating an index from a pre-build CAGRA graph
+   *
+   * @param metric      the distance type used
+   * @param graph       a previously built CAGRA graph
+   * @param dataset     the dataset used for indexing
+   * @param resources   an instance of {@link CuVSResources}
+   */
+  private CagraIndexImpl(
+      CagraIndexParams.CuvsDistanceType metric,
+      CuVSMatrix graph,
+      CuVSMatrix dataset,
+      CuVSResources resources) {
+    Objects.requireNonNull(graph);
+    Objects.requireNonNull(dataset);
+
+    this.resources = resources;
+
+    assert graph instanceof CuVSMatrixBaseImpl;
+    assert dataset instanceof CuVSMatrixBaseImpl;
+
+    this.cagraIndexReference =
+        fromGraph(metric, (CuVSMatrixBaseImpl) graph, (CuVSMatrixBaseImpl) dataset);
+  }
+
   private void checkNotDestroyed() {
     if (destroyed) {
       throw new IllegalStateException("destroyed");
@@ -164,9 +190,12 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImp
       omp_set_num_threads(numWriterThreads);
 
       MemorySegment dataSeg = dataset.memorySegment();
+      // TODO: type kDLCPU()/kDLCUDA() should be aligned with the CuVSMatrixBaseImpl type (host or
+      // device?)
 
       long[] datasetShape = {rows, cols};
-      MemorySegment datasetTensor = prepareTensor(localArena, dataSeg, datasetShape, 2, 32, 2, 1);
+      MemorySegment datasetTensor =
+          prepareTensor(localArena, dataSeg, datasetShape, kDLFloat(), 32, kDLCPU(), 1);
 
       var index = createCagraIndex();
 
@@ -381,6 +410,81 @@ public void serialize(OutputStream outputStream, Path tempFile, int bufferLength
     }
   }
 
+  @Override
+  public CuVSMatrix getGraph() {
+    try (var localArena = Arena.ofConfined()) {
+      var outPtr = localArena.allocate(uint32_t);
+      checkCuVSError(
+          cuvsCagraIndexGetGraphDegree(cagraIndexReference.getMemorySegment(), outPtr),
+          "cuvsCagraIndexGetGraphDegree");
+      long graphDegree = Util.dereferenceUnsignedInt(outPtr);
+
+      checkCuVSError(
+          cuvsCagraIndexGetSize(cagraIndexReference.getMemorySegment(), outPtr),
+          "cuvsCagraIndexGetSize");
+      long size = Util.dereferenceUnsignedInt(outPtr);
+
+      // TODO: use a "device" graph + tensor, avoid (defer) copy
+      var graph = new CuVSHostMatrixArenaImpl(size, graphDegree, CuVSMatrix.DataType.UINT);
+      var graphHostTensor = graph.toTensor(localArena);
+      var graphDeviceTensor =
+          prepareTensor(
+              localArena,
+              MemorySegment.NULL,
+              new long[] {size, graphDegree},
+              kDLUInt(),
+              32,
+              kDLCUDA(),
+              1);
+      checkCuVSError(
+          cuvsCagraIndexGetGraph(cagraIndexReference.getMemorySegment(), graphDeviceTensor),
+          "cuvsCagraIndexGetGraph");
+
+      try (var resourceAccess = resources.access()) {
+        var cuvsRes = resourceAccess.handle();
+        checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync");
+
+        checkCuVSError(
+            cuvsMatrixCopy(cuvsRes, graphDeviceTensor, graphHostTensor), "cuvsMatrixCopy");
+
+        checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync");
+      }
+
+      return graph;
+    }
+  }
+
+  private IndexReference fromGraph(
+      CagraIndexParams.CuvsDistanceType metric,
+      CuVSMatrixBaseImpl graph,
+      CuVSMatrixBaseImpl dataset) {
+    try (var localArena = Arena.ofConfined()) {
+      long rows = dataset.size();
+      long cols = dataset.columns();
+
+      var index = createCagraIndex();
+      try (var resourcesAccess = resources.access()) {
+        long cuvsRes = resourcesAccess.handle();
+
+        long[] datasetShape = {rows, cols};
+        MemorySegment datasetTensor =
+            prepareTensor(
+                localArena, dataset.memorySegment(), datasetShape, kDLFloat(), 32, kDLCPU(), 1);
+
+        long[] graphShape = {graph.size(), graph.columns()};
+        MemorySegment graphTensor =
+            prepareTensor(
+                localArena, graph.memorySegment(), graphShape, kDLUInt(), 32, kDLCPU(), 1);
+
+        checkCuVSError(
+            cuvsCagraIndexFromArgs(cuvsRes, metric.value, graphTensor, datasetTensor, index),
+            "cuvsCagraIndexFromArgs");
+      }
+
+      return new IndexReference(index, dataset);
+    }
+  }
+
   @Override
   public void serializeToHNSW(OutputStream outputStream) throws Throwable {
     Path path =
@@ -678,6 +782,7 @@ public static class Builder implements CagraIndex.Builder {
     private CagraIndexParams cagraIndexParams;
     private final CuVSResources cuvsResources;
     private InputStream inputStream;
+    private CuVSMatrix graph;
 
     public Builder(CuVSResources cuvsResources) {
       this.cuvsResources = cuvsResources;
@@ -689,6 +794,12 @@ public Builder from(InputStream inputStream) {
       return this;
     }
 
+    @Override
+    public Builder from(CuVSMatrix graph) {
+      this.graph = graph;
+      return this;
+    }
+
     @Override
     public Builder withDataset(float[][] vectors) {
       this.dataset = CuVSMatrix.ofArray(vectors);
@@ -712,7 +823,17 @@ public CagraIndexImpl build() throws Throwable {
       if (inputStream != null) {
         return new CagraIndexImpl(inputStream, cuvsResources);
       } else {
-        return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources);
+        if (graph != null) {
+          if (cagraIndexParams == null || dataset == null) {
+            throw new IllegalArgumentException(
+                "In order to reconstruct a CAGRA index from a graph, "
+                    + "you must specify the original dataset and the metric used.");
+          }
+          return new CagraIndexImpl(
+              cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources);
+        } else {
+          return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources);
+        }
       }
     }
   }
diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java
@@ -18,13 +18,12 @@
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT;
+import static com.nvidia.cuvs.internal.common.Util.prepareTensor;
+import static com.nvidia.cuvs.internal.panama.headers_h.*;
 
 import com.nvidia.cuvs.CuVSHostMatrix;
 import com.nvidia.cuvs.RowView;
-import java.lang.foreign.MemoryLayout;
-import java.lang.foreign.MemorySegment;
-import java.lang.foreign.SequenceLayout;
-import java.lang.foreign.ValueLayout;
+import java.lang.foreign.*;
 import java.lang.invoke.VarHandle;
 
 /**
@@ -61,7 +60,7 @@ protected CuVSHostMatrixImpl(
   protected static ValueLayout valueLayoutFromType(DataType dataType) {
     return switch (dataType) {
       case FLOAT -> C_FLOAT;
-      case INT -> C_INT;
+      case INT, UINT -> C_INT;
       case BYTE -> C_CHAR;
     };
   }
@@ -131,6 +130,12 @@ public ValueLayout valueLayout() {
     return valueLayout;
   }
 
+  @Override
+  public MemorySegment toTensor(Arena arena) {
+    return prepareTensor(
+        arena, memorySegment, new long[] {size, columns}, code(), bits(), kDLCPU(), 1);
+  }
+
   private static class SliceRowView implements RowView {
     private final MemorySegment memorySegment;
     private final long size;
diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java
@@ -15,7 +15,10 @@
  */
 package com.nvidia.cuvs.internal;
 
+import static com.nvidia.cuvs.internal.panama.headers_h.*;
+
 import com.nvidia.cuvs.CuVSMatrix;
+import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 
 public abstract class CuVSMatrixBaseImpl implements CuVSMatrix {
@@ -45,4 +48,21 @@ public long columns() {
   public MemorySegment memorySegment() {
     return memorySegment;
   }
+
+  protected int bits() {
+    return switch (dataType) {
+      case FLOAT, INT, UINT -> 32;
+      case BYTE -> 8;
+    };
+  }
+
+  protected int code() {
+    return switch (dataType) {
+      case FLOAT -> kDLFloat();
+      case INT -> kDLInt();
+      case UINT, BYTE -> kDLUInt();
+    };
+  }
+
+  public abstract MemorySegment toTensor(Arena arena);
 }
diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java
@@ -19,11 +19,7 @@
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG;
-import static com.nvidia.cuvs.internal.panama.headers_h.cudaGetDeviceCount;
-import static com.nvidia.cuvs.internal.panama.headers_h.cudaGetDeviceProperties_v2;
-import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemGetInfo;
-import static com.nvidia.cuvs.internal.panama.headers_h.cudaSetDevice;
-import static com.nvidia.cuvs.internal.panama.headers_h.size_t;
+import static com.nvidia.cuvs.internal.panama.headers_h.*;
 
 import com.nvidia.cuvs.GPUInfo;
 import com.nvidia.cuvs.internal.panama.DLDataType;
@@ -73,6 +69,13 @@ public static void checkCudaError(int value, String caller) {
     }
   }
 
+  private static final long UNSIGNED_INT_MASK = 0xFFFFFFFFL;
+
+  public static long dereferenceUnsignedInt(MemorySegment ptr) {
+    assert ptr.byteSize() == 4;
+    return ptr.get(uint32_t, 0) & UNSIGNED_INT_MASK;
+  }
+
   /**
    * Java analog to CUDA's cudaMemcpyKind, used for cudaMemcpy() calls.
    * @see <a href="https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html">CUDA Runtime API</a>
diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java
@@ -290,6 +290,48 @@ private Path createSerializedIndex(CuVSMatrix dataset) throws Throwable {
     }
   }
 
+  @Test
+  public void testReconstructIndexFromGraph() throws Throwable {
+    try (var dataset = CuVSMatrix.ofArray(createSampleData())) {
+      var queries = createSampleQueries();
+      List<Map<Integer, Float>> expectedResults = getExpectedResults();
+
+      try (CuVSResources resources = CuVSResources.create()) {
+        var index = indexOnce(dataset, resources);
+        var graph = index.getGraph();
+
+        var reconstructedIndex =
+            CagraIndex.newBuilder(resources)
+                .from(graph)
+                .withDataset(dataset)
+                .withIndexParams(
+                    new CagraIndexParams.Builder().withMetric(CuvsDistanceType.L2Expanded).build())
+                .build();
+        queryAndCompare(
+            index,
+            reconstructedIndex,
+            SearchResults.IDENTITY_MAPPING,
+            queries,
+            expectedResults,
+            resources);
+
+        var originalIndexPath = serializeOnce(index);
+        var reconstructedIndexPath = serializeOnce(reconstructedIndex);
+
+        var originalBytes = Files.readAllBytes(originalIndexPath);
+        var reconstructedBytes = Files.readAllBytes(reconstructedIndexPath);
+
+        assertArrayEquals(originalBytes, reconstructedBytes);
+
+        index.destroyIndex();
+        reconstructedIndex.destroyIndex();
+
+        Files.deleteIfExists(originalIndexPath);
+        Files.deleteIfExists(reconstructedIndexPath);
+      }
+    }
+  }
+
   @Test
   public void testIndexingAndSearchingFlowWithCustomMappingFunction() throws Throwable {
     var dataset = CuVSMatrix.ofArray(createSampleData());

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@ public interface CuVSMatrix extends AutoCloseable {`
`30`	`30`	`enum DataType {`
`31`	`31`	`FLOAT,`
`32`	`32`	`INT,`
	`33`	`+ UINT,`
`33`	`34`	`BYTE`
`34`	`35`	`}`
`35`	`36`