From f352930091b5b3e57868b5c8c7f034b4d706ddde Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Tue, 28 Apr 2026 21:13:52 +0000 Subject: [PATCH 1/3] cuvs-lucene-137: increase max-number of HNSW layers in CAGRA_HNSW's conversion from CAGRA to HNSW --- src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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; From baeadd1e3671bacf6eafc9c52504de829c2bfd65 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Wed, 29 Apr 2026 01:01:37 +0000 Subject: [PATCH 2/3] cuvs-lucene__139: This code allows us to construct the HNSW graph on GPU without loading the full set of data on the Java Heap, but instead allows us to stream the set of data to the Java Heap. --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 89 ++++++++++++++++- .../Lucene99AcceleratedHNSWVectorsWriter.java | 95 +++++++++++++++++-- .../java/com/nvidia/cuvs/lucene/Utils.java | 33 ++----- 3 files changed, 183 insertions(+), 34 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 2a8ab02..b518742 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -76,7 +76,10 @@ 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 + *

+ * This overload takes a {@code List} as the vector source and is used + * by the flush path where vectors are already materialised on the Java heap. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, @@ -170,6 +173,90 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } + /** + * 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 + *

+ * This overload takes a {@code CuVSMatrix} as the vector source and is used + * by the merge path. Vectors for higher-layer subsets are read directly from + * the native host matrix via {@link CuVSMatrix#getRow(long)} and + * {@link RowView#toArray(float[])}, avoiding any additional heap allocation + * of the full dataset. + */ + public static GPUBuiltHnswGraph createMultiLayerHnswGraph( + FieldInfo fieldInfo, + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + int graphDegree, + CagraIndexParams params, + QuantizationType quantization) + throws Throwable { + + int M = graphDegree / 2; + + List layerNodes = new ArrayList<>(); + List layerAdjacencies = new ArrayList<>(); + + // Layer 0: Use full CAGRA adjacency list + layerNodes.add(null); + layerAdjacencies.add(adjacencyListMatrix); + + int currentLayerSize = size; + int layerIndex = 1; + Random random = new Random(); + + while (layerIndex < hnswLayers && currentLayerSize > 1) { + int nextLayerSize = Math.max(2, currentLayerSize / M); + SortedSet selectedNodesSet = new TreeSet<>(); + + if (layerIndex == 1) { + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(random.nextInt(size)); + } + } else { + int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); + } + } + + int[] selectedNodes = + selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); + layerNodes.add(selectedNodes); + + if (quantization == QuantizationType.NONE) { + // 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++) { + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + } + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + } else { + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; + for (int i = 0; i < nextLayerSize; i++) { + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + } + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); + } + + currentLayerSize = nextLayerSize; + layerIndex++; + random = new Random(new Random().nextLong()); + } + + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } + /** * Builds a CAGRA graph for a subset of binary quantized vectors */ diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 35209ce..653bc14 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; @@ -140,6 +143,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException /** * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * Used by the flush path (operates on an in-memory List). * * @param fieldInfo instance of FieldInfo that has the field description * @param vectors vectors to index @@ -200,6 +204,74 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro } } + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * Used by the merge path (operates on a pre-built CuVSHostMatrix to avoid + * double-materializing the full dataset on heap and in native memory simultaneously). + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param dataset pre-built host-memory matrix of all merged vectors + * @param size number of vectors in the dataset + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, CuVSHostMatrix dataset, int size) + throws IOException { + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + if (size < 2) { + int dims = fieldInfo.getVectorDimension(); + float[] buf = new float[dims]; + dataset.getRow(0).toArray(buf); + writeSingleVectorGraph(fieldInfo, List.of(buf)); + return; + } + try { + CagraIndexParams params = + cagraIndexParams( + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getCagraGraphBuildAlgo(), + acceleratedHNSWParams.getCuVSIvfPqParams()); + CagraIndex cagraIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + int dimensions = fieldInfo.getVectorDimension(); + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, + size, + dimensions, + adjacencyListMatrix, + dataset, + acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), + params, + QuantizationType.NONE); + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); + cagraIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + /** * Build the indexes and writes it to the disk. */ @@ -276,14 +348,25 @@ 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)); - writeFieldInternal(fieldInfo, dataset); + 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, size); } catch (Throwable t) { Utils.handleThrowable(t); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index c55aa7c..731f9df 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -54,20 +54,12 @@ static void handleThrowable(Throwable t) throws IOException { * @return an instance of 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 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( // was: CuVSMatrix.deviceBuilder(resources, ... + data.size(), dimensions, CuVSMatrix.DataType.FLOAT); for (float[] vector : data) { builder.addVector(vector); } - return builder.build(); } @@ -84,20 +76,12 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSReso */ 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 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( // was: CuVSMatrix.deviceBuilder(resources, ... + data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } - return builder.build(); } @@ -112,13 +96,8 @@ static CuVSMatrix createByteMatrix( static CuVSMatrix createByteMatrixFromArray( byte[][] data, int bytesPerVector, CuVSResources resources) { 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( // was: CuVSMatrix.deviceBuilder(resources, ... + data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } From 9d4398219f6d1e7689860766355d4b5229d53e5b Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Thu, 2 Jul 2026 00:33:45 +0000 Subject: [PATCH 3/3] Consolidate duplicated HNSW graph and field-writing methods Address review feedback on duplicated code paths introduced by the host-streaming refactor: - Merge the two createMultiLayerHnswGraph overloads (List-based and CuVSMatrix-based) into a single CuVSMatrix-based method. The flush and quantized paths now build a matrix and sample rows from it, matching the merge path. - Merge the two writeFieldInternal overloads into one that takes a CuVSMatrix, with a thin List adapter for the flush/sorting path. - Drop the redundant `size` parameter throughout; it is derived from dataset.size(). - Remove the now-unused CuVSResources parameter from the Utils matrix builders (createFloatMatrix/createByteMatrix/createByteMatrixFromArray) and correct their docs, which still referred to device memory after the switch to hostBuilder. Updated all call sites (CuVS2510GPUVectorsWriter, binary/scalar quantized writers). While migrating the quantized writers to the unified graph method, fixed a latent byte-width bug: the quantized branch hardcoded (dimensions + 7) / 8 (binary packing), which is wrong for scalar quantization (one byte per dimension). The width is now taken from the matrix's column count, correct for both binary and scalar. Signed-off-by: Zack Meeks --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 122 ++---------------- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 6 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 75 ++--------- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 6 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 6 +- .../java/com/nvidia/cuvs/lucene/Utils.java | 41 +++--- 6 files changed, 44 insertions(+), 212 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index b518742..625464f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -78,116 +78,14 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * Each layer contains 1/M nodes from the previous layer * Creates layers until the highest layer has <= M nodes *

- * This overload takes a {@code List} as the vector source and is used - * by the flush path where vectors are already materialised on the Java heap. + * 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, - int hnswLayers, - int graphDegree, - CagraIndexParams params, - QuantizationType quantization) - throws Throwable { - - // Calculate M as cagraGraphDegree/2 - 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 - 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 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]); - } - } - - // 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][]; - for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (float[]) vectors.get(selectedNodes[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][]; - for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (byte[]) vectors.get(selectedNodes[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); - } - - /** - * 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 - *

- * This overload takes a {@code CuVSMatrix} as the vector source and is used - * by the merge path. Vectors for higher-layer subsets are read directly from - * the native host matrix via {@link CuVSMatrix#getRow(long)} and - * {@link RowView#toArray(float[])}, avoiding any additional heap allocation - * of the full dataset. - */ - public static GPUBuiltHnswGraph createMultiLayerHnswGraph( - FieldInfo fieldInfo, - int size, int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, @@ -197,6 +95,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { + int size = (int) vectorDataset.size(); int M = graphDegree / 2; List layerNodes = new ArrayList<>(); @@ -239,7 +138,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); } else { - int bytesPerVector = (dimensions + 7) / 8; + // 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++) { vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); @@ -272,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 653bc14..485f147 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -142,8 +142,8 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException } /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index. - * Used by the flush path (operates on an in-memory List). + * 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 @@ -158,83 +158,35 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeSingleVectorGraph(fieldInfo, vectors); return; } - try { - CuVSMatrix dataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); - - 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, - acceleratedHNSWParams.getHnswLayers(), - acceleratedHNSWParams.getGraphdegree(), - params, - QuantizationType.NONE); - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets, - acceleratedHNSWParams.getGraphdegree()); - cagraIndex.close(); - } catch (Throwable t) { - Utils.handleThrowable(t); - } + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeFieldInternal(fieldInfo, dataset); } /** * Builds the intermediate CAGRA index and builds and writes the HNSW index. - * Used by the merge path (operates on a pre-built CuVSHostMatrix to avoid - * double-materializing the full dataset on heap and in native memory simultaneously). + * 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 pre-built host-memory matrix of all merged vectors - * @param size number of vectors in the dataset + * @param dataset matrix of all vectors to index * @throws IOException */ - private void writeFieldInternal(FieldInfo fieldInfo, CuVSHostMatrix dataset, int size) - 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) { - int dims = fieldInfo.getVectorDimension(); - float[] buf = new float[dims]; + float[] buf = new float[fieldInfo.getVectorDimension()]; dataset.getRow(0).toArray(buf); writeSingleVectorGraph(fieldInfo, List.of(buf)); return; } try { CagraIndexParams params = - cagraIndexParams( - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getCagraGraphBuildAlgo(), - acceleratedHNSWParams.getCuVSIvfPqParams()); + CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) @@ -245,7 +197,6 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSHostMatrix dataset, int GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, dataset, @@ -366,7 +317,7 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws builder.addVector(mergedVectors.vectorValue(it.index())); } CuVSHostMatrix dataset = builder.build(); - writeFieldInternal(fieldInfo, dataset, size); + 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 731f9df..c8e8070 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -43,20 +43,18 @@ 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) { + static CuVSMatrix createFloatMatrix(List data, int dimensions) { CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder( // was: CuVSMatrix.deviceBuilder(resources, ... - data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); for (float[] vector : data) { builder.addVector(vector); } @@ -64,21 +62,15 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSReso } /** - * 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) { + static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder( // was: CuVSMatrix.deviceBuilder(resources, ... - data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } @@ -86,18 +78,15 @@ static CuVSMatrix createByteMatrix( } /** - * 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.hostBuilder( // was: CuVSMatrix.deviceBuilder(resources, ... - data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); }