Fix GPU and Java-heap OOM when building large vector indexes

src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java

@@ -42,7 +42,7 @@ public static enum Strategy {
   public static final int MIN_GRAPH_DEG = 1;
   public static final int MAX_GRAPH_DEG = 512;
   public static final int MIN_HNSW_LAYERS = 1;
-  public static final int MAX_HNSW_LAYERS = 3;
+  public static final int MAX_HNSW_LAYERS = 99;
   public static final int MIN_MAX_CONN = 1;
   public static final int MAX_MAX_CONN = 512;
   public static final int MIN_BEAM_WIDTH = 1;

src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java

@@ -76,97 +76,84 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens
    * Creates a multi-layer HNSW graph with dynamic number of layers.
    * M = cagraGraphDegree/2
    * Each layer contains 1/M nodes from the previous layer
-   * Creates layers until the highest layer has ≤ M nodes
+   * Creates layers until the highest layer has <= M nodes
+   * <p>
+   * Vectors for higher-layer subsets are read directly from the native matrix
+   * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])},
+   * avoiding any additional heap allocation of the full dataset. Used by both
+   * the flush and merge paths; the caller provides the vectors as a
+   * {@link CuVSMatrix}.
    */
   public static GPUBuiltHnswGraph createMultiLayerHnswGraph(
       FieldInfo fieldInfo,
-      int size,
       int dimensions,
       CuVSMatrix adjacencyListMatrix,
-      List<?> vectors,
+      CuVSMatrix vectorDataset,
       int hnswLayers,
       int graphDegree,
       CagraIndexParams params,
       QuantizationType quantization)
       throws Throwable {
 
-    // Calculate M as cagraGraphDegree/2
+    int size = (int) vectorDataset.size();
     int M = graphDegree / 2;
 
-    // Store all layers data
     List<int[]> layerNodes = new ArrayList<>();
     List<CuVSMatrix> layerAdjacencies = new ArrayList<>();
 
     // Layer 0: Use full CAGRA adjacency list
-    layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list
+    layerNodes.add(null);
     layerAdjacencies.add(adjacencyListMatrix);
 
     int currentLayerSize = size;
     int layerIndex = 1;
     Random random = new Random();
 
     while (layerIndex < hnswLayers && currentLayerSize > 1) {
-      // Calculate size for next layer (1/M of current layer)
       int nextLayerSize = Math.max(2, currentLayerSize / M);
-      // Select nodes for this layer
       SortedSet<Integer> selectedNodesSet = new TreeSet<>();
 
       if (layerIndex == 1) {
-        // Select from all nodes (Layer 0)
         while (selectedNodesSet.size() < nextLayerSize) {
           selectedNodesSet.add(random.nextInt(size));
         }
       } else {
-        // Select from previous layer nodes
         int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1);
         while (selectedNodesSet.size() < nextLayerSize) {
-          int idx = random.nextInt(prevLayerNodes.length);
-          selectedNodesSet.add(prevLayerNodes[idx]);
+          selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]);
         }
       }
 
-      // Convert to sorted array
       int[] selectedNodes =
           selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray();
-
       layerNodes.add(selectedNodes);
 
       if (quantization == QuantizationType.NONE) {
-        // Extract vectors for selected nodes
-        float[][] selectedVectors = new float[nextLayerSize][];
+        // Read only the sampled rows from the native matrix — no full-dataset heap copy
+        float[][] selectedVectors = new float[nextLayerSize][dimensions];
         for (int i = 0; i < nextLayerSize; i++) {
-          selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]);
+          vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]);
         }
-
-        // Build CAGRA graph for this layer
         layerAdjacencies.add(
             buildCagraGraphForSubset(
                 selectedVectors, selectedNodes, 0, params, dimensions, quantization));
-
       } else {
-
-        // Extract vectors for selected nodes
-        int bytesPerVector = (dimensions + 7) / 8;
-        byte[][] selectedVectors = new byte[nextLayerSize][];
+        // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim.
+        int bytesPerVector = (int) vectorDataset.columns();
+        byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector];
         for (int i = 0; i < nextLayerSize; i++) {
-          selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]);
+          vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]);
         }
-
-        // Build CAGRA graph for this layer
         layerAdjacencies.add(
             buildCagraGraphForSubset(
                 selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization));
       }
 
-      // Update for next iteration
       currentLayerSize = nextLayerSize;
       layerIndex++;
-
-      // Use different seed for each layer
       random = new Random(new Random().nextLong());
     }
 
-    // Create the multi-layer graph with all layers
     return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies);
   }
 
@@ -185,11 +172,9 @@ private static CuVSMatrix buildCagraGraphForSubset(
     CuVSMatrix subsetDataset;
 
     if (quantization == QuantizationType.BINARY) {
-      subsetDataset =
-          createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance());
+      subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector);
     } else if (quantization == QuantizationType.SCALAR) {
-      subsetDataset =
-          createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance());
+      subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions);
     } else {
       subsetDataset = CuVSMatrix.ofArray((float[][]) vectors);
     }

src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java

@@ -201,8 +201,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<float[]> vectors) thro
         var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex);
         try {
           CuVSMatrix cagraDataset =
-              Utils.createFloatMatrix(
-                  vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance());
+              Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension());
           writeCagraIndex(cagraIndexOutputStream, cagraDataset);
         } catch (Throwable t) {
           // Fallback to brute force in a few cases, for now.
@@ -215,8 +214,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<float[]> vectors) thro
       if (indexType.isBruteForce()) {
         var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex);
         CuVSMatrix bruteforceDataset =
-            Utils.createFloatMatrix(
-                vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance());
+            Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension());
 
         writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset);
         bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset;

src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java

@@ -16,12 +16,12 @@
 import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME;
 import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance;
 import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance;
-import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors;
 import static org.apache.lucene.index.VectorEncoding.FLOAT32;
 import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance;
 
 import com.nvidia.cuvs.CagraIndex;
 import com.nvidia.cuvs.CagraIndexParams;
+import com.nvidia.cuvs.CuVSHostMatrix;
 import com.nvidia.cuvs.CuVSMatrix;
 import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType;
 import java.io.IOException;
@@ -34,11 +34,14 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.DocsWithFieldSet;
 import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.MergeState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.apache.lucene.index.Sorter;
 import org.apache.lucene.index.Sorter.DocMap;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.InfoStream;
@@ -139,7 +142,8 @@ public KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOException
   }
 
   /**
-   * Builds the intermediate CAGRA index and builds and writes the HNSW index.
+   * Flush/sorting path: builds a host matrix from the heap vectors, then delegates
+   * to {@link #writeFieldInternal(FieldInfo, CuVSMatrix)}.
    *
    * @param fieldInfo instance of FieldInfo that has the field description
    * @param vectors vectors to index
@@ -154,29 +158,48 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<float[]> vectors) thro
       writeSingleVectorGraph(fieldInfo, vectors);
       return;
     }
-    try {
-      CuVSMatrix dataset =
-          Utils.createFloatMatrix(
-              vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance());
+    CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension());
+    writeFieldInternal(fieldInfo, dataset);
+  }
 
+  /**
+   * Builds the intermediate CAGRA index and builds and writes the HNSW index.
+   * Single implementation used by both the flush and merge paths. The dataset is a
+   * {@link CuVSMatrix} (host-backed on the merge path) so the full set of vectors is
+   * never double-materialised on the Java heap.
+   *
+   * @param fieldInfo instance of FieldInfo that has the field description
+   * @param dataset   matrix of all vectors to index
+   * @throws IOException
+   */
+  private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException {
+    int size = (int) dataset.size();
+    if (size == 0) {
+      writeEmpty(fieldInfo, hnswMeta);
+      return;
+    }
+    if (size < 2) {
+      float[] buf = new float[fieldInfo.getVectorDimension()];
+      dataset.getRow(0).toArray(buf);
+      writeSingleVectorGraph(fieldInfo, List.of(buf));
+      return;
+    }
+    try {
       CagraIndexParams params =
           CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns());
-
       CagraIndex cagraIndex =
           CagraIndex.newBuilder(getCuVSResourcesInstance())
               .withDataset(dataset)
               .withIndexParams(params)
               .build();
       CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph();
-      int size = (int) dataset.size();
       int dimensions = fieldInfo.getVectorDimension();
       GPUBuiltHnswGraph hnswGraph =
           createMultiLayerHnswGraph(
               fieldInfo,
-              size,
               dimensions,
               adjacencyListMatrix,
-              vectors,
+              dataset,
               acceleratedHNSWParams.getHnswLayers(),
               acceleratedHNSWParams.getGraphdegree(),
               params,
@@ -276,13 +299,24 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List<float[]> vectors)
   }
 
   /**
-   * Create combined data set for the merged segment and call writeFieldInternal.
+   * Streams merged vectors directly into a native host-memory matrix (CuVSHostMatrix)
+   * without materialising a List<float[]> on the Java heap, then calls writeFieldInternal.
+   * This avoids the double-copy OOM (heap list + native matrix simultaneously) that
+   * occurs when force-merging large segments.
    */
   private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
     try {
-      List<float[]> dataset =
-          createListFromMergedVectors(
-              KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState));
+      FloatVectorValues mergedVectors =
+          KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
+      int size = mergedVectors.size();
+      int dims = fieldInfo.getVectorDimension();
+      CuVSMatrix.Builder<CuVSHostMatrix> builder =
+          CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT);
+      KnnVectorValues.DocIndexIterator it = mergedVectors.iterator();
+      for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) {
+        builder.addVector(mergedVectors.vectorValue(it.index()));
+      }
+      CuVSHostMatrix dataset = builder.build();
       writeFieldInternal(fieldInfo, dataset);
     } catch (Throwable t) {
       Utils.handleThrowable(t);

src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java

@@ -155,8 +155,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<byte[]> vectors) throw
       int dimensions = fieldInfo.getVectorDimension();
       int bytesPerVector = (dimensions + 7) / 8;
 
-      CuVSMatrix dataset =
-          Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance());
+      CuVSMatrix dataset = Utils.createByteMatrix(vectors, bytesPerVector);
 
       if (dataset.size() < 2) {
         writeSingleVectorGraph(fieldInfo, vectors);
@@ -179,10 +178,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<byte[]> vectors) throw
       GPUBuiltHnswGraph hnswGraph =
           createMultiLayerHnswGraph(
               fieldInfo,
-              size,
               dimensions,
               adjacencyListMatrix,
-              vectors,
+              dataset,
               acceleratedHNSWParams.getHnswLayers(),
               acceleratedHNSWParams.getGraphdegree(),
               params,

src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java

@@ -181,8 +181,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<?> vectors) throws IOE
       }
 
       // Create CuVSMatrix with BYTE data type (unsigned bytes)
-      CuVSMatrix dataset =
-          Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance());
+      CuVSMatrix dataset = Utils.createByteMatrix(unsignedVectors, dimensions);
 
       if (dataset.size() < 2) {
         writeSingleVectorGraph(fieldInfo, unsignedVectors);
@@ -204,10 +203,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, List<?> vectors) throws IOE
       GPUBuiltHnswGraph hnswGraph =
           createMultiLayerHnswGraph(
               fieldInfo,
-              size,
               dimensions,
               adjacencyListMatrix,
-              unsignedVectors,
+              dataset,
               acceleratedHNSWParams.getHnswLayers(),
               acceleratedHNSWParams.getGraphdegree(),
               params,