diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index dffb500..d2df4e9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 2a8ab02..625464f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -76,29 +76,33 @@ 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 + *

+ * 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 layerNodes = new ArrayList<>(); List 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; @@ -106,67 +110,50 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( 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 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); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index e8fb304..717f3ce 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -201,8 +201,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List 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 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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 35209ce..485f147 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/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 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 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 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 dataset = - createListFromMergedVectors( - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + int size = mergedVectors.size(); + int dims = fieldInfo.getVectorDimension(); + CuVSMatrix.Builder 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); diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index d8a98cc..a51cb13 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -155,8 +155,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List 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 vectors) throw GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - vectors, + dataset, acceleratedHNSWParams.getHnswLayers(), acceleratedHNSWParams.getGraphdegree(), params, diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 21b4be3..dd40f21 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/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, diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index c55aa7c..c8e8070 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -43,82 +43,50 @@ static void handleThrowable(Throwable t) throws IOException { } /** - * A method to build a CuVSMatrix from a list of float vectors. + * Builds a host-memory CuVSMatrix from a list of float vectors. * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + *

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, + * without creating an intermediate {@code float[][]} on the heap. * * @param data The float vectors - * @param dimensions The number float elements in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix + * @param dimensions The number of float elements in each vector + * @return a host-memory CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { - // Use Builder pattern to avoid intermediate float[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createFloatMatrix(List data, int dimensions) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - dimensions, // columns (vector dimension) - CuVSMatrix.DataType.FLOAT); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); for (float[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). - * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). * * @param data The byte vectors (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrix( - List data, int bytesPerVector, CuVSResources resources) { - // Use Builder pattern to avoid intermediate byte[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). + * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). * * @param data The 2D byte array (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrixFromArray( - byte[][] data, int bytesPerVector, CuVSResources resources) { + static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.length, // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); }