From c579e871b20bd123ffb8cc0eff9b5138d48c280f Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Wed, 29 Apr 2026 01:01:37 +0000 Subject: [PATCH 01/20] 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/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 9c49c07fe0..a39cbd49be 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -77,7 +77,10 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree * (its column count). Ceil is used to accommodate odd graph degrees. * 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, @@ -169,6 +172,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/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 13edc64975..47a6028d61 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/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 @@ -198,6 +202,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. */ @@ -273,14 +345,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/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index e4a20d2b4d..3cd608e64b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/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 8598e4fd13303d1afb9ce68e15ef0bb205bc800e Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Thu, 2 Jul 2026 00:33:45 +0000 Subject: [PATCH 02/20] 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 | 120 ++---------------- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 6 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 73 ++--------- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 6 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 6 +- .../java/com/nvidia/cuvs/lucene/Utils.java | 41 +++--- 6 files changed, 44 insertions(+), 208 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index a39cbd49be..6005575034 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -79,114 +79,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, - CagraIndexParams params, - QuantizationType quantization) - throws Throwable { - - int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 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, @@ -196,6 +96,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { + int size = (int) vectorDataset.size(); int M = graphDegree / 2; List layerNodes = new ArrayList<>(); @@ -238,7 +139,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]); @@ -271,11 +173,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/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 97fedf924f..ce77f36dea 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/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/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 47a6028d61..7bca35d5f0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/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,81 +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(), - 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); - 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) @@ -243,7 +197,6 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSHostMatrix dataset, int GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, dataset, @@ -363,7 +316,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/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 87907d2cbb..10380ca538 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/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(), params, QuantizationType.BINARY); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 7141af56ee..f9dbbfb06d 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/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(), params, QuantizationType.SCALAR); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 3cd608e64b..0fa96ded3e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/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); } From f66bdcefcf4ada4b9ea5ca399f616cfff2b94884 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 14 Jul 2026 14:35:28 -0700 Subject: [PATCH 03/20] Add example: chunked sequential ingestion of large .fbin files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ChunkedFbinIngestExample, a reference for streaming a large vector file into an accelerated HNSW index without the two common pitfalls: - fd churn: reopening/seeking/closing the source file per vector, which is far slower than sequential I/O. - doubled memory: pre-loading the whole file onto the JVM heap on top of Lucene's own per-segment buffer. The nested ChunkedFbinReader opens the file once and serves vectors from a reusable buffer refilled by large sequential reads, holding at most one chunk. Vectors are streamed one at a time into addDocument to build a single-segment index. Because cuvs-lucene is a codec (below addDocument), how the source is read is application code — the example is a template to adapt to any vector source (DB, object store, stream). Runnable with no arguments (generates a small demo .fbin) or against a real file with a configurable chunk size. README updated with a run entry. --- java/cuvs-lucene/examples/README.md | 11 + .../examples/ChunkedFbinIngestExample.java | 272 ++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java diff --git a/java/cuvs-lucene/examples/README.md b/java/cuvs-lucene/examples/README.md index 19013675b6..1535d644ba 100644 --- a/java/cuvs-lucene/examples/README.md +++ b/java/cuvs-lucene/examples/README.md @@ -32,3 +32,14 @@ To run the Index and Search on GPU example do: ```sh mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample ``` + +To run the chunked `.fbin` ingestion example (reference pattern for streaming a large vector file +into an accelerated HNSW index without per-vector file reopening or holding the whole file in +memory) do: + +```sh +mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.08.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.ChunkedFbinIngestExample +``` + +With no arguments it generates and indexes a small demo `.fbin`; pass a real file and chunk size as +`... ChunkedFbinIngestExample `. diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java new file mode 100644 index 0000000000..a3ce8afad5 --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java @@ -0,0 +1,272 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** + * Reference pattern for ingesting a LARGE {@code .fbin} vector file into an accelerated HNSW index + * without the two pitfalls that dominate build time and memory when loading big vector files: + * + *

    + *
  1. fd churn — reopening/seeking/closing the file per vector is orders of magnitude + * slower than sequential reads. {@link ChunkedFbinReader} opens the file ONCE and reads it + * front-to-back in large sequential chunks. + *
  2. unbounded / doubled memory — pre-loading the whole file into the JVM heap holds a + * redundant copy on top of Lucene's own per-segment buffer (~2x peak). The chunked reader + * holds at most one chunk, streaming each vector straight into {@code addDocument}. + *
+ * + *

cuvs-lucene is a Lucene codec and sits below {@code addDocument}, so how you read your source + * data is application code — adapt {@link ChunkedFbinReader} to your own source (a DB, object + * store, or stream). The properties that matter are: open once, read sequentially, bound memory + * to a chunk. + * + *

Usage: {@code ChunkedFbinIngestExample [] []}. With no arguments a + * small demo {@code .fbin} is generated and indexed. + */ +public class ChunkedFbinIngestExample { + + private static final Logger log = Logger.getLogger(ChunkedFbinIngestExample.class.getName()); + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + public static void main(String[] args) throws Exception { + int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; + Path indexDirPath = Paths.get(UUID.randomUUID().toString()); + + Path fbinPath; + boolean generated = false; + if (args.length >= 1) { + fbinPath = Paths.get(args[0]); + } else { + fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); + writeDemoFbin(fbinPath, 5000, 32, new Random(222)); + generated = true; + log.info("No .fbin provided; generated a demo file at " + fbinPath); + } + + try { + buildIndex(fbinPath, indexDirPath, chunkSizeMB); + runSampleSearch(indexDirPath, fbinPath, 5); + } finally { + FileUtils.deleteDirectory(indexDirPath.toFile()); + if (generated) { + Files.deleteIfExists(fbinPath); + } + } + } + + /** + * Builds a single-segment accelerated HNSW index, streaming vectors from the {@code .fbin} in + * sequential chunks (never holding more than one chunk in memory). + */ + private static void buildIndex(Path fbinPath, Path indexDirPath, int chunkSizeMB) + throws Exception { + AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + try (ChunkedFbinReader reader = new ChunkedFbinReader(fbinPath, chunkSizeMB)) { + int n = reader.size(); + + // Keep the whole dataset in ONE segment: disable RAM-based flushing and raise the doc-count + // flush threshold above the document count so nothing flushes before commit. (This only + // controls Lucene's segment cadence; the vectors are not held here — they are streamed one + // at a time from the chunked reader into addDocument.) + // Order matters: enable the doc-count flush trigger BEFORE disabling the RAM trigger, since + // Lucene rejects a config where both are disabled at once. + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(Math.max(2, n + 1)) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + log.info( + "Indexing " + + n + + " vectors (" + + reader.dimension() + + "-dim) from " + + fbinPath + + " using " + + chunkSizeMB + + " MB sequential chunks"); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < n; i++) { + float[] vector = reader.get(i); // sequential access -> served from the current chunk + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); // single flush -> single segment; the GPU CAGRA build happens here + } + log.info("Index build complete: " + indexDirPath); + } + } + + /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ + private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { + float[] queryVector; + try (ChunkedFbinReader reader = new ChunkedFbinReader(fbinPath, 1)) { + queryVector = reader.get(0); + } + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + log.info("Sample search returned " + results.scoreDocs.length + " hits:"); + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc sd = results.scoreDocs[i]; + String id = searcher.storedFields().document(sd.doc).get(ID_FIELD); + log.info(" rank " + (i + 1) + ": id=" + id + " score=" + sd.score); + } + } + } + + /** Writes a small random {@code .fbin} so the example is runnable without external data. */ + private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) + throws IOException { + ByteBuffer buf = + ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] + buf.putInt(dim); + for (int i = 0; i < numVectors; i++) { + for (int j = 0; j < dim; j++) { + buf.putFloat(random.nextFloat() * 100); + } + } + buf.flip(); + try (FileChannel ch = + FileChannel.open( + path, + StandardOpenOption.CREATE, + StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING)) { + while (buf.hasRemaining()) { + ch.write(buf); + } + } + } + + /** + * Chunked, sequential reader for uncompressed {@code .fbin} files + * ({@code [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). + * + *

Opens the file ONCE and serves {@link #get(int)} from a reusable buffer that is refilled + * with a single large sequential read whenever the requested index leaves the current chunk. For + * sequential access (index = 0, 1, 2, ...) it reads the file front-to-back in {@code size/chunk} + * bulk reads while holding only one chunk in memory — the opposite of reopening the file per + * vector. + */ + static final class ChunkedFbinReader implements AutoCloseable { + + private static final long HEADER_BYTES = 8; + + private final FileChannel channel; + private final int dimension; + private final int vectorCount; + private final int vectorBytes; + private final int chunkVectors; + private final ByteBuffer chunkBuffer; + + private long chunkStart = -1; // first vector index currently buffered + private int chunkLen = 0; // number of vectors currently buffered + + ChunkedFbinReader(Path path, int chunkSizeMB) throws IOException { + this.channel = FileChannel.open(path, StandardOpenOption.READ); + ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + readFully(header, 0); + header.flip(); + this.vectorCount = header.getInt(); + this.dimension = header.getInt(); + this.vectorBytes = dimension * Float.BYTES; + + long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; + int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int + this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); + this.chunkBuffer = + ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN); + } + + int size() { + return vectorCount; + } + + int dimension() { + return dimension; + } + + float[] get(int index) throws IOException { + if (index < 0 || index >= vectorCount) { + throw new IndexOutOfBoundsException( + "Index " + index + " out of bounds [0, " + vectorCount + ")"); + } + if (chunkStart < 0 || index < chunkStart || index >= chunkStart + chunkLen) { + long start = (index / (long) chunkVectors) * chunkVectors; + int toRead = (int) Math.min(chunkVectors, vectorCount - start); + chunkBuffer.clear(); + chunkBuffer.limit(toRead * vectorBytes); + readFully(chunkBuffer, HEADER_BYTES + start * (long) vectorBytes); + chunkStart = start; + chunkLen = toRead; + } + int base = (int) (index - chunkStart) * vectorBytes; + float[] vector = new float[dimension]; + for (int i = 0; i < dimension; i++) { + vector[i] = chunkBuffer.getFloat(base + i * Float.BYTES); + } + return vector; + } + + private void readFully(ByteBuffer buf, long position) throws IOException { + long pos = position; + while (buf.hasRemaining()) { + int n = channel.read(buf, pos); + if (n < 0) { + throw new IOException("Unexpected EOF reading at position " + pos); + } + pos += n; + } + } + + @Override + public void close() throws IOException { + channel.close(); + } + } +} From d0b5eb9aa16b7b4e2418331fab668271e29a5021 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 14 Jul 2026 15:06:39 -0700 Subject: [PATCH 04/20] Parallelize level-0 HNSW graph serialization in writeGraph Level 0 (all N nodes) is now encoded in parallel: within memory-bounded waves, threads delta/VInt-encode contiguous node sub-ranges into per-thread buffers, which are concatenated to the IndexOutput in node order. The per-node encoding is independent and shared with the serial path, so the on-disk bytes are identical. Higher levels stay serial. Level-0 neighbors are read from the already-materialized layer0Neighbors array, so parallel reads are safe. Thread count comes from AcceleratedHNSWParams.getWriterThreads() (default 1 -> unchanged behavior; opt-in), and applies to CAGRA_HNSW and the quantized variants (shared writeGraph). Measured on deep1b-10M (degree 88): write-graph 36 s -> 12.8 s at 4 threads, 8.1 s at 16 (single-threaded concat is the ~5-6 s floor). Recall unchanged (within GPU-build run-to-run noise). --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 176 ++++++++++++++---- .../Lucene99AcceleratedHNSWVectorsWriter.java | 4 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 4 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 4 +- 4 files changed, 141 insertions(+), 47 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 6005575034..350ae2ad97 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -19,8 +19,14 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.ByteBuffersDataOutput; +import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -222,56 +228,144 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) - throws IOException { - // write vectors' neighbors on each level into the vectorIndex file + public static int[][] writeGraph( + GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) throws IOException { int countOnLevel0 = graph.size(); - int[][] offsets = new int[graph.numLevels()][]; - int[] scratch = new int[graph.maxConn() * 2]; - for (int level = 0; level < graph.numLevels(); level++) { + int numLevels = graph.numLevels(); + int[][] offsets = new int[numLevels][]; + + // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is + // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in + // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk bytes + // are identical to the fully-serial path (blocks in node order, offsets = per-node byte lengths). + int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); + offsets[0] = new int[level0Nodes.length]; + if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { + writeLevel0Parallel(graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, numThreads); + } else { + writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0); + } + + for (int level = 1; level < numLevels; level++) { int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; - int nodeOffsetId = 0; - - for (int node : sortedNodes) { - // Get node neighbors - NeighborArray neighbors = graph.getNeighbors(level, node); - // Get the size of the neighbor array - int size = neighbors.size(); - // Write size in VInt as the neighbors list is typically small - long offsetStart = vectorIndex.getFilePointer(); - // Get neighbors - int[] nnodes = neighbors.nodes(); - // Sort them - Arrays.sort(nnodes, 0, size); - // Now that we have sorted, do delta encoding to minimize the required bits to store the - // information - int actualSize = 0; - if (size > 0) { - scratch[0] = nnodes[0]; - actualSize = 1; - } - // De-duplication - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - // Sorting step helps here - if (nnodes[i - 1] == nnodes[i]) { + writeLevelSerial(graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0); + } + return offsets; + } + + /** Node count below which parallel level-0 serialization is not worth the overhead. */ + private static final int PARALLEL_MIN_NODES = 1 << 16; + + /** Nodes per wave — bounds the transient encode buffer regardless of dataset size. */ + private static final int WAVE_NODES = 1 << 20; + + /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ + private static void writeLevelSerial( + GPUBuiltHnswGraph graph, + IndexOutput out, + int level, + int[] sortedNodes, + int[] offsets, + int countOnLevel0) + throws IOException { + int[] scratch = new int[graph.maxConn() * 2]; + int idx = 0; + for (int node : sortedNodes) { + long start = out.getFilePointer(); + encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); + offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); + } + } + + /** + * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node + * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order + * (identical layout to the serial path). + */ + private static void writeLevel0Parallel( + GPUBuiltHnswGraph graph, + IndexOutput out, + int[] nodes, + int[] offsets, + int countOnLevel0, + int numThreads) + throws IOException { + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + try { + int n = nodes.length; + for (int waveStart = 0; waveStart < n; waveStart += WAVE_NODES) { + int waveEnd = Math.min(waveStart + WAVE_NODES, n); + int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; + + ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; + List> futures = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int subStart = waveStart + t * perThread; + final int subEnd = Math.min(subStart + perThread, waveEnd); + final int slot = t; + if (subStart >= subEnd) { continue; } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + futures.add( + pool.submit( + () -> { + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + int[] scratch = new int[graph.maxConn() * 2]; + for (int i = subStart; i < subEnd; i++) { + long before = buffer.size(); + encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); + offsets[i] = Math.toIntExact(buffer.size() - before); + } + buffers[slot] = buffer; + return null; + })); + } + for (Future f : futures) { + f.get(); } - // Write the size after duplicates are removed - vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); + // Concatenate in thread order (== node order), preserving the serial byte layout. + for (ByteBuffersDataOutput buffer : buffers) { + if (buffer != null) { + buffer.copyTo(out); + } } - offsets[level][nodeOffsetId++] = - Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted during parallel writeGraph", e); + } catch (ExecutionException e) { + throw new IOException("Parallel writeGraph failed", e.getCause()); + } finally { + pool.shutdown(); + } + } + + /** + * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt + * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. + */ + private static void encodeNode( + NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) throws IOException { + int size = neighbors == null ? 0 : neighbors.size(); + int actualSize = 0; + if (size > 0) { + int[] nnodes = neighbors.nodes(); + Arrays.sort(nnodes, 0, size); + scratch[0] = nnodes[0]; + actualSize = 1; + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + } + } + out.writeVInt(actualSize); + for (int i = 0; i < actualSize; i++) { + out.writeVInt(scratch[i]); } - // Return offsets (information written while writing the meta info) - return offsets; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 7bca35d5f0..e0652c5d82 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -205,7 +205,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws params, QuantizationType.NONE); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -281,7 +281,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 10380ca538..00c5f1d4cb 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -187,7 +187,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -275,7 +275,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index f9dbbfb06d..62b67657ed 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -213,7 +213,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -300,7 +300,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata From cd234900866f2e42ecd99f8ccf23971084bb6ea1 Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 15 Jul 2026 13:00:02 -0700 Subject: [PATCH 05/20] Add native flat buffering for single-segment CAGRA_HNSW builds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stream vectors directly into a native host matrix during indexing instead of buffering them as a heap List, so the full dataset is no longer held twice (Lucene's flat-writer list + the native CAGRA build input). This takes peak host memory from ~2x to ~1x, fixing OOM on out-of-core builds, and removes the per-vector createFloatMatrix copy (the matrix-assembly stage). Opt-in via AcceleratedHNSWParams.numInputVectors (default 0 = disabled → existing heap-buffered path unchanged). When set: - FieldWriter streams addValue() into a CuVSMatrix.hostBuilder host matrix (reusing the same builder the merge path already uses) and tracks its own DocsWithFieldSet. - NativeFlatVectorsWriter writes .vec/.vemf directly from that host matrix, byte-compatible with Lucene99FlatVectorsWriter (10.2.0) so the stock Lucene99FlatVectorsReader reads them; format constants are re-declared with a version-pinned attribution and an upgrade-check note. - The writer/format wire the native path into flush and skip the Lucene flat writer so it never opens the same .vec/.vemf. The host matrix is sized for exactly numInputVectors rows, so the value must equal the number of vectors added; the writer fails fast otherwise. Supported only for the unsorted single-segment flush path — merges and index-sorted segments throw. Validated on deep1b-10M: recall parity with the heap path (99.588 vs 99.594) confirms format correctness, and the matrix-assembly stage is eliminated. --- .../cuvs/lucene/AcceleratedHNSWParams.java | 44 +++- .../com/nvidia/cuvs/lucene/FieldWriter.java | 100 ++++++++- .../Lucene99AcceleratedHNSWVectorsFormat.java | 5 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 79 +++++++- .../cuvs/lucene/NativeFlatVectorsWriter.java | 190 ++++++++++++++++++ 5 files changed, 410 insertions(+), 8 deletions(-) create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index a5f164b70b..2d0cc64a97 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -66,6 +66,7 @@ public static enum Strategy { public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; public static final HnswHeuristicType DEFAULT_HNSW_HEURISTIC_TYPE = HnswHeuristicType.SAME_GRAPH_FOOTPRINT; + public static final int DEFAULT_NUM_INPUT_VECTORS = 0; public static final Supplier DEFAULT_IVF_PQ_PARAMS = () -> { @@ -91,6 +92,7 @@ public static enum Strategy { private final CuvsDistanceType cuvsDistanceType; private final int nnDescentNumIterations; private final HnswHeuristicType hnswHeuristicType; + private final int numInputVectors; /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. @@ -109,6 +111,7 @@ public static enum Strategy { * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. * @param hnswHeuristicType the heuristic cuVS applies when deriving the CAGRA build parameters from maxConn and beamWidth under the HEURISTIC strategy. + * @param numInputVectors exact number of vectors to be indexed, used to pre-size the native flat buffer (0 = disabled). */ private AcceleratedHNSWParams( int writerThreads, @@ -124,7 +127,8 @@ private AcceleratedHNSWParams( Strategy strategy, CuvsDistanceType cuvsDistanceType, int nnDescentNumIterations, - HnswHeuristicType hnswHeuristicType) { + HnswHeuristicType hnswHeuristicType, + int numInputVectors) { super(); this.writerThreads = writerThreads; this.intermediateGraphDegree = intermediateGraphDegree; @@ -140,6 +144,7 @@ private AcceleratedHNSWParams( this.cuvsDistanceType = cuvsDistanceType; this.nnDescentNumIterations = nnDescentNumIterations; this.hnswHeuristicType = hnswHeuristicType; + this.numInputVectors = numInputVectors; } /** @@ -272,6 +277,17 @@ public HnswHeuristicType getHnswHeuristicType() { return hnswHeuristicType; } + /** + * Get the number of input vectors used to pre-size the native flat buffer. A value of + * {@value DEFAULT_NUM_INPUT_VECTORS} means unset (the writer uses the default heap-buffered + * flat path). + * + * @return the number of vectors to be indexed, or 0 if unset + */ + public int getNumInputVectors() { + return numInputVectors; + } + @Override public String toString() { return "AcceleratedHNSWParams [writerThreads=" @@ -302,6 +318,8 @@ public String toString() { + nnDescentNumIterations + ", hnswHeuristicType=" + hnswHeuristicType + + ", numInputVectors=" + + numInputVectors + "]"; } @@ -324,6 +342,7 @@ public static class Builder { private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE; private int nnDescentNumIterations = DEFAULT_NN_DESCENT_NUM_ITERATIONS; private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE; + private int numInputVectors = DEFAULT_NUM_INPUT_VECTORS; /** * Set the number of cuVS writer threads while building the index @@ -507,6 +526,23 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) { return this; } + /** + * Set the exact number of vectors to be indexed, used to pre-allocate a single contiguous + * native flat buffer (avoiding the on-heap {@code List} and the extra host-matrix + * copy). The native buffer is sized for exactly this many rows, so the value MUST equal the + * number of vectors actually added; the writer fails fast otherwise. Only supported for the + * unsorted single-segment CAGRA_HNSW build (no merges). A value of + * {@value DEFAULT_NUM_INPUT_VECTORS} (the default) disables it and uses the default + * heap-buffered flat path. + * + * @param numInputVectors the exact number of vectors to be indexed, or 0 to disable + * @return instance of {@link Builder} + */ + public Builder withNumInputVectors(int numInputVectors) { + this.numInputVectors = numInputVectors; + return this; + } + /** * Validates the input parameters. * @@ -591,6 +627,9 @@ private void validate() throws IllegalArgumentException { + MAX_NN_DESCENT_NUM_ITERATIONS + "]"); } + if (numInputVectors < 0) { + throw new IllegalArgumentException("numInputVectors cannot be negative."); + } } /** @@ -620,7 +659,8 @@ public AcceleratedHNSWParams build() { strategy, cuvsDistanceType, nnDescentNumIterations, - hnswHeuristicType); + hnswHeuristicType, + numInputVectors); } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 8e2d9a70e2..64f75766ea 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -8,6 +8,8 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; import java.util.List; @@ -23,18 +25,57 @@ public class FieldWriter extends KnnFieldVectorsWriter { RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); private final FieldInfo fieldInfo; + private final int dimension; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; private QuantizationType quantizationType; + /** + * Native-buffering state. When {@code numInputVectors > 0} the writer streams incoming vectors + * directly into a native host matrix ({@link CuVSMatrix#hostBuilder}) instead of accumulating a + * {@code List} on the Java heap via {@link #flatFieldVectorsWriter}. That matrix is reused + * as the CAGRA build input, so the full dataset is never held twice (heap list + native copy). + * + *

The matrix is preallocated for exactly {@code numInputVectors} rows, so the hint must equal + * the number of vectors actually added (validated at build time in the caller). + */ + private final int numInputVectors; + + private final boolean nativeBuffering; + private final CuVSMatrix.Builder hostMatrixBuilder; + private final DocsWithFieldSet nativeDocsWithField; + private int nativeCount; + private CuVSHostMatrix builtMatrix; + @SuppressWarnings("unchecked") public FieldWriter( QuantizationType quantizationType, FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + this(quantizationType, fieldInfo, flatFieldVectorsWriter, 0); + } + + @SuppressWarnings("unchecked") + public FieldWriter( + QuantizationType quantizationType, + FieldInfo fieldInfo, + FlatFieldVectorsWriter flatFieldVectorsWriter, + int numInputVectors) { this.quantizationType = quantizationType; this.fieldInfo = fieldInfo; + this.dimension = fieldInfo.getVectorDimension(); this.flatFieldVectorsWriter = (FlatFieldVectorsWriter) flatFieldVectorsWriter; + this.numInputVectors = numInputVectors; + this.nativeBuffering = numInputVectors > 0; + if (nativeBuffering) { + // Preallocates one contiguous native region of numInputVectors * dimension * 4 bytes. + this.hostMatrixBuilder = + CuVSMatrix.hostBuilder(numInputVectors, dimension, CuVSMatrix.DataType.FLOAT); + this.nativeDocsWithField = new DocsWithFieldSet(); + } else { + this.hostMatrixBuilder = null; + this.nativeDocsWithField = null; + } } @Override @@ -46,7 +87,25 @@ public void addValue(int docID, Object vectorValue) throws IOException { + "\" appears more than once in this document (only one value is allowed per" + " field)"); } - flatFieldVectorsWriter.addValue(docID, (float[]) vectorValue); + if (nativeBuffering) { + if (nativeCount >= numInputVectors) { + throw new IllegalStateException( + "Buffered vectors (" + + (nativeCount + 1) + + ") exceed numInputVectors (" + + numInputVectors + + ") for field \"" + + fieldInfo.name + + "\""); + } + // hostMatrixBuilder.addVector validates the dimension and performs the native row copy. + hostMatrixBuilder.addVector((float[]) vectorValue); + nativeDocsWithField.add(docID); + nativeCount++; + lastDocID = docID; + } else { + flatFieldVectorsWriter.addValue(docID, (float[]) vectorValue); + } } List getByteVectors() { @@ -68,7 +127,41 @@ FieldInfo fieldInfo() { } DocsWithFieldSet getDocsWithFieldSet() { - return flatFieldVectorsWriter.getDocsWithFieldSet(); + return nativeBuffering ? nativeDocsWithField : flatFieldVectorsWriter.getDocsWithFieldSet(); + } + + /** Whether this writer streams vectors into a native host matrix (hint path). */ + boolean isNativeBuffering() { + return nativeBuffering; + } + + /** + * The native host matrix holding the buffered vectors. Valid only when {@link #isNativeBuffering()} + * is true. The matrix is built once and cached; the caller owns closing it via + * {@link #releaseNativeBuffer()} once the CAGRA build has consumed it. + */ + CuVSHostMatrix getHostMatrix() { + if (builtMatrix == null) { + builtMatrix = hostMatrixBuilder.build(); + } + return builtMatrix; + } + + /** Number of vectors buffered so far (hint path). */ + int getNativeVectorCount() { + return nativeCount; + } + + int dimension() { + return dimension; + } + + /** Closes the native host matrix. Safe to call multiple times; a no-op in the non-native path. */ + void releaseNativeBuffer() { + if (nativeBuffering) { + getHostMatrix().close(); + builtMatrix = null; + } } @Override @@ -78,6 +171,7 @@ public Object copyValue(Object vectorValue) { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed(); + // The native host matrix is off-heap and intentionally excluded from Lucene's heap RAM accounting. + return SHALLOW_SIZE + (nativeBuffering ? 0 : flatFieldVectorsWriter.ramBytesUsed()); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 61c9b41c3a..050f626657 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -74,7 +74,10 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + // In hint mode the accelerated writer owns the flat .vec/.vemf files, so the Lucene flat writer + // must not be created (it would open the same outputs). The fallback path below still needs it. + boolean nativeMode = isSupported() && acceleratedHNSWParams.getNumInputVectors() > 0; + var flatWriter = nativeMode ? null : FLAT_VECTORS_FORMAT.fieldsWriter(state); if (isSupported()) { log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index e0652c5d82..3236f6b1e5 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -64,6 +64,18 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + + /** + * Hint-path state. When {@code numInputVectors > 0}, vectors are streamed into a native host + * matrix (see {@link FieldWriter}) rather than a heap {@code List}, and the flat + * {@code .vec}/{@code .vemf} files are written by {@link #nativeFlat} instead of by + * {@link #flatVectorsWriter} (which is {@code null} in this mode). Supports only the unsorted + * single-segment flush path; merges and index-sorted flushes are rejected. + */ + private final int numInputVectors; + + private final boolean nativeMode; + private final NativeFlatVectorsWriter nativeFlat; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -96,6 +108,8 @@ public Lucene99AcceleratedHNSWVectorsWriter( this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; this.acceleratedHNSWParams = acceleratedHNSWParams; + this.numInputVectors = acceleratedHNSWParams.getNumInputVectors(); + this.nativeMode = numInputVectors > 0; vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); @@ -117,6 +131,9 @@ public Lucene99AcceleratedHNSWVectorsWriter( VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + // In hint mode we own the flat files; the Lucene flat writer must be absent to avoid opening + // the same .vec/.vemf outputs. + nativeFlat = nativeMode ? new NativeFlatVectorsWriter(state) : null; success = true; printInfoStream(infoStream, COMPONENT, "Lucene99AcceleratedHNSWVectorsWriter is initialized"); } finally { @@ -135,6 +152,14 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException if (encoding != FLOAT32) { throw new IllegalArgumentException("Expected float32, got:" + encoding); } + if (nativeMode) { + // Buffer directly into a native host matrix; return the FieldWriter itself so Lucene routes + // addValue() here rather than to a (nonexistent) Lucene flat field writer. + var cuvsFieldWriter = + new FieldWriter(QuantizationType.NONE, fieldInfo, null, numInputVectors); + fields.add(cuvsFieldWriter); + return cuvsFieldWriter; + } var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); var cuvsFieldWriter = new FieldWriter(QuantizationType.NONE, fieldInfo, writer); fields.add(cuvsFieldWriter); @@ -228,6 +253,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws */ @Override public void flush(int maxDoc, DocMap sortMap) throws IOException { + if (nativeMode) { + if (sortMap != null) { + throw new UnsupportedOperationException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" + + " index-sorted segments; unset it (0) to enable the sorted flush path"); + } + for (var field : fields) { + writeFieldNative(field, maxDoc); + } + return; + } flatVectorsWriter.flush(maxDoc, sortMap); for (var field : fields) { if (sortMap == null) { @@ -238,6 +274,36 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { } } + /** + * Hint-path flush for a single field: writes the flat {@code .vec}/{@code .vemf} from the native + * host matrix, builds the CAGRA/HNSW graph from the same matrix, then releases the matrix. Both + * consumers read the matrix before it is closed. + */ + private void writeFieldNative(FieldWriter fieldData, int maxDoc) throws IOException { + int count = fieldData.getNativeVectorCount(); + if (count != numInputVectors) { + throw new IllegalStateException( + "numInputVectors (" + + numInputVectors + + ") must equal the number of vectors added (" + + count + + ") for field \"" + + fieldData.fieldInfo().name + + "\"; the native host matrix is sized for the hint exactly"); + } + FieldInfo fieldInfo = fieldData.fieldInfo(); + try { + CuVSHostMatrix dataset = fieldData.getHostMatrix(); + long ts = StageTimers.start(); + nativeFlat.writeField(fieldInfo, dataset, maxDoc, fieldData.getDocsWithFieldSet()); + StageTimers.stop( + "flat-write [DISK]", ts, (long) count * fieldInfo.getVectorDimension() * Float.BYTES); + writeFieldInternal(fieldInfo, dataset); + } finally { + fieldData.releaseNativeBuffer(); + } + } + /** * Builds the index and writes it to the disk. * @@ -327,6 +393,11 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws */ @Override public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + if (nativeMode) { + throw new UnsupportedOperationException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) supports only the" + + " unsorted single-segment flush path; unset it (0) to enable merges"); + } flatVectorsWriter.mergeOneField(fieldInfo, mergeState); vectorBasedMerge(fieldInfo, mergeState); } @@ -340,7 +411,11 @@ public void finish() throws IOException { throw new IllegalStateException("already finished"); } finished = true; - flatVectorsWriter.finish(); + if (nativeMode) { + nativeFlat.finish(); + } else { + flatVectorsWriter.finish(); + } if (hnswMeta != null) { // write end of fields marker hnswMeta.writeInt(-1); @@ -357,7 +432,7 @@ public void finish() throws IOException { @Override public void close() throws IOException { printInfoStream(infoStream, COMPONENT, "Closing resources"); - IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter, nativeFlat); closeCuVSResourcesInstance(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java new file mode 100644 index 0000000000..31d1d67315 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java @@ -0,0 +1,190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSHostMatrix; +import java.io.Closeable; +import java.io.IOException; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.ByteOrder; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; + +/** + * Writes the flat vector files ({@code .vec} data + {@code .vemf} meta) directly from a native host + * matrix, byte-for-byte compatible with Lucene's {@code Lucene99FlatVectorsWriter} so the stock + * {@code Lucene99FlatVectorsReader} can read them. + * + *

This is the hint-path counterpart to delegating to Lucene's {@code FlatVectorsWriter}: the + * accelerated writer streams vectors into a {@link CuVSHostMatrix} during indexing (see + * {@link FieldWriter}) and never materialises the full dataset as a {@code List} on the + * Java heap, so the flat file is written here from that native matrix instead. + * + *

Ported code — version-pinned. The dense float32 layout (format constants, header, meta + * field order, and footer) is transcribed from Lucene 10.2.0, which must stay equal to the + * {@code lucene-core} version in {@code pom.xml}. Sources (tag {@code releases/lucene/10.2.0}): + * + *

    + *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat} — format constants: + * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsFormat.java + *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter} — write sequence: + * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsWriter.java + *
+ * + *

On a Lucene upgrade: re-verify this class against the new version's two files above. The + * constants are package-private in Lucene (hence re-declared here), and the file format is validated + * on read via {@code CodecUtil.checkIndexHeader} (codec name + version range) plus the fixed meta + * layout — so if the new version bumps {@code VERSION_CURRENT}, renames a codec, or changes the meta + * field order, files written here will be silently incompatible and the stock + * {@code Lucene99FlatVectorsReader} will reject or misread them. Update the mirrored constants and + * the {@code writeField}/{@code writeMeta} sequence to match, or bind to Lucene's own writer. + */ +final class NativeFlatVectorsWriter implements Closeable { + + // Mirrors org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat (10.2.0) so the standard + // Lucene99FlatVectorsReader accepts the header/codec of the files written here. + private static final String META_CODEC_NAME = "Lucene99FlatVectorsFormatMeta"; + private static final String VECTOR_DATA_CODEC_NAME = "Lucene99FlatVectorsFormatData"; + private static final String META_EXTENSION = "vemf"; + private static final String VECTOR_DATA_EXTENSION = "vec"; + private static final int VERSION_CURRENT = 0; + private static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; + + // Little-endian float layout matching Lucene's on-disk .vec byte order. Must be UNALIGNED: the + // destination is a heap byte[]-backed MemorySegment whose max alignment is 1 byte, so a 4-byte + // aligned JAVA_FLOAT layout is rejected with "incompatible with alignment constraints". + private static final ValueLayout.OfFloat LE_FLOAT = + ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN); + + // Byte granularity for a single writeBytes call; bounds the transient encode buffer. + private static final int CHUNK_BYTES = 1 << 18; // 256 KiB + + private final IndexOutput meta; + private final IndexOutput vectorData; + private boolean finished; + + NativeFlatVectorsWriter(SegmentWriteState state) throws IOException { + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + String vectorDataFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); + boolean success = false; + try { + meta = state.directory.createOutput(metaFileName, state.context); + vectorData = state.directory.createOutput(vectorDataFileName, state.context); + CodecUtil.writeIndexHeader( + meta, + META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + vectorData, + VECTOR_DATA_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Writes one dense float32 field: the raw vectors to {@code .vec} and the field metadata (plus the + * ordinal-to-doc mapping) to {@code .vemf}. Vectors are read from {@code matrix} in ordinal order, + * which matches the ascending-docID order in which {@code docsWithField} was populated. + * + * @param field the field being written + * @param matrix the native host matrix holding {@code docsWithField.cardinality()} rows of {@code + * field.getVectorDimension()} floats each + * @param maxDoc the segment's maxDoc, used to build the ordinal-to-doc mapping + * @param docsWithField the set of docs that have a value for this field + */ + void writeField( + FieldInfo field, CuVSHostMatrix matrix, int maxDoc, DocsWithFieldSet docsWithField) + throws IOException { + // Mirrors Lucene99FlatVectorsWriter#writeField (see class-level version pin). + int count = docsWithField.cardinality(); + int dim = field.getVectorDimension(); + long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES); + writeFloat32Vectors(matrix, count, dim); + long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset; + writeMeta(field, maxDoc, count, vectorDataOffset, vectorDataLength, docsWithField); + } + + private void writeFloat32Vectors(CuVSHostMatrix matrix, int count, int dim) throws IOException { + int rowBytes = dim * Float.BYTES; + int chunkRows = Math.max(1, CHUNK_BYTES / rowBytes); + byte[] chunk = new byte[chunkRows * rowBytes]; + MemorySegment chunkSeg = MemorySegment.ofArray(chunk); + float[] rowBuf = new float[dim]; + int r = 0; + for (int ord = 0; ord < count; ord++) { + matrix.getRow(ord).toArray(rowBuf); // native -> heap float[] (bulk) + MemorySegment.copy(rowBuf, 0, chunkSeg, LE_FLOAT, (long) r * rowBytes, dim); // -> LE bytes + if (++r == chunkRows) { + vectorData.writeBytes(chunk, r * rowBytes); + r = 0; + } + } + if (r > 0) { + vectorData.writeBytes(chunk, r * rowBytes); + } + } + + private void writeMeta( + FieldInfo field, + int maxDoc, + int count, + long vectorDataOffset, + long vectorDataLength, + DocsWithFieldSet docsWithField) + throws IOException { + // Mirrors Lucene99FlatVectorsWriter#writeMeta (see class-level version pin); field order is + // load-bearing and must match Lucene's reader. + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(field.getVectorSimilarityFunction().ordinal()); + meta.writeVLong(vectorDataOffset); + meta.writeVLong(vectorDataLength); + meta.writeVInt(field.getVectorDimension()); + meta.writeInt(count); + OrdToDocDISIReaderConfiguration.writeStoredMeta( + DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, count, maxDoc, docsWithField); + } + + /** Writes the end-of-fields marker and footers. Mirrors {@code Lucene99FlatVectorsWriter.finish}. */ + void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + if (meta != null) { + meta.writeInt(-1); + CodecUtil.writeFooter(meta); + } + if (vectorData != null) { + CodecUtil.writeFooter(vectorData); + } + } + + @Override + public void close() throws IOException { + IOUtils.close(meta, vectorData); + } +} From 3ad0d470743168d6b436d2ff40ec4aac58c5931b Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 15 Jul 2026 15:04:39 -0700 Subject: [PATCH 06/20] Parallelize CAGRA-to-HNSW conversion in GPUBuiltHnswGraph fillNeighborArray materialized the CAGRA adjacency into on-heap NeighborArrays in a serial per-node loop (~8.5s at 10M x degree-88). Parallelize it under the existing writerThreads knob. CagraIndex.getGraph() returns a CuVSDeviceMatrix whose getRow uses a shared, stateful buffered reader that is not safe for concurrent access, so the parallel path pulls layer 0 to host once (bulk device->host copy) and materializes disjoint node ranges concurrently over the stateless host matrix, closing the temp copy after. The serial path (writerThreads=1) is unchanged and reads the device matrix directly. Applies to all CAGRA_HNSW variants. deep1b-10M @ 16 threads: hnsw-convert 8.5s -> 3.0s, recall unchanged. --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 7 +- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 98 +++++++++++++++++-- .../Lucene99AcceleratedHNSWVectorsWriter.java | 3 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 3 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 3 +- 5 files changed, 98 insertions(+), 16 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 350ae2ad97..50642b6142 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -75,7 +75,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens layerAdjacencies.add(adjacencyMatrix); // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); } /** @@ -99,7 +99,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( int hnswLayers, int graphDegree, CagraIndexParams params, - QuantizationType quantization) + QuantizationType quantization, + int numThreads) throws Throwable { int size = (int) vectorDataset.size(); @@ -161,7 +162,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( random = new Random(new Random().nextLong()); } - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 7e9f888e32..beb0f0cc77 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -6,10 +6,16 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -38,9 +44,14 @@ public class GPUBuiltHnswGraph extends HnswGraph { * @param dimensions the vector dimension * @param layerNodes the nodes on the layer * @param layerAdjacencies adjacency list + * @param numThreads threads to use for materializing the adjacency (1 = serial) */ public GPUBuiltHnswGraph( - int size, int dimensions, List layerNodes, List layerAdjacencies) { + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int numThreads) { this.size = size; this.dimensions = dimensions; @@ -50,38 +61,105 @@ public GPUBuiltHnswGraph( // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size, numThreads); // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { int[] nodes = layerNodes.get(level); CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); - this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length, numThreads)); } } + /** Node count below which parallel materialization is not worth the thread overhead. */ + private static final int PARALLEL_MIN_NODES = 1 << 16; + /** - * Fills the neighbor array using the adjacency matrix. + * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. + * + *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe + * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose + * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so + * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node + * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are + * read directly in both paths. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes + * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) * @return the NeighborArray */ - private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + private static NeighborArray[] fillNeighborArray( + CuVSMatrix adjacency, int size, int numThreads) { NeighborArray[] neighbors = new NeighborArray[size]; - for (int i = 0; i < size; i++) { - RowView rv = adjacency.getRow(i); + if (numThreads <= 1 || size < PARALLEL_MIN_NODES) { + fillNeighborRange(adjacency, neighbors, 0, size); + return neighbors; + } + CuVSMatrix source = adjacency; + CuVSHostMatrix hostCopy = null; + if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { + hostCopy = deviceAdjacency.toHost(); + source = hostCopy; + } + try { + fillNeighborArrayParallel(source, neighbors, size, numThreads); + return neighbors; + } finally { + if (hostCopy != null) { + hostCopy.close(); + } + } + } + + /** + * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code + * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code + * source} must be a host matrix (stateless {@code getRow}). + */ + private static void fillNeighborArrayParallel( + CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) { + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + try { + int perThread = (size + numThreads - 1) / numThreads; + List> futures = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int start = t * perThread; + final int end = Math.min(start + perThread, size); + if (start >= end) { + break; + } + futures.add(pool.submit(() -> fillNeighborRange(source, neighbors, start, end))); + } + for (Future f : futures) { + f.get(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted during parallel HNSW conversion", e); + } catch (ExecutionException e) { + throw new RuntimeException("Failed during parallel HNSW conversion", e.getCause()); + } finally { + pool.shutdown(); + } + } + + /** Fills {@code neighbors[start, end)} from the adjacency rows. */ + private static void fillNeighborRange( + CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { + for (int i = start; i < end; i++) { + RowView rv = source.getRow(i); if (rv != null && rv.size() > 0) { - neighbors[i] = new NeighborArray((int) rv.size(), true); + NeighborArray na = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { - neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + na.addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); } + neighbors[i] = na; } else { neighbors[i] = new NeighborArray(0, true); } } - return neighbors; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 3236f6b1e5..e1f707edaf 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -228,7 +228,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws acceleratedHNSWParams.getHnswLayers(), acceleratedHNSWParams.getGraphdegree(), params, - QuantizationType.NONE); + QuantizationType.NONE, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 00c5f1d4cb..82cd4ed5ba 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -183,7 +183,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.BINARY); + QuantizationType.BINARY, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 62b67657ed..0a53bfbd15 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -208,7 +208,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.SCALAR); + QuantizationType.SCALAR, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); From 948b41dac887edcfd4c6ccfd0f265a3dfde11d43 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 30 Jul 2026 19:19:23 -0700 Subject: [PATCH 07/20] Add prefetching + reused-array to the fbin ingest example Upgrade the large-.fbin ingestion reference from a single-buffered, fresh-array-per-vector reader to the optimized ingestion pattern used by the benchmark harness, and rename ChunkedFbinIngestExample -> OptimizedFbinIngestExample to reflect it. PrefetchingFbinReader now demonstrates the four properties that keep ingestion from bottlenecking the GPU build: - open once, read sequentially in large chunks (no fd churn) - bounded memory (at most two chunks resident) - overlap: a background reader thread fills the next chunk while the ingest thread feeds the current one into addDocument, hiding the disk read behind per-document indexing - reuse the vector array: get(int, float[]) unpacks into a caller-owned array, safe because Lucene copies each vector eagerly at addDocument Scope is ingestion only; the codec stays at defaults (no native flat buffering or other config-gated behavior). README updated (description, run command, jar version 26.08.0 -> 26.10.0). --- java/cuvs-lucene/examples/README.md | 11 +- ...e.java => OptimizedFbinIngestExample.java} | 196 +++++++++++++----- 2 files changed, 154 insertions(+), 53 deletions(-) rename java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/{ChunkedFbinIngestExample.java => OptimizedFbinIngestExample.java} (51%) diff --git a/java/cuvs-lucene/examples/README.md b/java/cuvs-lucene/examples/README.md index 1535d644ba..c2c724bac4 100644 --- a/java/cuvs-lucene/examples/README.md +++ b/java/cuvs-lucene/examples/README.md @@ -33,13 +33,14 @@ To run the Index and Search on GPU example do: mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample ``` -To run the chunked `.fbin` ingestion example (reference pattern for streaming a large vector file -into an accelerated HNSW index without per-vector file reopening or holding the whole file in -memory) do: +To run the chunked `.fbin` ingestion example (reference pattern for efficiently streaming a large +vector file into an accelerated HNSW index — open the file once, read sequential prefetched chunks +that overlap the disk read with indexing, hold at most two chunks in memory, and reuse a single +vector array across all documents) do: ```sh -mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.08.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.ChunkedFbinIngestExample +mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.OptimizedFbinIngestExample ``` With no arguments it generates and indexes a small demo `.fbin`; pass a real file and chunk size as -`... ChunkedFbinIngestExample `. +`... OptimizedFbinIngestExample `. diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java similarity index 51% rename from java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java rename to java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java index a3ce8afad5..559159005a 100644 --- a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/ChunkedFbinIngestExample.java +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java @@ -18,6 +18,8 @@ import java.nio.file.StandardOpenOption; import java.util.Random; import java.util.UUID; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; import java.util.logging.Logger; import org.apache.commons.io.FileUtils; import org.apache.lucene.codecs.Codec; @@ -36,29 +38,36 @@ import org.apache.lucene.store.FSDirectory; /** - * Reference pattern for ingesting a LARGE {@code .fbin} vector file into an accelerated HNSW index - * without the two pitfalls that dominate build time and memory when loading big vector files: + * Reference pattern for efficiently ingesting a LARGE {@code .fbin} vector file into an accelerated + * HNSW index. cuvs-lucene is a Lucene codec and sits below {@code addDocument}, so how you read your + * source data is application code — this example shows the four properties that keep ingestion from + * bottlenecking the GPU build: * *

    - *
  1. fd churn — reopening/seeking/closing the file per vector is orders of magnitude - * slower than sequential reads. {@link ChunkedFbinReader} opens the file ONCE and reads it - * front-to-back in large sequential chunks. - *
  2. unbounded / doubled memory — pre-loading the whole file into the JVM heap holds a - * redundant copy on top of Lucene's own per-segment buffer (~2x peak). The chunked reader - * holds at most one chunk, streaming each vector straight into {@code addDocument}. + *
  3. Open once, read sequentially. Reopening/seeking/closing the file per vector ("fd + * churn") is orders of magnitude slower than sequential reads. {@link PrefetchingFbinReader} + * opens the file ONCE and reads it front-to-back in large sequential chunks. + *
  4. Bounded memory. Pre-loading the whole file onto the JVM heap holds a redundant copy + * on top of Lucene's own per-segment buffer (~2x peak). This reader holds at most two chunks. + *
  5. Overlap read with consumption. A background reader thread fills the NEXT chunk (into a + * second buffer) while the ingest thread is still feeding the current chunk into {@code + * addDocument}, so the sequential disk read is hidden behind the per-document indexing work + * instead of serializing in front of it. + *
  6. Reuse the vector array. {@link PrefetchingFbinReader#get(int, float[])} unpacks + * directly into a caller-supplied array, avoiding a fresh {@code float[]} allocation per + * vector. This is safe because Lucene copies the vector value eagerly inside {@code + * addDocument} — the array may be reused as soon as {@code addDocument} returns. *
* - *

cuvs-lucene is a Lucene codec and sits below {@code addDocument}, so how you read your source - * data is application code — adapt {@link ChunkedFbinReader} to your own source (a DB, object - * store, or stream). The properties that matter are: open once, read sequentially, bound memory - * to a chunk. + *

Adapt {@link PrefetchingFbinReader} to your own source (a DB, object store, or stream). The + * properties above are what matter, not the {@code .fbin} specifics. * - *

Usage: {@code ChunkedFbinIngestExample [] []}. With no arguments a + *

Usage: {@code OptimizedFbinIngestExample [] []}. With no arguments a * small demo {@code .fbin} is generated and indexed. */ -public class ChunkedFbinIngestExample { +public class OptimizedFbinIngestExample { - private static final Logger log = Logger.getLogger(ChunkedFbinIngestExample.class.getName()); + private static final Logger log = Logger.getLogger(OptimizedFbinIngestExample.class.getName()); private static final String ID_FIELD = "id"; private static final String VECTOR_FIELD = "vector_field"; @@ -90,14 +99,15 @@ public static void main(String[] args) throws Exception { /** * Builds a single-segment accelerated HNSW index, streaming vectors from the {@code .fbin} in - * sequential chunks (never holding more than one chunk in memory). + * sequential prefetched chunks (never holding more than two chunks in memory) and reusing a single + * vector array across all documents. */ private static void buildIndex(Path fbinPath, Path indexDirPath, int chunkSizeMB) throws Exception { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); Codec codec = new Lucene101AcceleratedHNSWCodec(params); - try (ChunkedFbinReader reader = new ChunkedFbinReader(fbinPath, chunkSizeMB)) { + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, chunkSizeMB)) { int n = reader.size(); // Keep the whole dataset in ONE segment: disable RAM-based flushing and raise the doc-count @@ -122,16 +132,18 @@ private static void buildIndex(Path fbinPath, Path indexDirPath, int chunkSizeMB + fbinPath + " using " + chunkSizeMB - + " MB sequential chunks"); + + " MB prefetched sequential chunks"); + // One reusable array for the whole build — refilled per vector, copied eagerly by Lucene. + float[] vector = new float[reader.dimension()]; try (Directory dir = FSDirectory.open(indexDirPath); IndexWriter writer = new IndexWriter(dir, config)) { for (int i = 0; i < n; i++) { - float[] vector = reader.get(i); // sequential access -> served from the current chunk + reader.get(i, vector); // sequential access -> served from the prefetched chunk, no alloc Document doc = new Document(); doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); - writer.addDocument(doc); + writer.addDocument(doc); // copies the vector -> 'vector' is safe to reuse next iteration } writer.commit(); // single flush -> single segment; the GPU CAGRA build happens here } @@ -142,7 +154,7 @@ private static void buildIndex(Path fbinPath, Path indexDirPath, int chunkSizeMB /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { float[] queryVector; - try (ChunkedFbinReader reader = new ChunkedFbinReader(fbinPath, 1)) { + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, 1)) { queryVector = reader.get(0); } try (Directory dir = FSDirectory.open(indexDirPath); @@ -184,30 +196,50 @@ private static void writeDemoFbin(Path path, int numVectors, int dim, Random ran } /** - * Chunked, sequential reader for uncompressed {@code .fbin} files - * ({@code [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). + * Prefetching, double-buffered sequential reader for uncompressed {@code .fbin} files ({@code + * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). * - *

Opens the file ONCE and serves {@link #get(int)} from a reusable buffer that is refilled - * with a single large sequential read whenever the requested index leaves the current chunk. For - * sequential access (index = 0, 1, 2, ...) it reads the file front-to-back in {@code size/chunk} - * bulk reads while holding only one chunk in memory — the opposite of reopening the file per - * vector. + *

Opens the file ONCE. A background thread reads the file front-to-back into two reusable direct + * buffers: while the caller consumes the current chunk, the reader fills the next one, so the disk + * read overlaps with the caller's per-vector work. {@link #get(int, float[])} unpacks directly into + * a caller-supplied array (no per-vector allocation). + * + *

Forward-only, single-consumer. {@code get} must be called with non-decreasing indices + * from a single thread — the intended pattern for streaming ingestion. */ - static final class ChunkedFbinReader implements AutoCloseable { + static final class PrefetchingFbinReader implements AutoCloseable { private static final long HEADER_BYTES = 8; + private static final class Chunk { + final ByteBuffer buf; + final long start; + final int len; + + Chunk(ByteBuffer buf, long start, int len) { + this.buf = buf; + this.start = start; + this.len = len; + } + } + + /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ + private static final Chunk POISON = new Chunk(null, -1, 0); + private final FileChannel channel; private final int dimension; private final int vectorCount; private final int vectorBytes; private final int chunkVectors; - private final ByteBuffer chunkBuffer; - private long chunkStart = -1; // first vector index currently buffered - private int chunkLen = 0; // number of vectors currently buffered + private final BlockingQueue free = new ArrayBlockingQueue<>(2); + private final BlockingQueue ready = new ArrayBlockingQueue<>(2); + private final Thread reader; + private volatile IOException readerError; + + private Chunk current; // consumer-owned; the chunk currently being served - ChunkedFbinReader(Path path, int chunkSizeMB) throws IOException { + PrefetchingFbinReader(Path path, int chunkSizeMB) throws IOException { this.channel = FileChannel.open(path, StandardOpenOption.READ); ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); readFully(header, 0); @@ -219,8 +251,16 @@ static final class ChunkedFbinReader implements AutoCloseable { long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); - this.chunkBuffer = - ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN); + + // Two reusable direct buffers: the reader fills one while the consumer drains the other. + for (int i = 0; i < 2; i++) { + free.add( + ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); + } + + this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); + this.reader.setDaemon(true); + this.reader.start(); } int size() { @@ -231,26 +271,85 @@ int dimension() { return dimension; } - float[] get(int index) throws IOException { + /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ + private void readLoop() { + long next = 0; + try { + while (next < vectorCount) { + ByteBuffer buf = free.take(); + int toRead = (int) Math.min(chunkVectors, vectorCount - next); + buf.clear(); + buf.limit(toRead * vectorBytes); + readFully(buf, HEADER_BYTES + next * (long) vectorBytes); + ready.put(new Chunk(buf, next, toRead)); + next += toRead; + } + ready.put(POISON); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); // close() requested; stop quietly + } catch (IOException e) { + readerError = e; + try { + ready.put(POISON); // unblock the consumer so it can observe the error + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } + } + + private void advance() throws IOException { + if (current != null) { + try { + free.put(current.buf); // hand the drained buffer back to the reader + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted returning chunk buffer", e); + } + current = null; + } + Chunk next; + try { + next = ready.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted awaiting next chunk", e); + } + if (next == POISON) { + if (readerError != null) { + throw new IOException("Prefetch reader failed", readerError); + } + throw new IOException("No more chunks available (unexpected EOF in prefetch)"); + } + current = next; + } + + /** Fills {@code dst} with the vector at {@code index} (no allocation). */ + void get(int index, float[] dst) throws IOException { if (index < 0 || index >= vectorCount) { throw new IndexOutOfBoundsException( "Index " + index + " out of bounds [0, " + vectorCount + ")"); } - if (chunkStart < 0 || index < chunkStart || index >= chunkStart + chunkLen) { - long start = (index / (long) chunkVectors) * chunkVectors; - int toRead = (int) Math.min(chunkVectors, vectorCount - start); - chunkBuffer.clear(); - chunkBuffer.limit(toRead * vectorBytes); - readFully(chunkBuffer, HEADER_BYTES + start * (long) vectorBytes); - chunkStart = start; - chunkLen = toRead; + while (current == null || index >= current.start + current.len) { + advance(); + } + if (index < current.start) { + throw new UnsupportedOperationException( + "PrefetchingFbinReader requires forward-only sequential access; got index " + + index + + " before current chunk start " + + current.start); } - int base = (int) (index - chunkStart) * vectorBytes; - float[] vector = new float[dimension]; + int base = (int) (index - current.start) * vectorBytes; for (int i = 0; i < dimension; i++) { - vector[i] = chunkBuffer.getFloat(base + i * Float.BYTES); + dst[i] = current.buf.getFloat(base + i * Float.BYTES); } - return vector; + } + + /** Convenience allocating variant (e.g. for a one-off query vector). */ + float[] get(int index) throws IOException { + float[] dst = new float[dimension]; + get(index, dst); + return dst; } private void readFully(ByteBuffer buf, long position) throws IOException { @@ -266,6 +365,7 @@ private void readFully(ByteBuffer buf, long position) throws IOException { @Override public void close() throws IOException { + reader.interrupt(); channel.close(); } } From 33c85a21c9fb28218a2336fa643f9543761e4b91 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 31 Jul 2026 09:08:30 -0700 Subject: [PATCH 08/20] Honor cagraGraphBuildAlgo override in HEURISTIC strategy HEURISTIC previously always chose the CAGRA build algo by row count (NN_DESCENT below 5M, else IVF_PQ) and ignored cagraGraphBuildAlgo. Now an explicit IVF_PQ or NN_DESCENT overrides that choice, while AUTO_SELECT keeps the row-count determination; params for the chosen algo stay auto-tuned (getIVFPQParams / getNNDescentParams). Flip DEFAULT_CAGRA_GRAPH_BUILD_ALGO from NN_DESCENT to AUTO_SELECT so callers that use the default keep the row-count behavior (the switch above would otherwise force NN_DESCENT for them). CUSTOM strategy is unchanged. --- .../cuvs/lucene/AcceleratedHNSWParams.java | 2 +- .../cuvs/lucene/CagraIndexParamsFactory.java | 168 +++++++++++++++--- 2 files changed, 147 insertions(+), 23 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 2d0cc64a97..67f69289df 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -59,7 +59,7 @@ public static enum Strategy { public static final int DEFAULT_MAX_CONN = 32; public static final int DEFAULT_BEAM_WIDTH = 32; public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = - CagraGraphBuildAlgo.NN_DESCENT; + CagraGraphBuildAlgo.AUTO_SELECT; public static final int DEFAULT_NUM_MERGE_WORKERS = 1; public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 8bb46faf61..5be2451348 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -6,6 +6,13 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CodebookGen; +import com.nvidia.cuvs.CagraIndexParams.CudaDataType; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSIvfPqIndexParams; +import com.nvidia.cuvs.CuVSIvfPqParams; +import com.nvidia.cuvs.CuVSIvfPqSearchParams; /** * A centralized place for producing {@link CagraIndexParams} from the cuvs-lucene input parameter @@ -16,8 +23,104 @@ */ public class CagraIndexParamsFactory { + private static final int ALGO_SWITCH_THRESHOLD = 5_000_000; + private CagraIndexParamsFactory() {} + private static CuVSIvfPqParams getCuVSIvfPqParams(long rows, long dimension) { + int pqDim; + int pqBits; + if (dimension <= 32) { + pqDim = 16; + pqBits = 8; + } else { + pqBits = 4; + if (dimension <= 64) { + pqDim = 32; + } else if (dimension <= 128) { + pqDim = 64; + } else if (dimension <= 192) { + pqDim = 96; + } else { + pqDim = (int) roundUpSafe(dimension / 2, 128); + } + } + int nLists = (int) Math.max(1, rows / 2000); + final int kmeansNIters = 10; + final double kMinPointsPerCluster = 32; + double minKmeansTrainsetPoints = kMinPointsPerCluster * nLists; + final double maxKmeansTrainsetFraction = 1.0; + double minKmeansTrainsetFraction = + Math.min(maxKmeansTrainsetFraction, minKmeansTrainsetPoints / rows); + double kmeansTrainsetFraction = + Math.clamp( + 1.0 / Math.sqrt(rows * 1e-5), minKmeansTrainsetFraction, maxKmeansTrainsetFraction); + int nProbes = (int) Math.round(Math.sqrt(nLists) / 20 + 4); + CuVSIvfPqIndexParams cuVSIvfPqIndexParams = + new CuVSIvfPqIndexParams.Builder() + .withCodebookKind(CodebookGen.PER_SUBSPACE) + .withKmeansNIters(kmeansNIters) + .withKmeansTrainsetFraction(kmeansTrainsetFraction) + .withNLists(nLists) + .withPqBits(pqBits) + .withPqDim(pqDim) + .withAddDataOnBuild(true) + .withConservativeMemoryAllocation(true) + .build(); + CuVSIvfPqSearchParams cuVSIvfPqSearchParams = + new CuVSIvfPqSearchParams.Builder() + .withLutDtype(CudaDataType.CUDA_R_16F) + .withInternalDistanceDtype(CudaDataType.CUDA_R_16F) + .withNProbes(nProbes) + .build(); + return new CuVSIvfPqParams.Builder() + .withCuVSIvfPqIndexParams(cuVSIvfPqIndexParams) + .withCuVSIvfPqSearchParams(cuVSIvfPqSearchParams) + .withRefinementRate(1) + .build(); + } + + private static long roundUpSafe(long numberToRound, long modulus) { + long remainder = numberToRound % modulus; + if (remainder == 0) { + return numberToRound; + } + return numberToRound - remainder + modulus; + } + + private static CagraIndexParams getNNDescentParams( + int graphDegree, + int intGraphDegree, + int writerThreads, + long nnDescentNumIterations, + CuvsDistanceType cuvsDistanceType) { + return new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(graphDegree) + .withIntermediateGraphDegree(intGraphDegree) + .withNNDescentNumIterations(nnDescentNumIterations) + .withNumWriterThreads(writerThreads) + .withMetric(cuvsDistanceType) + .build(); + } + + private static CagraIndexParams getIVFPQParams( + int graphDegree, + int intGraphDegree, + int writerThreads, + long rows, + long dimension, + CuvsDistanceType cuvsDistanceType) { + return new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.IVF_PQ) + .withCuVSIvfPqParams(getCuVSIvfPqParams(rows, dimension)) + .withNumWriterThreads(writerThreads) + .withIntermediateGraphDegree(intGraphDegree) + .withGraphDegree(graphDegree) + .withMetric(cuvsDistanceType) + .build(); + } + /** * Creates an instance of {@link CagraIndexParams} for the GPU-native CAGRA index based on the * chosen strategy in the {@link GPUSearchParams}. @@ -74,29 +177,50 @@ public static CagraIndexParams create( public static CagraIndexParams create( AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { - // Delegate the derivation of the graph degrees, build algorithm and its parameters to cuVS, - // expressed in terms of the HNSW-equivalent maxConn/beamWidth. - CagraIndexParams derived = - CagraIndexParams.fromHnswParams( - rows, - dimension, - acceleratedHNSWParams.getMaxConn(), - acceleratedHNSWParams.getBeamWidth(), - acceleratedHNSWParams.getHnswHeuristicType(), - acceleratedHNSWParams.getCuvsDistanceType()); - // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS default - // (not a heuristic value). We can rebuild the CagraIndexParams with the caller-supplied - // writerThreads for now but should fix this in cuVS in the future. - return new CagraIndexParams.Builder() - .withGraphDegree(derived.getGraphDegree()) - .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) - .withCagraGraphBuildAlgo(derived.getCagraGraphBuildAlgo()) - .withCuVSIvfPqParams(derived.getCuVSIvfPqParams()) - .withNNDescentNumIterations(derived.getNNDescentNumIterations()) - .withMetric(acceleratedHNSWParams.getCuvsDistanceType()) - .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) - .build(); + // An explicit cagraGraphBuildAlgo of IVF_PQ or NN_DESCENT overrides the heuristic choice; + // AUTO_SELECT (the default) delegates all parameter derivation to cuVS via fromHnswParams. + CagraGraphBuildAlgo algo = acceleratedHNSWParams.getCagraGraphBuildAlgo(); + if (algo == CagraGraphBuildAlgo.IVF_PQ) { + return getIVFPQParams( + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getWriterThreads(), + rows, + dimension, + acceleratedHNSWParams.getCuvsDistanceType()); + } else if (algo == CagraGraphBuildAlgo.NN_DESCENT) { + return getNNDescentParams( + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getNNDescentNumIterations(), + acceleratedHNSWParams.getCuvsDistanceType()); + } else { + // AUTO_SELECT: delegate the derivation of the graph degrees, build algorithm and its + // parameters to cuVS, expressed in terms of the HNSW-equivalent maxConn/beamWidth. + CagraIndexParams derived = + CagraIndexParams.fromHnswParams( + rows, + dimension, + acceleratedHNSWParams.getMaxConn(), + acceleratedHNSWParams.getBeamWidth(), + acceleratedHNSWParams.getHnswHeuristicType(), + acceleratedHNSWParams.getCuvsDistanceType()); + // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS + // default (not a heuristic value). We can rebuild the CagraIndexParams with the + // caller-supplied writerThreads for now but should fix this in cuVS in the future. + return new CagraIndexParams.Builder() + .withGraphDegree(derived.getGraphDegree()) + .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) + .withCagraGraphBuildAlgo(derived.getCagraGraphBuildAlgo()) + .withCuVSIvfPqParams(derived.getCuVSIvfPqParams()) + .withNNDescentNumIterations(derived.getNNDescentNumIterations()) + .withMetric(acceleratedHNSWParams.getCuvsDistanceType()) + .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) + .build(); + } } + // CUSTOM: forward the caller's algorithm and the parameters it consumes. return new CagraIndexParams.Builder() .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) From 1b9b9d8f85b996a7ff4194767c4edacc56bebefb Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 31 Jul 2026 14:06:08 -0700 Subject: [PATCH 09/20] Expand the fbin ingest example into a full optimized CAGRA-HNSW build Rename OptimizedFbinIngestExample -> OptimizedCagraHnswBuildExample and turn on every build-side optimization, so the example is a single golden reference for building a large accelerated HNSW index rather than just an ingestion pattern: - Native flat buffering: size the codec's numInputVectors to each segment, streaming vectors into the native host matrix instead of a heap List (single-segment; ~halves peak host memory). - Automatic graph-build algorithm: HEURISTIC + AUTO_SELECT lets cuVS pick NN_DESCENT vs IVF_PQ by dataset size and auto-tune its parameters. - Partitioned multi-segment build with a user-specified segment count: - sequential: K single-segment passes appended to one directory, peak host = one slice (N/K); - overlapped: a bounded pool (PIPELINE_DEPTH) builds segments into their own dirs with the GPU commit serialized on a permit, then hardlinks them into the final directory via addIndexes (no bulk copy). Peak host = depth * (N/K). The class Javadoc states the assumptions behind the segment count: native flat buffering forces one segment per slice, the GPU build is serialized (more segments buy host-memory headroom + ingest overlap, not device parallelism), and search fans out across all K segment graphs. Adds a lucene-misc dependency (HardlinkCopyDirectoryWrapper) and updates the examples README with the new name and 4-arg usage. --- java/cuvs-lucene/examples/README.md | 16 +- java/cuvs-lucene/examples/pom.xml | 5 + .../OptimizedCagraHnswBuildExample.java | 628 ++++++++++++++++++ .../examples/OptimizedFbinIngestExample.java | 372 ----------- 4 files changed, 642 insertions(+), 379 deletions(-) create mode 100644 java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java delete mode 100644 java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java diff --git a/java/cuvs-lucene/examples/README.md b/java/cuvs-lucene/examples/README.md index c2c724bac4..0d61d0a7c0 100644 --- a/java/cuvs-lucene/examples/README.md +++ b/java/cuvs-lucene/examples/README.md @@ -33,14 +33,16 @@ To run the Index and Search on GPU example do: mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample ``` -To run the chunked `.fbin` ingestion example (reference pattern for efficiently streaming a large -vector file into an accelerated HNSW index — open the file once, read sequential prefetched chunks -that overlap the disk read with indexing, hold at most two chunks in memory, and reuse a single -vector array across all documents) do: +To run the optimized CAGRA-HNSW build example (reference pattern for efficiently building an +accelerated HNSW index from a large `.fbin` with every ingest-side knob on — open the file once and +stream sequential prefetched chunks that overlap the disk read with indexing, hold at most two chunks +in memory, reuse a single vector array, size a native flat buffer per segment, auto-select the CAGRA +graph-build algorithm, and optionally partition into K segments built sequentially or overlapped) do: ```sh -mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.OptimizedFbinIngestExample +mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.OptimizedCagraHnswBuildExample ``` -With no arguments it generates and indexes a small demo `.fbin`; pass a real file and chunk size as -`... OptimizedFbinIngestExample `. +With no arguments it generates and indexes a small demo `.fbin` as a single segment; pass a real file, +chunk size, segment count, and overlap flag as +`... OptimizedCagraHnswBuildExample `. diff --git a/java/cuvs-lucene/examples/pom.xml b/java/cuvs-lucene/examples/pom.xml index 4432883f64..e8b1833c8a 100644 --- a/java/cuvs-lucene/examples/pom.xml +++ b/java/cuvs-lucene/examples/pom.xml @@ -92,6 +92,11 @@ lucene-backward-codecs 10.2.0 + + org.apache.lucene + lucene-misc + 10.2.0 + commons-io commons-io diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java new file mode 100644 index 0000000000..9e93e99508 --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java @@ -0,0 +1,628 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.misc.store.HardlinkCopyDirectoryWrapper; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** + * Reference pattern for building an accelerated HNSW index (whose graph is built on the GPU with + * CAGRA) from a LARGE {@code .fbin} vector file with every optimization turned on — spanning + * both how vectors are ingested and how the GPU build is configured and partitioned. cuvs-lucene is a + * Lucene codec that sits below {@code addDocument}, so how you read your source data and how you shape + * segments is application code — this example shows the full recipe. + * + *

The knobs, and why each one matters

+ * + *
    + *
  1. Streaming, prefetched, bounded-memory reads. {@link PrefetchingFbinReader} opens the + * file ONCE, reads it front-to-back in large sequential chunks (no per-vector seek/close "fd + * churn"), holds at most two chunks (no whole-file heap copy), fills the NEXT chunk on a + * background thread while the ingest thread drains the current one (disk read hidden behind + * indexing), and unpacks into a caller-reused {@code float[]} (no per-vector allocation — safe + * because Lucene copies the value eagerly inside {@code addDocument}). + *
  2. Native flat buffering. {@link AcceleratedHNSWParams.Builder#withNumInputVectors} sizes + * a native host matrix to exactly the segment's vector count, so vectors stream straight into + * the buffer the GPU build consumes instead of piling up as a {@code List} on the JVM + * heap that must then be assembled. This removes the assembly copy and roughly halves peak host + * memory. It requires a single-segment build (no flush before commit, no merge). + *
  3. Automatic graph-build algorithm. {@code HEURISTIC} + {@code AUTO_SELECT} lets cuVS + * pick the CAGRA build algorithm by dataset size (NN_DESCENT below ~5M vectors, IVF_PQ at or + * above) and auto-tune its parameters. NN_DESCENT generally reaches higher recall but takes + * longer to build; IVF_PQ is faster but slightly lower recall. Force one explicitly only under + * expert guidance. + *
  4. Partitioned multi-segment build. Because native flat buffering is single-segment, "K + * segments" means K independent single-segment builds over contiguous slices, combined at the + * end. This is a deliberate memory/throughput lever — see the assumptions below. + *
+ * + *

Assumptions for the user-specified number of segments

+ * + *

{@code numSegments} is your choice and it is a trade-off, not a free speedup: + * + *

    + *
  • Each segment is one native-flat build over a {@code 1/K} slice, so peak host memory scales + * as {@code 1/K} (sequential mode) — this is the point of partitioning for large or + * memory-bounded datasets. + *
  • The GPU is serialized: only one CAGRA build runs on the device at a time. Extra + * segments do NOT parallelize the graph build; they only reduce host memory and (in overlap + * mode) let one segment's host-side ingest run during a prior segment's GPU commit. + *
  • Search fans out across all K segment graphs, so more segments trade a little query throughput + * (and can shift recall) for lower build-time memory. Pick K to fit your memory budget, not + * higher. + *
  • Each slice must stay a single segment: {@code maxBufferedDocs > sliceSize} and {@link + * NoMergePolicy} (no flush, no merge) so native flat buffering stays valid. + *
+ * + *

Sequential vs overlapped

+ * + *

With {@code overlap=false} the K slices are built as K sequential passes appended to one + * directory; peak host memory is a single slice's buffer ({@code N/K}). With {@code overlap=true} a + * bounded pool builds up to {@code PIPELINE_DEPTH} segments at once into their own directories, so a + * segment's ingest overlaps a prior segment's (serialized) GPU commit; the finished per-segment + * indexes are then combined by hardlinking their files into the final directory ({@link + * HardlinkCopyDirectoryWrapper} + {@code addIndexes}, no bulk copy of the vector data). Overlap costs + * up to {@code PIPELINE_DEPTH * (N/K)} peak host memory and can hide most of the ingest time behind + * the GPU build, bounded by disk contention between concurrent reads and the committing segment's + * writes. + * + *

Adapt {@link PrefetchingFbinReader} to your own source (a DB, object store, or stream). The + * properties above are what matter, not the {@code .fbin} specifics. + * + *

Usage: {@code OptimizedCagraHnswBuildExample [] [] [] + * []}. With no arguments a small demo {@code .fbin} is generated and indexed as a + * single segment. + */ +public class OptimizedCagraHnswBuildExample { + + private static final Logger log = + Logger.getLogger(OptimizedCagraHnswBuildExample.class.getName()); + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + /** Max segments built concurrently in overlap mode; peak host memory is this many slice buffers. */ + private static final int PIPELINE_DEPTH = 2; + + public static void main(String[] args) throws Exception { + int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; + int numSegments = args.length >= 3 ? Math.max(1, Integer.parseInt(args[2])) : 1; + boolean overlap = args.length >= 4 && Boolean.parseBoolean(args[3]); + Path indexDirPath = Paths.get(UUID.randomUUID().toString()); + + Path fbinPath; + boolean generated = false; + if (args.length >= 1) { + fbinPath = Paths.get(args[0]); + } else { + fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); + writeDemoFbin(fbinPath, 5000, 32, new Random(222)); + generated = true; + log.info("No .fbin provided; generated a demo file at " + fbinPath); + } + + try { + buildIndex(fbinPath, indexDirPath, chunkSizeMB, numSegments, overlap); + runSampleSearch(indexDirPath, fbinPath, 5); + } finally { + FileUtils.deleteDirectory(indexDirPath.toFile()); + if (generated) { + Files.deleteIfExists(fbinPath); + } + } + } + + /** + * Builds an accelerated HNSW index from {@code fbinPath} into {@code numSegments} contiguous + * slices, dispatching to the sequential or overlapped partitioned build. {@code numSegments == 1} + * is the plain single-segment native-flat build. + */ + private static void buildIndex( + Path fbinPath, Path indexDirPath, int chunkSizeMB, int numSegments, boolean overlap) + throws Exception { + int total; + int dim; + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, 1)) { + total = reader.size(); + dim = reader.dimension(); + } + List slices = sliceEvenly(total, numSegments); // [start, size] per segment + + log.info( + "Indexing " + + total + + " vectors (" + + dim + + "-dim) from " + + fbinPath + + " into " + + slices.size() + + " segment(s), " + + (overlap && slices.size() > 1 ? "overlapped" : "sequential") + + " build, " + + chunkSizeMB + + " MB prefetched chunks"); + + if (overlap && slices.size() > 1) { + buildOverlapped(fbinPath, indexDirPath, chunkSizeMB, slices, dim); + } else { + buildSequential(fbinPath, indexDirPath, chunkSizeMB, slices, dim); + } + log.info("Index build complete: " + indexDirPath); + } + + /** + * Sequential partitioned build: one forward-only prefetch reader streamed front-to-back across all + * K slices, each slice built as a single native-flat segment appended to the same directory (first + * pass {@code CREATE}, later passes {@code APPEND}). Peak host memory is one slice's native buffer. + */ + private static void buildSequential( + Path fbinPath, Path indexDirPath, int chunkSizeMB, List slices, int dim) + throws Exception { + float[] scratch = new float[dim]; // one reusable array for the whole build + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, chunkSizeMB); + Directory dir = FSDirectory.open(indexDirPath)) { + for (int p = 0; p < slices.size(); p++) { + int[] slice = slices.get(p); + log.info( + "Building segment " + (p + 1) + "/" + slices.size() + ": docs [" + slice[0] + ", " + + (slice[0] + slice[1]) + ")"); + buildSegment(dir, reader, scratch, slice[0], slice[1], p == 0, null); + } + } + } + + /** + * Overlapped partitioned build: a bounded pool builds up to {@link #PIPELINE_DEPTH} segments at + * once, each into its OWN directory with its OWN prefetch reader over just its slice, so a + * segment's ingest overlaps a prior segment's GPU commit. The GPU build is serialized on a single + * permit. The finished per-segment indexes are then hardlinked into the final directory (no bulk + * copy of the vector data). + */ + private static void buildOverlapped( + Path fbinPath, Path indexDirPath, int chunkSizeMB, List slices, int dim) + throws Exception { + int depth = Math.min(slices.size(), PIPELINE_DEPTH); + int maxSlice = slices.stream().mapToInt(s -> s[1]).max().orElse(0); + double peakHostGb = (double) depth * maxSlice * dim * Float.BYTES / 1e9; + log.info( + "Overlapped build: " + + slices.size() + + " segment(s), pipeline depth " + + depth + + " (up to " + + depth + + " co-resident native host buffers, ~" + + String.format("%.2f", peakHostGb) + + " GB peak host)"); + + List segDirs = new ArrayList<>(); + for (int p = 0; p < slices.size(); p++) { + segDirs.add(Paths.get(indexDirPath + "_p" + p)); + } + // Start from fresh per-segment temp dirs, and always remove them afterwards (even on failure) so + // a crashed build does not leave orphaned per-segment indexes behind. + for (Path segDir : segDirs) { + FileUtils.deleteQuietly(segDir.toFile()); + } + try { + Semaphore gpuPermit = new Semaphore(1); // serialize the GPU CAGRA build across segments + ExecutorService pool = Executors.newFixedThreadPool(depth); + List> futures = new ArrayList<>(); + for (int p = 0; p < slices.size(); p++) { + int[] slice = slices.get(p); + Path segDir = segDirs.get(p); + futures.add( + pool.submit( + () -> { + float[] scratch = new float[dim]; + try (PrefetchingFbinReader reader = + new PrefetchingFbinReader(fbinPath, slice[0], slice[1], chunkSizeMB); + Directory d = FSDirectory.open(segDir)) { + // createNew=true: each segment is a fresh single-segment index in its own dir. + buildSegment(d, reader, scratch, slice[0], slice[1], true, gpuPermit); + } + return null; + })); + } + pool.shutdown(); + try { + for (Future f : futures) { + f.get(); // propagate any build failure + } + } finally { + pool.shutdownNow(); + } + combineByHardlink(indexDirPath, segDirs); + } finally { + for (Path segDir : segDirs) { + FileUtils.deleteQuietly(segDir.toFile()); + } + } + } + + /** + * Builds one segment from the contiguous slice {@code [start, start + size)}: a single-threaded + * {@link IndexWriter} whose codec's {@code numInputVectors} is sized to the slice (native flat + * buffering, single segment). When {@code gpuPermit} is non-null the commit — which runs the GPU + * CAGRA build — is serialized on it while other segments' host-side ingest may proceed. + */ + private static void buildSegment( + Directory dir, + PrefetchingFbinReader reader, + float[] scratch, + int start, + int size, + boolean createNew, + Semaphore gpuPermit) + throws Exception { + Codec codec = codecFor(size); + + // Keep this slice in ONE segment: raise the doc-count flush threshold above the slice size and + // disable RAM-based flushing so nothing flushes before commit, and forbid merges. This is what + // makes native flat buffering valid. Order matters: enable the doc-count trigger BEFORE disabling + // the RAM trigger, since Lucene rejects a config where both are disabled at once. + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(Math.max(2, size + 1)) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE) + .setOpenMode( + createNew + ? IndexWriterConfig.OpenMode.CREATE + : IndexWriterConfig.OpenMode.APPEND); + + try (IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < size; i++) { + int id = start + i; + reader.get(id, scratch); // sequential access -> served from the prefetched chunk, no alloc + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, Integer.toString(id), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, scratch, EUCLIDEAN)); + writer.addDocument(doc); // copies the vector -> 'scratch' is safe to reuse next iteration + } + // The single flush at commit is where the GPU CAGRA build runs; serialize it if asked. + if (gpuPermit != null) { + gpuPermit.acquire(); + try { + writer.commit(); + } finally { + gpuPermit.release(); + } + } else { + writer.commit(); + } + } + } + + /** + * Combines the per-segment indexes into {@code indexDirPath} by hardlinking their files (same + * filesystem) rather than copying the vector data. {@link HardlinkCopyDirectoryWrapper} falls back + * to a byte copy automatically if the segment dirs and the final dir are on different filesystems. + */ + private static void combineByHardlink(Path indexDirPath, List segDirs) throws IOException { + Directory[] sources = new Directory[segDirs.size()]; + try { + for (int i = 0; i < segDirs.size(); i++) { + sources[i] = FSDirectory.open(segDirs.get(i)); + } + IndexWriterConfig iwc = + new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE); // keep segments separate + try (Directory target = new HardlinkCopyDirectoryWrapper(FSDirectory.open(indexDirPath)); + IndexWriter combiner = new IndexWriter(target, iwc)) { + combiner.addIndexes(sources); + } + } finally { + for (Directory s : sources) { + if (s != null) { + s.close(); + } + } + } + } + + /** Builds a codec with all knobs on, sizing the native flat buffer to {@code numInputVectors}. */ + private static Codec codecFor(int numInputVectors) throws Exception { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) // auto-tune graph params + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.AUTO_SELECT) // pick algo by dataset size + .withNumInputVectors(numInputVectors) // native flat buffering (single segment) + .build(); + return new Lucene101AcceleratedHNSWCodec(params); + } + + /** Splits {@code total} into {@code k} contiguous [start, size] slices, spreading the remainder. */ + private static List sliceEvenly(int total, int k) { + List slices = new ArrayList<>(); + int base = total / k; + int rem = total % k; + int start = 0; + for (int p = 0; p < k; p++) { + int size = base + (p < rem ? 1 : 0); // spread the remainder over the first slices + if (size <= 0) { + continue; + } + slices.add(new int[] {start, size}); + start += size; + } + return slices; + } + + /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ + private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { + float[] queryVector; + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, 1)) { + queryVector = reader.get(0); + } + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + log.info("Sample search returned " + results.scoreDocs.length + " hits:"); + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc sd = results.scoreDocs[i]; + String id = searcher.storedFields().document(sd.doc).get(ID_FIELD); + log.info(" rank " + (i + 1) + ": id=" + id + " score=" + sd.score); + } + } + } + + /** Writes a small random {@code .fbin} so the example is runnable without external data. */ + private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) + throws IOException { + ByteBuffer buf = + ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] + buf.putInt(dim); + for (int i = 0; i < numVectors; i++) { + for (int j = 0; j < dim; j++) { + buf.putFloat(random.nextFloat() * 100); + } + } + buf.flip(); + try (FileChannel ch = + FileChannel.open( + path, + StandardOpenOption.CREATE, + StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING)) { + while (buf.hasRemaining()) { + ch.write(buf); + } + } + } + + /** + * Prefetching, double-buffered sequential reader for uncompressed {@code .fbin} files ({@code + * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). + * + *

Opens the file ONCE. A background thread reads the (optionally sliced) range front-to-back + * into two reusable direct buffers: while the caller consumes the current chunk, the reader fills + * the next one, so the disk read overlaps with the caller's per-vector work. {@link #get(int, + * float[])} unpacks directly into a caller-supplied array (no per-vector allocation). + * + *

Forward-only, single-consumer. {@code get} must be called with non-decreasing indices + * from a single thread — the intended pattern for streaming ingestion. To read only a slice (e.g. + * one segment of a partitioned build), construct it with a {@code [firstVector, count)} range so + * each segment streams just its own portion of the file. + */ + static final class PrefetchingFbinReader implements AutoCloseable { + + private static final long HEADER_BYTES = 8; + + private static final class Chunk { + final ByteBuffer buf; + final long start; + final int len; + + Chunk(ByteBuffer buf, long start, int len) { + this.buf = buf; + this.start = start; + this.len = len; + } + } + + /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ + private static final Chunk POISON = new Chunk(null, -1, 0); + + private final FileChannel channel; + private final int dimension; + private final int firstVector; // absolute index of the first vector this reader serves + private final int endVector; // absolute index just past the last vector this reader serves + private final int vectorBytes; + private final int chunkVectors; + + private final BlockingQueue free = new ArrayBlockingQueue<>(2); + private final BlockingQueue ready = new ArrayBlockingQueue<>(2); + private final Thread reader; + private volatile IOException readerError; + + private Chunk current; // consumer-owned; the chunk currently being served + + /** Reads the whole file from index 0. */ + PrefetchingFbinReader(Path path, int chunkSizeMB) throws IOException { + this(path, 0, -1, chunkSizeMB); + } + + /** + * Reads the contiguous range {@code [firstVector, firstVector + count)}, or to end of file if + * {@code count <= 0}. {@link #get} then serves absolute file indices within that range. + */ + PrefetchingFbinReader(Path path, int firstVector, int count, int chunkSizeMB) + throws IOException { + this.channel = FileChannel.open(path, StandardOpenOption.READ); + ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + readFully(header, 0); + header.flip(); + int numVectors = header.getInt(); + this.dimension = header.getInt(); + this.vectorBytes = dimension * Float.BYTES; + this.firstVector = firstVector; + this.endVector = + count > 0 ? (int) Math.min((long) firstVector + count, numVectors) : numVectors; + + long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; + int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int + this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); + + // Two reusable direct buffers: the reader fills one while the consumer drains the other. + for (int i = 0; i < 2; i++) { + free.add( + ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); + } + + this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); + this.reader.setDaemon(true); + this.reader.start(); + } + + int size() { + return endVector - firstVector; + } + + int dimension() { + return dimension; + } + + /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ + private void readLoop() { + long next = firstVector; + try { + while (next < endVector) { + ByteBuffer buf = free.take(); + int toRead = (int) Math.min(chunkVectors, endVector - next); + buf.clear(); + buf.limit(toRead * vectorBytes); + readFully(buf, HEADER_BYTES + next * (long) vectorBytes); + ready.put(new Chunk(buf, next, toRead)); + next += toRead; + } + ready.put(POISON); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); // close() requested; stop quietly + } catch (IOException e) { + readerError = e; + try { + ready.put(POISON); // unblock the consumer so it can observe the error + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } + } + + private void advance() throws IOException { + if (current != null) { + try { + free.put(current.buf); // hand the drained buffer back to the reader + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted returning chunk buffer", e); + } + current = null; + } + Chunk next; + try { + next = ready.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted awaiting next chunk", e); + } + if (next == POISON) { + if (readerError != null) { + throw new IOException("Prefetch reader failed", readerError); + } + throw new IOException("No more chunks available (unexpected EOF in prefetch)"); + } + current = next; + } + + /** Fills {@code dst} with the vector at absolute {@code index} (no allocation). */ + void get(int index, float[] dst) throws IOException { + if (index < firstVector || index >= endVector) { + throw new IndexOutOfBoundsException( + "Index " + index + " out of bounds [" + firstVector + ", " + endVector + ")"); + } + while (current == null || index >= current.start + current.len) { + advance(); + } + if (index < current.start) { + throw new UnsupportedOperationException( + "PrefetchingFbinReader requires forward-only sequential access; got index " + + index + + " before current chunk start " + + current.start); + } + int base = (int) (index - current.start) * vectorBytes; + for (int i = 0; i < dimension; i++) { + dst[i] = current.buf.getFloat(base + i * Float.BYTES); + } + } + + /** Convenience allocating variant (e.g. for a one-off query vector). */ + float[] get(int index) throws IOException { + float[] dst = new float[dimension]; + get(index, dst); + return dst; + } + + private void readFully(ByteBuffer buf, long position) throws IOException { + long pos = position; + while (buf.hasRemaining()) { + int n = channel.read(buf, pos); + if (n < 0) { + throw new IOException("Unexpected EOF reading at position " + pos); + } + pos += n; + } + } + + @Override + public void close() throws IOException { + reader.interrupt(); + channel.close(); + } + } +} diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java deleted file mode 100644 index 559159005a..0000000000 --- a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedFbinIngestExample.java +++ /dev/null @@ -1,372 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene.examples; - -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; -import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.nio.file.StandardOpenOption; -import java.util.Random; -import java.util.UUID; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; -import java.util.logging.Logger; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; - -/** - * Reference pattern for efficiently ingesting a LARGE {@code .fbin} vector file into an accelerated - * HNSW index. cuvs-lucene is a Lucene codec and sits below {@code addDocument}, so how you read your - * source data is application code — this example shows the four properties that keep ingestion from - * bottlenecking the GPU build: - * - *

    - *
  1. Open once, read sequentially. Reopening/seeking/closing the file per vector ("fd - * churn") is orders of magnitude slower than sequential reads. {@link PrefetchingFbinReader} - * opens the file ONCE and reads it front-to-back in large sequential chunks. - *
  2. Bounded memory. Pre-loading the whole file onto the JVM heap holds a redundant copy - * on top of Lucene's own per-segment buffer (~2x peak). This reader holds at most two chunks. - *
  3. Overlap read with consumption. A background reader thread fills the NEXT chunk (into a - * second buffer) while the ingest thread is still feeding the current chunk into {@code - * addDocument}, so the sequential disk read is hidden behind the per-document indexing work - * instead of serializing in front of it. - *
  4. Reuse the vector array. {@link PrefetchingFbinReader#get(int, float[])} unpacks - * directly into a caller-supplied array, avoiding a fresh {@code float[]} allocation per - * vector. This is safe because Lucene copies the vector value eagerly inside {@code - * addDocument} — the array may be reused as soon as {@code addDocument} returns. - *
- * - *

Adapt {@link PrefetchingFbinReader} to your own source (a DB, object store, or stream). The - * properties above are what matter, not the {@code .fbin} specifics. - * - *

Usage: {@code OptimizedFbinIngestExample [] []}. With no arguments a - * small demo {@code .fbin} is generated and indexed. - */ -public class OptimizedFbinIngestExample { - - private static final Logger log = Logger.getLogger(OptimizedFbinIngestExample.class.getName()); - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - public static void main(String[] args) throws Exception { - int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; - Path indexDirPath = Paths.get(UUID.randomUUID().toString()); - - Path fbinPath; - boolean generated = false; - if (args.length >= 1) { - fbinPath = Paths.get(args[0]); - } else { - fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); - writeDemoFbin(fbinPath, 5000, 32, new Random(222)); - generated = true; - log.info("No .fbin provided; generated a demo file at " + fbinPath); - } - - try { - buildIndex(fbinPath, indexDirPath, chunkSizeMB); - runSampleSearch(indexDirPath, fbinPath, 5); - } finally { - FileUtils.deleteDirectory(indexDirPath.toFile()); - if (generated) { - Files.deleteIfExists(fbinPath); - } - } - } - - /** - * Builds a single-segment accelerated HNSW index, streaming vectors from the {@code .fbin} in - * sequential prefetched chunks (never holding more than two chunks in memory) and reusing a single - * vector array across all documents. - */ - private static void buildIndex(Path fbinPath, Path indexDirPath, int chunkSizeMB) - throws Exception { - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params); - - try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, chunkSizeMB)) { - int n = reader.size(); - - // Keep the whole dataset in ONE segment: disable RAM-based flushing and raise the doc-count - // flush threshold above the document count so nothing flushes before commit. (This only - // controls Lucene's segment cadence; the vectors are not held here — they are streamed one - // at a time from the chunked reader into addDocument.) - // Order matters: enable the doc-count flush trigger BEFORE disabling the RAM trigger, since - // Lucene rejects a config where both are disabled at once. - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(Math.max(2, n + 1)) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - - log.info( - "Indexing " - + n - + " vectors (" - + reader.dimension() - + "-dim) from " - + fbinPath - + " using " - + chunkSizeMB - + " MB prefetched sequential chunks"); - - // One reusable array for the whole build — refilled per vector, copied eagerly by Lucene. - float[] vector = new float[reader.dimension()]; - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < n; i++) { - reader.get(i, vector); // sequential access -> served from the prefetched chunk, no alloc - Document doc = new Document(); - doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - doc.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); - writer.addDocument(doc); // copies the vector -> 'vector' is safe to reuse next iteration - } - writer.commit(); // single flush -> single segment; the GPU CAGRA build happens here - } - log.info("Index build complete: " + indexDirPath); - } - } - - /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ - private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { - float[] queryVector; - try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, 1)) { - queryVector = reader.get(0); - } - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); - log.info("Sample search returned " + results.scoreDocs.length + " hits:"); - for (int i = 0; i < results.scoreDocs.length; i++) { - ScoreDoc sd = results.scoreDocs[i]; - String id = searcher.storedFields().document(sd.doc).get(ID_FIELD); - log.info(" rank " + (i + 1) + ": id=" + id + " score=" + sd.score); - } - } - } - - /** Writes a small random {@code .fbin} so the example is runnable without external data. */ - private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) - throws IOException { - ByteBuffer buf = - ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); - buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] - buf.putInt(dim); - for (int i = 0; i < numVectors; i++) { - for (int j = 0; j < dim; j++) { - buf.putFloat(random.nextFloat() * 100); - } - } - buf.flip(); - try (FileChannel ch = - FileChannel.open( - path, - StandardOpenOption.CREATE, - StandardOpenOption.WRITE, - StandardOpenOption.TRUNCATE_EXISTING)) { - while (buf.hasRemaining()) { - ch.write(buf); - } - } - } - - /** - * Prefetching, double-buffered sequential reader for uncompressed {@code .fbin} files ({@code - * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). - * - *

Opens the file ONCE. A background thread reads the file front-to-back into two reusable direct - * buffers: while the caller consumes the current chunk, the reader fills the next one, so the disk - * read overlaps with the caller's per-vector work. {@link #get(int, float[])} unpacks directly into - * a caller-supplied array (no per-vector allocation). - * - *

Forward-only, single-consumer. {@code get} must be called with non-decreasing indices - * from a single thread — the intended pattern for streaming ingestion. - */ - static final class PrefetchingFbinReader implements AutoCloseable { - - private static final long HEADER_BYTES = 8; - - private static final class Chunk { - final ByteBuffer buf; - final long start; - final int len; - - Chunk(ByteBuffer buf, long start, int len) { - this.buf = buf; - this.start = start; - this.len = len; - } - } - - /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ - private static final Chunk POISON = new Chunk(null, -1, 0); - - private final FileChannel channel; - private final int dimension; - private final int vectorCount; - private final int vectorBytes; - private final int chunkVectors; - - private final BlockingQueue free = new ArrayBlockingQueue<>(2); - private final BlockingQueue ready = new ArrayBlockingQueue<>(2); - private final Thread reader; - private volatile IOException readerError; - - private Chunk current; // consumer-owned; the chunk currently being served - - PrefetchingFbinReader(Path path, int chunkSizeMB) throws IOException { - this.channel = FileChannel.open(path, StandardOpenOption.READ); - ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); - readFully(header, 0); - header.flip(); - this.vectorCount = header.getInt(); - this.dimension = header.getInt(); - this.vectorBytes = dimension * Float.BYTES; - - long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; - int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int - this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); - - // Two reusable direct buffers: the reader fills one while the consumer drains the other. - for (int i = 0; i < 2; i++) { - free.add( - ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); - } - - this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); - this.reader.setDaemon(true); - this.reader.start(); - } - - int size() { - return vectorCount; - } - - int dimension() { - return dimension; - } - - /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ - private void readLoop() { - long next = 0; - try { - while (next < vectorCount) { - ByteBuffer buf = free.take(); - int toRead = (int) Math.min(chunkVectors, vectorCount - next); - buf.clear(); - buf.limit(toRead * vectorBytes); - readFully(buf, HEADER_BYTES + next * (long) vectorBytes); - ready.put(new Chunk(buf, next, toRead)); - next += toRead; - } - ready.put(POISON); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); // close() requested; stop quietly - } catch (IOException e) { - readerError = e; - try { - ready.put(POISON); // unblock the consumer so it can observe the error - } catch (InterruptedException ignored) { - Thread.currentThread().interrupt(); - } - } - } - - private void advance() throws IOException { - if (current != null) { - try { - free.put(current.buf); // hand the drained buffer back to the reader - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted returning chunk buffer", e); - } - current = null; - } - Chunk next; - try { - next = ready.take(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted awaiting next chunk", e); - } - if (next == POISON) { - if (readerError != null) { - throw new IOException("Prefetch reader failed", readerError); - } - throw new IOException("No more chunks available (unexpected EOF in prefetch)"); - } - current = next; - } - - /** Fills {@code dst} with the vector at {@code index} (no allocation). */ - void get(int index, float[] dst) throws IOException { - if (index < 0 || index >= vectorCount) { - throw new IndexOutOfBoundsException( - "Index " + index + " out of bounds [0, " + vectorCount + ")"); - } - while (current == null || index >= current.start + current.len) { - advance(); - } - if (index < current.start) { - throw new UnsupportedOperationException( - "PrefetchingFbinReader requires forward-only sequential access; got index " - + index - + " before current chunk start " - + current.start); - } - int base = (int) (index - current.start) * vectorBytes; - for (int i = 0; i < dimension; i++) { - dst[i] = current.buf.getFloat(base + i * Float.BYTES); - } - } - - /** Convenience allocating variant (e.g. for a one-off query vector). */ - float[] get(int index) throws IOException { - float[] dst = new float[dimension]; - get(index, dst); - return dst; - } - - private void readFully(ByteBuffer buf, long position) throws IOException { - long pos = position; - while (buf.hasRemaining()) { - int n = channel.read(buf, pos); - if (n < 0) { - throw new IOException("Unexpected EOF reading at position " + pos); - } - pos += n; - } - } - - @Override - public void close() throws IOException { - reader.interrupt(); - channel.close(); - } - } -} From 56bc26035a3ec95421cfffd1d5438c3ebfe146fc Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 12 Aug 2026 18:14:55 -0700 Subject: [PATCH 10/20] Expand OptimizedCagraHnswBuildExample and add early index-sort check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add enableRMMAsyncMemory() call with a note that it must not be used with CPU-only codecs - Expose the primary tuning knobs in codecFor(): withMaxConn/withBeamWidth (recall/graph-size), withCuvsDistanceType, and withWriterThreads (seeded from availableProcessors()) - Drop explicit withCagraGraphBuildAlgo(AUTO_SELECT) — it is the default under HEURISTIC; update class Javadoc accordingly - Expand withNumInputVectors comment to cover the exact-count constraint, the filtered-ingest fallback, and index-sorted segment incompatibility - In Lucene99AcceleratedHNSWVectorsWriter, detect index-sorted segments at construction time rather than failing later at flush --- .../OptimizedCagraHnswBuildExample.java | 38 ++++++++++++++----- .../Lucene99AcceleratedHNSWVectorsWriter.java | 5 +++ 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java index 9e93e99508..e9a6761471 100644 --- a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java @@ -6,9 +6,10 @@ import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; -import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; @@ -67,11 +68,11 @@ * the buffer the GPU build consumes instead of piling up as a {@code List} on the JVM * heap that must then be assembled. This removes the assembly copy and roughly halves peak host * memory. It requires a single-segment build (no flush before commit, no merge). - *

  • Automatic graph-build algorithm. {@code HEURISTIC} + {@code AUTO_SELECT} lets cuVS - * pick the CAGRA build algorithm by dataset size (NN_DESCENT below ~5M vectors, IVF_PQ at or - * above) and auto-tune its parameters. NN_DESCENT generally reaches higher recall but takes - * longer to build; IVF_PQ is faster but slightly lower recall. Force one explicitly only under - * expert guidance. + *
  • Automatic graph-build algorithm. {@code HEURISTIC} strategy defaults to + * {@code AUTO_SELECT}, letting cuVS pick the CAGRA build algorithm by dataset size (NN_DESCENT + * below ~5M vectors, IVF_PQ at or above) and auto-tune its parameters. NN_DESCENT generally + * reaches higher recall but takes longer to build; IVF_PQ is faster but slightly lower recall. + * Force one explicitly via {@code withCagraGraphBuildAlgo} only under expert guidance. *
  • Partitioned multi-segment build. Because native flat buffering is single-segment, "K * segments" means K independent single-segment builds over contiguous slices, combined at the * end. This is a deliberate memory/throughput lever — see the assumptions below. @@ -141,6 +142,9 @@ public static void main(String[] args) throws Exception { log.info("No .fbin provided; generated a demo file at " + fbinPath); } + // Must not be called when using a CPU-only Lucene codec: those paths never load the cuVS + // native library, so CuVSProvider resolves to UnsupportedProvider and throws. + CuVSProvider.provider().enableRMMAsyncMemory(); try { buildIndex(fbinPath, indexDirPath, chunkSizeMB, numSegments, overlap); runSampleSearch(indexDirPath, fbinPath, 5); @@ -367,9 +371,25 @@ private static void combineByHardlink(Path indexDirPath, List segDirs) thr private static Codec codecFor(int numInputVectors) throws Exception { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) // auto-tune graph params - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.AUTO_SELECT) // pick algo by dataset size - .withNumInputVectors(numInputVectors) // native flat buffering (single segment) + // HEURISTIC lets cuVS pick the build algorithm and auto-tune its parameters based on + // maxConn and beamWidth below. + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + // Primary recall/graph-size knobs. Higher values improve recall at the cost of a + // larger graph and longer build. Match to your dataset and recall target. + .withMaxConn(32) + .withBeamWidth(32) + // Must match the distance metric used when querying the index. + .withCuvsDistanceType(CuvsDistanceType.L2Expanded) + // Starting point: one thread per logical CPU. Profile and tune for your hardware. + .withWriterThreads(Runtime.getRuntime().availableProcessors()) + // Native flat buffering: the value MUST equal the number of vectors actually ingested + // into this segment; the writer fails fast if they differ. If some input vectors are + // excluded (e.g. filtered during ingest), either pre-count the survivors in a separate + // pass or omit withNumInputVectors (pass 0) to fall back to the heap-buffered path, + // which buffers all vectors in a List on the JVM heap before building and + // therefore uses more peak host memory. Index-sorted segments (IndexWriterConfig + // .setIndexSort) are also unsupported. + .withNumInputVectors(numInputVectors) .build(); return new Lucene101AcceleratedHNSWCodec(params); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index e1f707edaf..bea2ce0ff0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -110,6 +110,11 @@ public Lucene99AcceleratedHNSWVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.numInputVectors = acceleratedHNSWParams.getNumInputVectors(); this.nativeMode = numInputVectors > 0; + if (nativeMode && state.segmentInfo.getIndexSort() != null) { + throw new IllegalArgumentException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" + + " index-sorted segments; unset it (0) to use the heap-buffered path"); + } vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); From 8ffba097623861f103c4c4c886a03d354f491279 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 13 Aug 2026 07:12:47 -0700 Subject: [PATCH 11/20] Fix merge errors --- .../java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java | 2 +- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java | 6 +----- .../LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java | 1 + .../LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java | 1 + 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 50642b6142..a343a3149c 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -83,7 +83,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree * (its column count). Ceil is used to accommodate odd graph degrees. * 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[])}, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index bea2ce0ff0..fa02446dc0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -246,8 +246,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws vectorIndexLength, size, hnswGraph, - graphLevelNodeOffsets, - acceleratedHNSWParams.getGraphdegree()); + graphLevelNodeOffsets); cagraIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); @@ -300,10 +299,7 @@ private void writeFieldNative(FieldWriter fieldData, int maxDoc) throws IOExcept FieldInfo fieldInfo = fieldData.fieldInfo(); try { CuVSHostMatrix dataset = fieldData.getHostMatrix(); - long ts = StageTimers.start(); nativeFlat.writeField(fieldInfo, dataset, maxDoc, fieldData.getDocsWithFieldSet()); - StageTimers.stop( - "flat-write [DISK]", ts, (long) count * fieldInfo.getVectorDimension() * Float.BYTES); writeFieldInternal(fieldInfo, dataset); } finally { fieldData.releaseNativeBuffer(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 82cd4ed5ba..52c995b9f7 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -182,6 +182,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw adjacencyListMatrix, dataset, acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), params, QuantizationType.BINARY, acceleratedHNSWParams.getWriterThreads()); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 0a53bfbd15..fd558108a9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -207,6 +207,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE adjacencyListMatrix, dataset, acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), params, QuantizationType.SCALAR, acceleratedHNSWParams.getWriterThreads()); From cfd34409a201dad21a93a144566412a78d97249b Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 13 Aug 2026 09:06:53 -0700 Subject: [PATCH 12/20] Add a Lucene-version tripwire and round-trip test for NativeFlatVectorsWriter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NativeFlatVectorsWriter hand-transcribes Lucene 10.2.0's .vec/.vemf flat vector format so it can write directly from native memory, bypassing the per-vector FloatVectorValues indirection Lucene's own writer requires. That port had no automated safety net: a Lucene upgrade could silently change the format constants or write sequence with nothing catching it before files became unreadable by the stock Lucene99FlatVectorsReader. Add two tests, referenced from NativeFlatVectorsWriter's class javadoc: - TestNativeFlatVectorsWriterFormatConstants asserts the resolved lucene-core classpath version still matches the pinned 10.2.0 (via the jar's Specification-Version manifest attribute, not a pom.xml text parse). Fast and environment-independent, so it runs even without cuVS/GPU available. Fires on any version bump, by design — an upgrade always needs a manual re-verification pass regardless of whether it happens to touch the mirrored constants. - TestNativeFlatVectorsWriterRoundTrip positively verifies the format still works by building a single-segment index with numInputVectors set (native flat buffering) and asserting every vector round-trips byte-exact through the real, unmodified Lucene99FlatVectorsReader. Requires cuVS/GPU. Expand NativeFlatVectorsWriter's "on a Lucene upgrade" class javadoc into a concrete checklist that points at both tests. --- .../cuvs/lucene/NativeFlatVectorsWriter.java | 23 ++- ...ativeFlatVectorsWriterFormatConstants.java | 46 ++++++ .../TestNativeFlatVectorsWriterRoundTrip.java | 135 ++++++++++++++++++ 3 files changed, 201 insertions(+), 3 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterFormatConstants.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java index 31d1d67315..8c5adb08b4 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java @@ -32,7 +32,11 @@ * *

    Ported code — version-pinned. The dense float32 layout (format constants, header, meta * field order, and footer) is transcribed from Lucene 10.2.0, which must stay equal to the - * {@code lucene-core} version in {@code pom.xml}. Sources (tag {@code releases/lucene/10.2.0}): + * {@code lucene-core} version in {@code pom.xml}. This is enforced by + * {@code TestNativeFlatVectorsWriterFormatConstants}, which fails the build if the resolved + * {@code lucene-core} classpath version drifts from the pin below — update + * {@code PINNED_LUCENE_VERSION} there together with this class on a verified upgrade. Sources (tag + * {@code releases/lucene/10.2.0}): * *

      *
    • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat} — format constants: @@ -46,8 +50,21 @@ * on read via {@code CodecUtil.checkIndexHeader} (codec name + version range) plus the fixed meta * layout — so if the new version bumps {@code VERSION_CURRENT}, renames a codec, or changes the meta * field order, files written here will be silently incompatible and the stock - * {@code Lucene99FlatVectorsReader} will reject or misread them. Update the mirrored constants and - * the {@code writeField}/{@code writeMeta} sequence to match, or bind to Lucene's own writer. + * {@code Lucene99FlatVectorsReader} will reject or misread them. Concretely: + * + *
        + *
      1. Diff the new version's {@code Lucene99FlatVectorsFormat}/{@code Lucene99FlatVectorsWriter} + * sources against the ones linked above and mirror any changed constants and any change to + * the {@code writeField}/{@code writeMeta} byte sequence (header, meta field order, footer) + * here. This class writes directly from native memory to avoid the per-vector {@code + * FloatVectorValues} indirection Lucene's own writer requires — that's the reason to keep + * hand-porting the format rather than delegating to it. + *
      2. Bump {@code PINNED_LUCENE_VERSION} in {@code TestNativeFlatVectorsWriterFormatConstants} to + * clear the tripwire. + *
      3. Note: {@code TestNativeFlatVectorsWriterRoundTrip} positively verifies the result + * by building a small index with {@code numInputVectors} set and asserting every + * vector round-trips byte-exact through the stock {@code Lucene99FlatVectorsReader}. + *
      */ final class NativeFlatVectorsWriter implements Closeable { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterFormatConstants.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterFormatConstants.java new file mode 100644 index 0000000000..af8d2c004a --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterFormatConstants.java @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.junit.Assert.assertEquals; + +import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; +import org.junit.Test; + +/** + * Tripwire for {@link NativeFlatVectorsWriter}'s hand-transcribed Lucene 10.2.0 format. + * + *

      {@link NativeFlatVectorsWriter}'s class javadoc says its dense float32 layout (format + * constants, header, meta field order, footer) is transcribed from Lucene 10.2.0 and "must stay + * equal to the lucene-core version in pom.xml." This asserts that invariant against the resolved + * {@code lucene-core} jar actually on the classpath (its manifest {@code Specification-Version}, + * not a parse of {@code pom.xml}'s text), so a Lucene upgrade fails this test immediately rather + * than silently writing {@code .vec}/{@code .vemf} files the stock + * {@code Lucene99FlatVectorsReader} may no longer read correctly. + * + *

      This intentionally fires on any version change, not just ones that actually alter + * the format — per the class javadoc, every upgrade needs a manual re-verification pass against + * the new version's {@code Lucene99FlatVectorsFormat}/{@code Lucene99FlatVectorsWriter} sources. + */ +public class TestNativeFlatVectorsWriterFormatConstants { + + private static final String PINNED_LUCENE_VERSION = "10.2.0"; + + @Test + public void lucenePinMatchesResolvedClasspathVersion() { + String resolved = Lucene99FlatVectorsFormat.class.getPackage().getSpecificationVersion(); + assertEquals( + "lucene-core on the classpath is " + + resolved + + ", but NativeFlatVectorsWriter is pinned to " + + PINNED_LUCENE_VERSION + + " (per its class javadoc). Re-verify NativeFlatVectorsWriter against the new" + + " version's Lucene99FlatVectorsFormat/Lucene99FlatVectorsWriter sources, update the" + + " mirrored constants and write sequence if needed, then bump PINNED_LUCENE_VERSION" + + " here.", + PINNED_LUCENE_VERSION, + resolved); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java new file mode 100644 index 0000000000..de4bdf6659 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Positive round-trip check for {@link NativeFlatVectorsWriter}: builds a single-segment index + * with {@code AcceleratedHNSWParams.numInputVectors} set (native flat buffering) and reads the + * {@code .vec}/{@code .vemf} files back through the stock {@code Lucene99FlatVectorsReader}, + * asserting every vector round-trips byte-exact. + * + *

      This is the check called for by {@link NativeFlatVectorsWriter}'s "on a Lucene upgrade" + * class javadoc: {@link TestNativeFlatVectorsWriterFormatConstants} catches a version-pin drift, + * but only this test actually confirms the hand-transcribed format is still readable by Lucene's + * real reader. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatVectorsWriterRoundTrip extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + @Test + public void vectorsRoundTripThroughStockLucene99FlatVectorsReader() throws Exception { + int numDocs = 500; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + // Force everything into a single unsorted, unmerged flush: native flat buffering requires + // numInputVectors to equal the exact number of vectors landing in that one flush. + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + FloatVectorValues values = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(values); + assertEquals(numDocs, values.size()); + assertEquals(dimension, values.dimension()); + + int seen = 0; + KnnVectorValues.DocIndexIterator it = values.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + String id = leafReader.storedFields().document(doc).get(ID_FIELD); + float[] roundTripped = values.vectorValue(it.index()); + assertArrayEquals( + "vector for id=" + + id + + " did not round-trip byte-exact through the stock" + + " Lucene99FlatVectorsReader", + dataset[Integer.parseInt(id)], + roundTripped, + 0f); + seen++; + } + assertEquals("did not visit every vector", numDocs, seen); + } + } +} From 4b4a0d72eb268980de6754def9e1829fe04fa5a4 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 13 Aug 2026 13:59:34 -0700 Subject: [PATCH 13/20] Add equivalence test for writerThreads-parallelized graph construction/write Neither of the two writerThreads-tuned parallelizations had test coverage: materializing the CAGRA adjacency into NeighborArrays (GPUBuiltHnswGraph's constructor) and encoding level 0 to disk (AcceleratedHNSWUtils.writeGraph). Both are only supposed to reformat the same data across threads, with no change in output versus the serial path. Add TestWriterThreadsGraphEquivalence, which feeds a synthetic adjacency (CuVSMatrix.ofArray, matching the construction the higher-layer subset builder already uses) directly into both code paths at writerThreads=1 and writerThreads=4, and asserts identical results: matching per-node arc lists for the graph construction, and matching per-node byte-length offsets plus byte-identical encoded output for writeGraph. Using a synthetic adjacency rather than a real CAGRA build isolates these two parallelizations from CAGRA's own build-to-build variance, since writerThreads is also forwarded into the actual GPU build parameters. The synthetic dataset is sized just above PARALLEL_MIN_NODES (1 << 16, the threshold duplicated in both classes) so the parallel branches are actually exercised rather than silently falling through to serial. --- .../TestWriterThreadsGraphEquivalence.java | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java new file mode 100644 index 0000000000..fd4f9079f4 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java @@ -0,0 +1,154 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSMatrix; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.junit.Test; + +/** + * Verifies that {@code writerThreads > 1} produces the same result as the serial path for the + * two parallelizations gated on {@code AcceleratedHNSWUtils}/{@code GPUBuiltHnswGraph}'s {@code + * PARALLEL_MIN_NODES} threshold: materializing the CAGRA adjacency into {@code NeighborArray}s + * (the {@code GPUBuiltHnswGraph} constructor), and encoding level 0 to disk ({@code + * AcceleratedHNSWUtils#writeGraph}). + * + *

      Both tests use a synthetic adjacency ({@link CuVSMatrix#ofArray(int[][])}, the same host-matrix + * construction the higher-layer subset builder already uses) rather than a real CAGRA build, so + * that the comparison isolates these two parallelizations from CAGRA's own build-to-build + * variance -- a real GPU build is not guaranteed to produce the identical graph twice even with the + * same input and thread count, which would make an end-to-end build comparison unreliable for this + * purpose. This needs cuVS/GPU only to allocate the host matrix, not to run a build. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestWriterThreadsGraphEquivalence extends LuceneTestCase { + + // Must be >= PARALLEL_MIN_NODES (1 << 16) in both AcceleratedHNSWUtils and + // GPUBuiltHnswGraph, or the "parallel" runs below silently fall through to the serial branch and + // the test would pass without exercising anything. + private static final int NUM_NODES = (1 << 16) + 1000; + private static final int DEGREE = 12; + private static final int NUM_THREADS = 4; + + @Test + public void fillNeighborArrayParallelMatchesSerial() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); + + try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency)) { + GPUBuiltHnswGraph serial = newSingleLayerGraph(matrix, 1); + GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); + assertGraphsEqual(serial, parallel); + } + } + + @Test + public void writeGraphParallelMatchesSerial() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); + + try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency); + Directory dir = new ByteBuffersDirectory()) { + // Materialize once, serially, so any difference found below is attributable only to + // writeGraph's own parallelization, not to fillNeighborArray's. + GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, 1); + + int[][] serialOffsets; + try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { + serialOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, 1); + } + int[][] parallelOffsets; + try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { + parallelOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, NUM_THREADS); + } + + assertEquals(serialOffsets.length, parallelOffsets.length); + for (int level = 0; level < serialOffsets.length; level++) { + assertArrayEquals( + "per-node byte-length offsets differ for level " + level, + serialOffsets[level], + parallelOffsets[level]); + } + + assertArrayEquals( + "writeGraph's parallel level-0 encoding produced different bytes than the serial path", + readAllBytes(dir, "serial"), + readAllBytes(dir, "parallel")); + } + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) { + // A single layer (layer 0 only): the constructor never consults layerNodes in that case, so + // the placeholder null entry mirrors the convention used elsewhere for "layer 0 needs no node + // list" without actually being read. + return new GPUBuiltHnswGraph( + NUM_NODES, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(layer0Adjacency), + numThreads); + } + + /** Every node/level's in-order arc list must match exactly between the two graphs. */ + private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { + assertEquals(a.numLevels(), b.numLevels()); + for (int level = 0; level < a.numLevels(); level++) { + int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); + for (int node : nodes) { + assertArrayEquals( + "node " + node + " at level " + level + " has different neighbors", + arcsOf(a, level, node), + arcsOf(b, level, node)); + } + } + } + + private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { + graph.seek(level, node); + List arcs = new ArrayList<>(); + for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { + arcs.add(n); + } + return arcs.stream().mapToInt(Integer::intValue).toArray(); + } + + private static byte[] readAllBytes(Directory dir, String name) throws Exception { + try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { + byte[] bytes = new byte[(int) in.length()]; + in.readBytes(bytes, 0, bytes.length); + return bytes; + } + } + + /** + * A deterministic, seeded pseudo-adjacency. It doesn't need to be a real CAGRA graph -- only a + * realistic shape (fixed degree, valid node ids) -- since {@code GPUBuiltHnswGraph} and {@code + * AcceleratedHNSWUtils#writeGraph} don't interpret the neighbor ids semantically. + */ + private static int[][] randomAdjacency(int numNodes, int degree, Random random) { + int[][] adjacency = new int[numNodes][degree]; + for (int[] row : adjacency) { + for (int j = 0; j < degree; j++) { + row[j] = random.nextInt(numNodes); + } + } + return adjacency; + } +} From 64ee08ddd6e8c1b541d23c77e417f4b6baa114c2 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 13 Aug 2026 15:59:40 -0700 Subject: [PATCH 14/20] Add functional coverage for native flat buffering (numInputVectors) TestNativeFlatVectorsWriterRoundTrip only verifies the flat .vec file round-trips through Lucene's reader -- it doesn't cover whether a natively-buffered index actually searches correctly, tolerates deletions, composes with the odd-graph-degree fix, or correctly rejects the usage patterns it doesn't support. None of that had test coverage before. Add two classes: - TestNativeFlatBufferingIndexAndSearch: end-to-end build + KNN search through the native-flat-buffered path; deletion applied after a natively buffered flush (with the javadoc spelling out that Lucene-level deletion is orthogonal to, and not caught by, the count-mismatch guard -- deletion only updates liveDocs and never touches FieldWriter's native buffer); and the M = ceil(cagraGraphDegree / 2) odd-degree fix combined with numInputVectors, a distinct call path (writeFieldNative) from the one TestAcceleratedHNSWOddGraphDegree exercises. - TestNativeFlatBufferingGuardRails: the three guard rails in Lucene99AcceleratedHNSWVectorsWriter around numInputVectors -- a count mismatch (simulating an ingest-time filter that skips addValue for some docs, as distinct from a delete), an index-sorted segment, and a merge attempt. The index-sort and count checks fire at different points (writer construction on the first addDocument, vs. flush), so each is asserted around the call that actually triggers it. --- .../TestNativeFlatBufferingGuardRails.java | 221 +++++++++++++++++ ...TestNativeFlatBufferingIndexAndSearch.java | 225 ++++++++++++++++++ 2 files changed, 446 insertions(+) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java new file mode 100644 index 0000000000..2645128419 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java @@ -0,0 +1,221 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.SerialMergeScheduler; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.search.Sort; +import org.apache.lucene.search.SortField; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Negative-path coverage for the three guard rails in {@code Lucene99AcceleratedHNSWVectorsWriter} + * around {@code AcceleratedHNSWParams.numInputVectors} (native flat buffering): a count mismatch, + * an index-sorted segment, and a merge attempt. None of these had test coverage before. + * + *

      Positive-path coverage (does a natively-buffered index actually search correctly, tolerate + * deletions, etc.) lives separately in {@link TestNativeFlatBufferingIndexAndSearch}. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatBufferingGuardRails extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + /** + * The count-mismatch guard exists for a caller-side bookkeeping error: declaring {@code + * numInputVectors} against the pre-filter document count instead of the number of vectors that + * actually reach {@code addValue} (e.g. an ingest-time filter skips some documents' vector + * field). It is unrelated to, and not triggered by, Lucene-level deletion -- see the javadoc on + * {@link TestNativeFlatBufferingIndexAndSearch#testDeletedDocsAfterNativeFlatBufferedFlush}. + */ + @Test + public void testCountMismatchFromIngestTimeFilterIsRejected() throws Exception { + int declaredNumInputVectors = 100; + int actuallyIndexed = declaredNumInputVectors - 1; // one doc "filtered out" before addValue + int dimension = 32; + float[][] dataset = generateDataset(random, actuallyIndexed, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(declaredNumInputVectors).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(declaredNumInputVectors + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < actuallyIndexed; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + IllegalStateException thrown = expectThrows(IllegalStateException.class, writer::commit); + assertTrue( + "unexpected message: " + thrown.getMessage(), + thrown.getMessage().contains("numInputVectors")); + } + } + + /** Native flat buffering pre-sizes a single flush's buffer and cannot support sorted flushes. */ + @Test + public void testIndexSortedSegmentIsRejected() throws Exception { + int numDocs = 50; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + Sort indexSort = new Sort(new SortField("sort_key", SortField.Type.LONG)); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setIndexSort(indexSort) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + // KnnVectorsFormat#fieldsWriter (and so the writer's index-sort check in its constructor) + // is invoked on the first addDocument() for the segment, not at commit() -- so the guard + // must be expected around the whole indexing loop, not just the flush. + IllegalArgumentException thrown = + expectThrows( + IllegalArgumentException.class, + () -> { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new NumericDocValuesField("sort_key", numDocs - i)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + }); + assertTrue( + "unexpected message: " + thrown.getMessage(), + thrown.getMessage().contains("index-sorted")); + } + } + + /** + * Native flat buffering supports only the unsorted single-segment flush path; merging two + * natively-buffered segments must be rejected rather than silently mis-sizing the native buffer. + */ + @Test + public void testMergeIsRejected() throws Exception { + int segmentSize = 40; + int dimension = 32; + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(segmentSize).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(segmentSize + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + // Keep the two flushes below as separate segments; NoMergePolicy also blocks forced + // merges, so it is swapped out before forceMerge() is called. + .setMergePolicy(NoMergePolicy.INSTANCE) + // Force merges to run synchronously on the calling thread, so the guard's exception + // (or whatever IndexWriter/SegmentMerger wraps it as) surfaces directly from + // forceMerge() instead of on a background merge thread. + .setMergeScheduler(new SerialMergeScheduler()); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + addSegment(writer, 0, segmentSize, dimension); + writer.commit(); // segment 1: exactly segmentSize vectors, matching numInputVectors + addSegment(writer, segmentSize, segmentSize, dimension); + writer.commit(); // segment 2: exactly segmentSize vectors, matching numInputVectors + + writer.getConfig().setMergePolicy(new TieredMergePolicy()); + Throwable thrown = expectThrows(Throwable.class, () -> writer.forceMerge(1)); + assertTrue( + "expected UnsupportedOperationException somewhere in the cause chain of: " + thrown, + causedBy(thrown, UnsupportedOperationException.class)); + } + } + + private void addSegment(IndexWriter writer, int startId, int count, int dimension) + throws Exception { + float[][] dataset = generateDataset(random, count, dimension); + for (int i = 0; i < count; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + } + + private static boolean causedBy(Throwable t, Class type) { + for (Throwable cur = t; cur != null; cur = cur.getCause()) { + if (type.isInstance(cur)) { + return true; + } + for (Throwable suppressed : cur.getSuppressed()) { + if (causedBy(suppressed, type)) { + return true; + } + } + } + return false; + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java new file mode 100644 index 0000000000..1829e4db62 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java @@ -0,0 +1,225 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Positive-path functional coverage for native flat buffering ({@code + * AcceleratedHNSWParams.numInputVectors}) beyond {@code TestNativeFlatVectorsWriterRoundTrip}, + * which only checks that the flat {@code .vec} file round-trips -- not that the resulting index is + * actually searchable, tolerates deletions, or composes correctly with the odd-graph-degree fix. + * + *

      The negative/guard-rail paths (count mismatch, index-sorted segments, merges) are covered + * separately in {@link TestNativeFlatBufferingGuardRails}. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatBufferingIndexAndSearch extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + /** A natively-buffered index must still be searchable through the normal Lucene KNN query API. */ + @Test + public void testIndexAndSearch() throws Exception { + int numDocs = 500; + int dimension = 32; + int topK = 10; + float[][] dataset = generateDataset(random, numDocs, dimension); + + buildNativeFlatBufferedIndex(numDocs, dimension, dataset); + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateDataset(random, 1, dimension)[0]; + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + + assertEquals("expected topK results", topK, results.scoreDocs.length); + for (var scoreDoc : results.scoreDocs) { + String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); + int idValue = Integer.parseInt(id); + assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); + } + } + } + + /** + * Deletion is orthogonal to native flat buffering: {@code IndexWriter.deleteDocuments} only + * updates Lucene's liveDocs bitset at search time -- it never touches {@code FieldWriter} or the + * native host matrix, so it cannot trip (and isn't meant to be caught by) the count-mismatch + * guard rail tested in {@link TestNativeFlatBufferingGuardRails}. This test instead confirms that + * deletions applied after a natively-buffered flush are still honored correctly at search time. + */ + @Test + public void testDeletedDocsAfterNativeFlatBufferedFlush() throws Exception { + int numDocs = 300; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); // single natively-buffered flush: FieldWriter's count matches numDocs + + // Delete every 3rd doc. No new vectors are added, so this does not trigger another flush of + // the vector field and cannot interact with the numInputVectors hint. + for (int i = 0; i < numDocs; i += 3) { + writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(i))); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected the deletions to land in the same single segment", 1, reader.leaves().size()); + assertTrue("expected some deleted docs", reader.numDeletedDocs() > 0); + + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], numDocs), numDocs); + for (var scoreDoc : results.scoreDocs) { + String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); + assertNotEquals( + "deleted doc id=" + id + " was still returned by search", 0, Integer.parseInt(id) % 3); + } + } + } + + /** + * The M = ceil(cagraGraphDegree / 2) fix ({@link TestAcceleratedHNSWOddGraphDegree}) must also + * hold on the native-flat-buffered write path ({@code writeFieldNative}), which is a distinct + * call path from the heap-buffered one that test exercises. + */ + @Test + public void testOddGraphDegreeWithNativeFlatBuffering() throws Exception { + int numDocs = 200; + int dimension = 32; + int oddGraphDegree = 63; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withIntermediateGraphDegree(128) + .withGraphDegree(oddGraphDegree) + .withNumInputVectors(numDocs) + .build(); + + buildNativeFlatBufferedIndex(numDocs, dimension, dataset, params); + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); + assertEquals(5, results.scoreDocs.length); + } + } + + private void buildNativeFlatBufferedIndex(int numDocs, int dimension, float[][] dataset) + throws Exception { + buildNativeFlatBufferedIndex( + numDocs, + dimension, + dataset, + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build()); + } + + private void buildNativeFlatBufferedIndex( + int numDocs, int dimension, float[][] dataset, AcceleratedHNSWParams params) + throws Exception { + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); + } + } +} From 22e10fe018ffbe73fa5f07ed6d68c487b03fd708 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 14 Aug 2026 06:05:23 -0700 Subject: [PATCH 15/20] Fix merge-time vector-count bug causing intermittent EOF during search FloatVectorValues#size() on KnnVectorsWriter.MergedVectorValues's merged view is the raw sum of every source segment's on-disk vector count -- not the number of live (non-deleted) vectors the iterator actually yields (Lucene's own Lucene99FlatVectorsWriter never calls .size() for exactly this reason, only iterates). vectorBasedMerge used that count to pre-size CuVSMatrix.hostBuilder's fixed native buffer, so whenever a merge dropped deleted docs, the buffer ended up larger than the number of vectors actually written into it, and the resulting HNSW graph was built with more nodes than the flat vector file actually contains. Symptom: an intermittent EOFException in OffHeapFloatVectorValues during concurrent search, whenever traversal happened to reach one of the phantom out-of-bounds nodes. Fix: count live vectors via a throwaway iteration before allocating the host buffer, rather than trusting size() -- but only when at least one segment being merged actually has deletions (mergeState.liveDocs has a non-null entry). size() is trustworthy whenever no segment has deletions, since per-segment vector counts already exclude docs without a value for the field independent of deletions, so that common case (most merges have no prior deletes) skips the extra pass entirely and costs exactly what it did before this fix. Add TestMergedGraphOrdinalBounds: rather than depending on a concurrent search happening to traverse a bad node, this deterministically merges two segments with interspersed deletions and walks the entire resulting HNSW graph, asserting every neighbor ordinal is in-bounds relative to the merged segment's actual live vector count. --- .../Lucene99AcceleratedHNSWVectorsWriter.java | 42 +++- .../lucene/TestMergedGraphOrdinalBounds.java | 195 ++++++++++++++++++ 2 files changed, 236 insertions(+), 1 deletion(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index fa02446dc0..868e27c2ce 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -43,6 +43,7 @@ import org.apache.lucene.index.Sorter.DocMap; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Bits; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; @@ -373,9 +374,48 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { + // FloatVectorValues#size() on the merged view is the raw sum of every source segment's + // on-disk vector count (MergedVectorValues.MergedFloat32VectorValues computes it once at + // construction from each sub-reader's unfiltered size) -- NOT the number of live + // (non-deleted) vectors the iterator below will actually yield, which is what + // CuVSMatrix.hostBuilder needs since it preallocates a fixed-size native buffer. Using + // size() here under-fills that buffer whenever the merge drops deleted docs, leaving the + // graph built over more rows than were actually populated. + // + // size() IS trustworthy when no segment being merged has any deletions: per-segment vector + // counts already exclude docs without a value for this field (sparse fields are handled at + // the single-segment level, independent of deletions), so the raw sum equals the live count + // in that case and the extra counting pass below can be skipped. + boolean anySegmentHasDeletions = false; + for (Bits liveDocs : mergeState.liveDocs) { + if (liveDocs != null) { + anySegmentHasDeletions = true; + break; + } + } + + int size; + if (anySegmentHasDeletions) { + // Count the live vectors via a throwaway iteration first (mergeFloatVectorValues + // constructs a fresh, independent view each call, so this doesn't disturb the real build + // pass below). + size = 0; + FloatVectorValues counting = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + KnnVectorValues.DocIndexIterator countingIt = counting.iterator(); + for (int doc = countingIt.nextDoc(); + doc != DocIdSetIterator.NO_MORE_DOCS; + doc = countingIt.nextDoc()) { + size++; + } + } else { + size = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) + .size(); + } + 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); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java new file mode 100644 index 0000000000..1fd8d2a3af --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java @@ -0,0 +1,195 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import java.util.Random; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.junit.Test; + +/** + * Repro for the CI-observed {@code EOFException} in {@code OffHeapFloatVectorValues} during + * concurrent KNN search over an accelerated-HNSW index built via {@link + * TestAcceleratedHNSWDeletedDocuments}/{@link TestCuVSAcceleratedHNSWDeletedDocuments} (both: + * deletions + a real merge, heap-buffered path, no {@code numInputVectors}). + * + *

      Rather than relying on a random concurrent search happening to traverse a bad graph node + * (which only reproduced intermittently, on one CI node), this walks the entire merged + * HNSW graph directly and asserts every neighbor ordinal is within the merged segment's actual + * flat-vector count. This targets the suspected root cause: {@code + * Lucene99AcceleratedHNSWVectorsWriter#mergeOneField} derives the merged vector set twice, + * independently -- once via the real {@code flatVectorsWriter.mergeOneField} (the authoritative + * flat {@code .vec} file) and again via {@code vectorBasedMerge}'s own call to {@code + * KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues} to build the CAGRA/HNSW graph. If + * those two independently-derived views of "the merged, post-deletion vector set" ever disagree in + * count or ordinal order, the graph ends up referencing ordinals the flat file doesn't actually + * have, which is exactly what an out-of-bounds read (EOFException) during traversal would look + * like. + * + *

      This test is deterministic: it fails on any disagreement between the graph and the flat file, + * rather than depending on a search happening to reach the bad node. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestMergedGraphOrdinalBounds extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String FIELD = "vector"; + + @Test + public void testMergedGraphOrdinalsStayWithinFlatVectorBounds() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + + Random random = new Random(1234); + int segmentSize = 300; + int dimension = 32; + // Interspersed deletions on both segments, so the merge must drop a scattered subset of + // ordinals from each -- not just a contiguous prefix/suffix -- when it re-derives the merged + // vector set. + int deleteEveryNth = 4; + + Codec codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + IndexWriterConfig config = + new IndexWriterConfig().setCodec(codec).setMergePolicy(NoMergePolicy.INSTANCE); + + int expectedLiveVectors; + try (Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, config)) { + int deletedFromSegment1 = + addSegmentWithInterspersedDeletions(writer, 0, segmentSize, dimension, deleteEveryNth, random); + writer.commit(); // segment 1, alone + int deletedFromSegment2 = + addSegmentWithInterspersedDeletions( + writer, segmentSize, segmentSize, dimension, deleteEveryNth, random); + writer.commit(); // segment 2, alone + + expectedLiveVectors = 2 * segmentSize - deletedFromSegment1 - deletedFromSegment2; + + // NoMergePolicy blocks forced merges too, so swap it out now that the two segments (each + // with their own interspersed deletions already committed) are set up. + writer.getConfig().setMergePolicy(new TieredMergePolicy()); + writer.forceMerge(1); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected the forced merge to produce a single segment", 1, reader.leaves().size()); + LeafReader leaf = reader.leaves().get(0).reader(); + + FloatVectorValues flatValues = leaf.getFloatVectorValues(FIELD); + assertEquals( + "merged flat vector count should equal (added - deleted)", + expectedLiveVectors, + flatValues.size()); + + HnswGraph graph = graphOf(leaf); + int level0NodeCount = graph.getNodesOnLevel(0).size(); + assertEquals( + "HNSW graph's level-0 node count disagrees with the merged flat vector file's actual" + + " count -- the graph and the flat file were derived independently by" + + " vectorBasedMerge and flatVectorsWriter.mergeOneField and disagree", + flatValues.size(), + level0NodeCount); + + assertAllNeighborOrdinalsInBounds(graph, flatValues.size()); + } + } + } + + /** + * Every neighbor referenced anywhere in the graph, at every level, must be a valid ordinal into + * the merged segment's actual flat vector data -- otherwise a reader resolving that neighbor's + * vector (e.g. mid-search, to score it) reads past the end of the flat file. + */ + private static void assertAllNeighborOrdinalsInBounds(HnswGraph graph, int liveVectorCount) + throws Exception { + for (int level = 0; level < graph.numLevels(); level++) { + HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(level); + while (nodes.hasNext()) { + int node = nodes.nextInt(); + assertTrue( + "node " + node + " at level " + level + " is itself out of bounds (live vectors: " + + liveVectorCount + ")", + node >= 0 && node < liveVectorCount); + graph.seek(level, node); + for (int neighbor = graph.nextNeighbor(); neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue( + "node " + + node + + " at level " + + level + + " has a neighbor ordinal " + + neighbor + + " out of bounds for the merged segment's " + + liveVectorCount + + " live vectors", + neighbor >= 0 && neighbor < liveVectorCount); + } + } + } + } + + private static HnswGraph graphOf(LeafReader leaf) throws Exception { + KnnVectorsReader knnReader = ((CodecReader) leaf).getVectorReader(); + if (knnReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnReader = fieldsReader.getFieldReader(FIELD); + } + return ((HnswGraphProvider) knnReader).getGraph(FIELD); + } + + /** + * Adds {@code count} documents (global ids {@code [startId, startId + count)}), then deletes + * every {@code deleteEveryNth}-th one by id, scattering the deletions across the segment rather + * than leaving a contiguous surviving range. + * + * @return the number of documents deleted from this segment + */ + private static int addSegmentWithInterspersedDeletions( + IndexWriter writer, + int startId, + int count, + int dimension, + int deleteEveryNth, + Random random) + throws Exception { + float[][] dataset = generateDataset(random, count, dimension); + for (int i = 0; i < count; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); + document.add(new KnnFloatVectorField(FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + int deleted = 0; + for (int i = 0; i < count; i += deleteEveryNth) { + writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(startId + i))); + deleted++; + } + return deleted; + } +} From 725fc9d5b006a4ed53e7fe214e77d03972dfecc6 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 14 Aug 2026 06:32:40 -0700 Subject: [PATCH 16/20] Restore unintended M-derivation change in createMultiLayerHnswGraph --- .../main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java | 3 +-- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java | 1 - .../LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java | 1 - .../LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java | 1 - 4 files changed, 1 insertion(+), 5 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index a343a3149c..4db4d379c0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -97,14 +97,13 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, int hnswLayers, - int graphDegree, CagraIndexParams params, QuantizationType quantization, int numThreads) throws Throwable { int size = (int) vectorDataset.size(); - int M = graphDegree / 2; + int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); List layerNodes = new ArrayList<>(); List layerAdjacencies = new ArrayList<>(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 868e27c2ce..8fc2b7f17a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -232,7 +232,6 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws adjacencyListMatrix, dataset, acceleratedHNSWParams.getHnswLayers(), - acceleratedHNSWParams.getGraphdegree(), params, QuantizationType.NONE, acceleratedHNSWParams.getWriterThreads()); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 52c995b9f7..82cd4ed5ba 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -182,7 +182,6 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw adjacencyListMatrix, dataset, acceleratedHNSWParams.getHnswLayers(), - acceleratedHNSWParams.getGraphdegree(), params, QuantizationType.BINARY, acceleratedHNSWParams.getWriterThreads()); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index fd558108a9..0a53bfbd15 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -207,7 +207,6 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE adjacencyListMatrix, dataset, acceleratedHNSWParams.getHnswLayers(), - acceleratedHNSWParams.getGraphdegree(), params, QuantizationType.SCALAR, acceleratedHNSWParams.getWriterThreads()); From 9faae9192415aa4410a6a8107d070064f0172c49 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 14 Aug 2026 06:41:30 -0700 Subject: [PATCH 17/20] Restore unnecessary removals --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 4db4d379c0..fc73651be7 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -105,11 +105,12 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( int size = (int) vectorDataset.size(); int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); + // Store all layers data List layerNodes = new ArrayList<>(); List layerAdjacencies = new ArrayList<>(); // Layer 0: Use full CAGRA adjacency list - layerNodes.add(null); + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list layerAdjacencies.add(adjacencyListMatrix); int currentLayerSize = size; @@ -117,22 +118,28 @@ 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) { 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) { @@ -141,6 +148,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( for (int i = 0; i < nextLayerSize; i++) { vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } + + // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); @@ -151,16 +160,22 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( for (int i = 0; i < nextLayerSize; 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, numThreads); } @@ -230,6 +245,7 @@ private static CuVSMatrix buildCagraGraphForSubset( */ public static int[][] writeGraph( GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) throws IOException { + // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); int numLevels = graph.numLevels(); int[][] offsets = new int[numLevels][]; @@ -251,6 +267,7 @@ public static int[][] writeGraph( offsets[level] = new int[sortedNodes.length]; writeLevelSerial(graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0); } + // Return offsets (information written while writing the meta info) return offsets; } From 77cbc65c066ae1fac18434e026cfedee57f07d43 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 14 Aug 2026 07:39:12 -0700 Subject: [PATCH 18/20] Fix inconsistent handling of explicit CuVSIvfPqParams under HEURISTIC+IVF_PQ Under Strategy.HEURISTIC with an explicit cagraGraphBuildAlgo override, getIVFPQParams honored every other caller-supplied field (graphDegree, intermediateGraphDegree, writerThreads, cuvsDistanceType) but always discarded a caller-configured CuVSIvfPqParams in favor of the auto-tuned one -- an inconsistency within the same branch, since nothing signaled that this one field behaved differently from the others. Add AcceleratedHNSWParams.isCuVSIvfPqParamsExplicit(), tracked via a new Builder flag set in withCuVSIvfPqParams(...), so getIVFPQParams can honor an explicitly-set value and only auto-tune when the caller left it unset. Also, while auditing this file: - Remove ALGO_SWITCH_THRESHOLD, dead since both create() overloads now fully delegate algorithm selection to cuVS's own fromDataset/ fromHnswParams heuristics rather than a local row-count threshold. - Restore getCuVSIvfPqParams' and roundUpSafe's original source-attribution comments (lost when getCuVSIvfPqParams was reintroduced in a prior commit after having been deleted upstream), rephrasing the ivf_pq_params reference to name the struct instead of a line-numbered URL that will drift from the real source. Drop roundUpSafe's comment; the arithmetic doesn't need it. - Add a TODO noting cuvs-java has no standalone binding for the C++ ivf_pq_params(dataset_extents, metric) constructor this method translates -- only fromHnswParams/fromDataset, which bundle algorithm selection with parameter tuning and can't be used to auto-tune IVF-PQ while pinning a different algorithm choice. --- .../cuvs/lucene/AcceleratedHNSWParams.java | 19 +++++++++++ .../cuvs/lucene/CagraIndexParamsFactory.java | 32 ++++++++++++++++--- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 67f69289df..873c1767db 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -86,6 +86,7 @@ public static enum Strategy { private final int beamWidth; private final CagraGraphBuildAlgo cagraGraphBuildAlgo; private final CuVSIvfPqParams cuVSIvfPqParams; + private final boolean cuVSIvfPqParamsExplicit; private final int numMergeWorkers; private final ExecutorService mergeExec; private final Strategy strategy; @@ -105,6 +106,9 @@ public static enum Strategy { * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism. * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. + * @param cuVSIvfPqParamsExplicit whether cuVSIvfPqParams was set explicitly by the caller, as + * opposed to defaulted; consulted under HEURISTIC with an explicit IVF_PQ override so a + * caller-supplied value is honored instead of silently replaced by the auto-tuned one. * @param numMergeWorkers The number of merge workers to use with the fallback mechanism. * @param mergeExec The instance of {@link ExecutorService} to use with the fallback mechanism. * @param strategy either HEURISTIC [Default] that delegates the CAGRA build parameters to cuVS (derived from the HNSW-equivalent maxConn and beamWidth) or CUSTOM that uses the parameters passed through this class. @@ -122,6 +126,7 @@ private AcceleratedHNSWParams( int beamWidth, CagraGraphBuildAlgo cagraGraphBuildAlgo, CuVSIvfPqParams cuVSIvfPqParams, + boolean cuVSIvfPqParamsExplicit, int numMergeWorkers, ExecutorService mergeExec, Strategy strategy, @@ -138,6 +143,7 @@ private AcceleratedHNSWParams( this.beamWidth = beamWidth; this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; this.cuVSIvfPqParams = cuVSIvfPqParams; + this.cuVSIvfPqParamsExplicit = cuVSIvfPqParamsExplicit; this.numMergeWorkers = numMergeWorkers; this.mergeExec = mergeExec; this.strategy = strategy; @@ -219,6 +225,16 @@ public CuVSIvfPqParams getCuVSIvfPqParams() { return cuVSIvfPqParams; } + /** + * Whether {@link #getCuVSIvfPqParams()} was set explicitly via {@link + * Builder#withCuVSIvfPqParams(CuVSIvfPqParams)}, as opposed to defaulted. + * + * @return true if the caller explicitly set cuVSIvfPqParams + */ + public boolean isCuVSIvfPqParamsExplicit() { + return cuVSIvfPqParamsExplicit; + } + /** * Get the number of merge workers set to be used in the fallback mechanism * @@ -337,6 +353,7 @@ public static class Builder { private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO; private int numMergeWorkers = DEFAULT_NUM_MERGE_WORKERS; private CuVSIvfPqParams cuVSIvfPqParams = null; + private boolean cuVSIvfPqParamsExplicit = false; private ExecutorService mergeExec = null; private Strategy strategy = DEFAULT_STRATEGY; private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE; @@ -442,6 +459,7 @@ public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo) */ public Builder withCuVSIvfPqParams(CuVSIvfPqParams cuVSIvfPqParams) { this.cuVSIvfPqParams = cuVSIvfPqParams; + this.cuVSIvfPqParamsExplicit = true; return this; } @@ -654,6 +672,7 @@ public AcceleratedHNSWParams build() { beamWidth, cagraGraphBuildAlgo, cuVSIvfPqParams, + cuVSIvfPqParamsExplicit, numMergeWorkers, mergeExec, strategy, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 5be2451348..c257f0be28 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -23,10 +23,22 @@ */ public class CagraIndexParamsFactory { - private static final int ALGO_SWITCH_THRESHOLD = 5_000_000; - private CagraIndexParamsFactory() {} + /** + * Translation of the internal logic found in the constructor of {@code struct ivf_pq_params} in + * cuVS's {@code cpp/include/cuvs/neighbors/ivf_pq.hpp}. + * + * Ideally we should hook into the internal API but this is currently replicated to avoid complications + * in other parts of code base. + * + * TODO: cuvs-java has no standalone binding for that C++ constructor -- only fromHnswParams and + * fromDataset, which bundle algorithm selection together with parameter tuning and offer no way + * to pin the algorithm while still getting cuVS's own auto-tuned IVF-PQ parameters. If cuvs-java + * exposed the ivf_pq_params(dataset_extents, metric) constructor directly (or an equivalent + * narrower entry point that derives IVF-PQ parameters without also choosing the algorithm), this + * method could delegate to it instead of reimplementing the heuristic here. + */ private static CuVSIvfPqParams getCuVSIvfPqParams(long rows, long dimension) { int pqDim; int pqBits; @@ -104,16 +116,23 @@ private static CagraIndexParams getNNDescentParams( .build(); } + /** + * @param explicitIvfPqParams the caller's own IVF-PQ params if they set one ({@link + * AcceleratedHNSWParams#isCuVSIvfPqParamsExplicit()}), honored as-is; otherwise {@code + * null}, in which case params are auto-tuned from {@code rows}/{@code dimension}. + */ private static CagraIndexParams getIVFPQParams( int graphDegree, int intGraphDegree, int writerThreads, long rows, long dimension, - CuvsDistanceType cuvsDistanceType) { + CuvsDistanceType cuvsDistanceType, + CuVSIvfPqParams explicitIvfPqParams) { return new CagraIndexParams.Builder() .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.IVF_PQ) - .withCuVSIvfPqParams(getCuVSIvfPqParams(rows, dimension)) + .withCuVSIvfPqParams( + explicitIvfPqParams != null ? explicitIvfPqParams : getCuVSIvfPqParams(rows, dimension)) .withNumWriterThreads(writerThreads) .withIntermediateGraphDegree(intGraphDegree) .withGraphDegree(graphDegree) @@ -187,7 +206,10 @@ public static CagraIndexParams create( acceleratedHNSWParams.getWriterThreads(), rows, dimension, - acceleratedHNSWParams.getCuvsDistanceType()); + acceleratedHNSWParams.getCuvsDistanceType(), + acceleratedHNSWParams.isCuVSIvfPqParamsExplicit() + ? acceleratedHNSWParams.getCuVSIvfPqParams() + : null); } else if (algo == CagraGraphBuildAlgo.NN_DESCENT) { return getNNDescentParams( acceleratedHNSWParams.getGraphdegree(), From 3ab2330c3092481f792f4fa69377fd23824ae55a Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 14 Aug 2026 08:15:28 -0700 Subject: [PATCH 19/20] Guard against native flat buffering with quantized fields FieldWriter's native-buffering path always preallocates its host matrix as DataType.FLOAT (4 bytes/dimension), regardless of quantizationType. This is currently safe -- the only call site that ever passes numInputVectors > 0 (Lucene99AcceleratedHNSWVectorsWriter) always passes QuantizationType.NONE, and the quantized writers never enable native buffering -- but nothing enforced it: it held only because no caller combines them today. Add a check in FieldWriter's constructor so that combination fails loudly (matching the existing numInputVectors guard message convention) instead of silently mis-sizing the native buffer if a future caller ever does combine them. Document the constraint (currently unsupported, not fixable via a caller workaround) everywhere numInputVectors' requirements are already described: AcceleratedHNSWParams.withNumInputVectors's javadoc, FieldWriter's own field javadoc, and OptimizedCagraHnswBuildExample's inline comment alongside the existing index-sort/merge constraints. --- .../lucene/examples/OptimizedCagraHnswBuildExample.java | 3 ++- .../java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java | 6 +++--- .../src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java | 7 +++++++ 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java index e9a6761471..a26d2349bc 100644 --- a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java @@ -388,7 +388,8 @@ private static Codec codecFor(int numInputVectors) throws Exception { // pass or omit withNumInputVectors (pass 0) to fall back to the heap-buffered path, // which buffers all vectors in a List on the JVM heap before building and // therefore uses more peak host memory. Index-sorted segments (IndexWriterConfig - // .setIndexSort) are also unsupported. + // .setIndexSort) are also unsupported, and binary/scalar quantized fields are not yet + // supported. .withNumInputVectors(numInputVectors) .build(); return new Lucene101AcceleratedHNSWCodec(params); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 873c1767db..0cdead8ecf 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -549,9 +549,9 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) { * native flat buffer (avoiding the on-heap {@code List} and the extra host-matrix * copy). The native buffer is sized for exactly this many rows, so the value MUST equal the * number of vectors actually added; the writer fails fast otherwise. Only supported for the - * unsorted single-segment CAGRA_HNSW build (no merges). A value of - * {@value DEFAULT_NUM_INPUT_VECTORS} (the default) disables it and uses the default - * heap-buffered flat path. + * unsorted single-segment CAGRA_HNSW build (no merges). Not yet supported for the + * binary/scalar quantized writers. A value of {@value DEFAULT_NUM_INPUT_VECTORS} (the + * default) disables it and uses the default heap-buffered flat path. * * @param numInputVectors the exact number of vectors to be indexed, or 0 to disable * @return instance of {@link Builder} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 64f75766ea..b6b1d9a0cd 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -38,6 +38,8 @@ public class FieldWriter extends KnnFieldVectorsWriter { * *

      The matrix is preallocated for exactly {@code numInputVectors} rows, so the hint must equal * the number of vectors actually added (validated at build time in the caller). + * + *

      Not yet supported for quantized fields ({@code quantizationType != QuantizationType.NONE}). */ private final int numInputVectors; @@ -68,6 +70,11 @@ public FieldWriter( this.numInputVectors = numInputVectors; this.nativeBuffering = numInputVectors > 0; if (nativeBuffering) { + if (quantizationType != QuantizationType.NONE) { + throw new UnsupportedOperationException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" + + " quantized fields; unset it (0) to use the heap-buffered path"); + } // Preallocates one contiguous native region of numInputVectors * dimension * 4 bytes. this.hostMatrixBuilder = CuVSMatrix.hostBuilder(numInputVectors, dimension, CuVSMatrix.DataType.FLOAT); From 563428cfa03b1673d768caafe6220502de1a96c3 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 18 Aug 2026 14:32:33 -0700 Subject: [PATCH 20/20] Apply Spotless formatting --- .../OptimizedCagraHnswBuildExample.java | 24 ++++++++++++------ .../cuvs/lucene/AcceleratedHNSWUtils.java | 13 ++++++---- .../com/nvidia/cuvs/lucene/FieldWriter.java | 3 ++- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 3 +-- .../Lucene99AcceleratedHNSWVectorsWriter.java | 6 +++-- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 6 +++-- ...ratedHNSWScalarQuantizedVectorsWriter.java | 6 +++-- .../cuvs/lucene/NativeFlatVectorsWriter.java | 9 ++----- .../lucene/TestMergedGraphOrdinalBounds.java | 25 +++++++++++-------- ...TestNativeFlatBufferingIndexAndSearch.java | 6 ++--- 10 files changed, 58 insertions(+), 43 deletions(-) diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java index a26d2349bc..f76b899243 100644 --- a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java @@ -209,8 +209,15 @@ private static void buildSequential( for (int p = 0; p < slices.size(); p++) { int[] slice = slices.get(p); log.info( - "Building segment " + (p + 1) + "/" + slices.size() + ": docs [" + slice[0] + ", " - + (slice[0] + slice[1]) + ")"); + "Building segment " + + (p + 1) + + "/" + + slices.size() + + ": docs [" + + slice[0] + + ", " + + (slice[0] + slice[1]) + + ")"); buildSegment(dir, reader, scratch, slice[0], slice[1], p == 0, null); } } @@ -244,7 +251,8 @@ private static void buildOverlapped( for (int p = 0; p < slices.size(); p++) { segDirs.add(Paths.get(indexDirPath + "_p" + p)); } - // Start from fresh per-segment temp dirs, and always remove them afterwards (even on failure) so + // Start from fresh per-segment temp dirs, and always remove them afterwards (even on failure) + // so // a crashed build does not leave orphaned per-segment indexes behind. for (Path segDir : segDirs) { FileUtils.deleteQuietly(segDir.toFile()); @@ -304,7 +312,8 @@ private static void buildSegment( // Keep this slice in ONE segment: raise the doc-count flush threshold above the slice size and // disable RAM-based flushing so nothing flushes before commit, and forbid merges. This is what - // makes native flat buffering valid. Order matters: enable the doc-count trigger BEFORE disabling + // makes native flat buffering valid. Order matters: enable the doc-count trigger BEFORE + // disabling // the RAM trigger, since Lucene rejects a config where both are disabled at once. IndexWriterConfig config = new IndexWriterConfig() @@ -314,9 +323,7 @@ private static void buildSegment( .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) .setMergePolicy(NoMergePolicy.INSTANCE) .setOpenMode( - createNew - ? IndexWriterConfig.OpenMode.CREATE - : IndexWriterConfig.OpenMode.APPEND); + createNew ? IndexWriterConfig.OpenMode.CREATE : IndexWriterConfig.OpenMode.APPEND); try (IndexWriter writer = new IndexWriter(dir, config)) { for (int i = 0; i < size; i++) { @@ -421,7 +428,8 @@ private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) try (Directory dir = FSDirectory.open(indexDirPath); DirectoryReader reader = DirectoryReader.open(dir)) { IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); log.info("Sample search returned " + results.scoreDocs.length + " hits:"); for (int i = 0; i < results.scoreDocs.length; i++) { ScoreDoc sd = results.scoreDocs[i]; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index fc73651be7..730940ca61 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -243,8 +243,8 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph( - GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) throws IOException { + public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) + throws IOException { // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); int numLevels = graph.numLevels(); @@ -252,8 +252,10 @@ public static int[][] writeGraph( // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in - // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk bytes - // are identical to the fully-serial path (blocks in node order, offsets = per-node byte lengths). + // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk + // bytes + // are identical to the fully-serial path (blocks in node order, offsets = per-node byte + // lengths). int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); offsets[0] = new int[level0Nodes.length]; if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { @@ -363,7 +365,8 @@ private static void writeLevel0Parallel( * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. */ private static void encodeNode( - NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) throws IOException { + NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) + throws IOException { int size = neighbors == null ? 0 : neighbors.size(); int actualSize = 0; if (size > 0) { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index b6b1d9a0cd..92df78224b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -178,7 +178,8 @@ public Object copyValue(Object vectorValue) { @Override public long ramBytesUsed() { - // The native host matrix is off-heap and intentionally excluded from Lucene's heap RAM accounting. + // The native host matrix is off-heap and intentionally excluded from Lucene's heap RAM + // accounting. return SHALLOW_SIZE + (nativeBuffering ? 0 : flatFieldVectorsWriter.ramBytesUsed()); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index beb0f0cc77..c4f8a3ea8f 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -90,8 +90,7 @@ public GPUBuiltHnswGraph( * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) * @return the NeighborArray */ - private static NeighborArray[] fillNeighborArray( - CuVSMatrix adjacency, int size, int numThreads) { + private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) { NeighborArray[] neighbors = new NeighborArray[size]; if (numThreads <= 1 || size < PARALLEL_MIN_NODES) { fillNeighborRange(adjacency, neighbors, 0, size); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 8fc2b7f17a..f617a23890 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -236,7 +236,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws QuantizationType.NONE, acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -349,7 +350,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 82cd4ed5ba..cf6d19f7ca 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -188,7 +188,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -276,7 +277,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 0a53bfbd15..aebfbea7c6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -214,7 +214,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -301,7 +302,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java index 8c5adb08b4..c47954fbd1 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java @@ -92,8 +92,7 @@ final class NativeFlatVectorsWriter implements Closeable { NativeFlatVectorsWriter(SegmentWriteState state) throws IOException { String metaFileName = - IndexFileNames.segmentFileName( - state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); String vectorDataFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); @@ -102,11 +101,7 @@ final class NativeFlatVectorsWriter implements Closeable { meta = state.directory.createOutput(metaFileName, state.context); vectorData = state.directory.createOutput(vectorDataFileName, state.context); CodecUtil.writeIndexHeader( - meta, - META_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); + meta, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); CodecUtil.writeIndexHeader( vectorData, VECTOR_DATA_CODEC_NAME, diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java index 1fd8d2a3af..12b00c6e12 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java @@ -82,7 +82,8 @@ public void testMergedGraphOrdinalsStayWithinFlatVectorBounds() throws Exception try (Directory dir = newDirectory(); IndexWriter writer = new IndexWriter(dir, config)) { int deletedFromSegment1 = - addSegmentWithInterspersedDeletions(writer, 0, segmentSize, dimension, deleteEveryNth, random); + addSegmentWithInterspersedDeletions( + writer, 0, segmentSize, dimension, deleteEveryNth, random); writer.commit(); // segment 1, alone int deletedFromSegment2 = addSegmentWithInterspersedDeletions( @@ -98,7 +99,8 @@ public void testMergedGraphOrdinalsStayWithinFlatVectorBounds() throws Exception writer.commit(); try (DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected the forced merge to produce a single segment", 1, reader.leaves().size()); + assertEquals( + "expected the forced merge to produce a single segment", 1, reader.leaves().size()); LeafReader leaf = reader.leaves().get(0).reader(); FloatVectorValues flatValues = leaf.getFloatVectorValues(FIELD); @@ -133,11 +135,17 @@ private static void assertAllNeighborOrdinalsInBounds(HnswGraph graph, int liveV while (nodes.hasNext()) { int node = nodes.nextInt(); assertTrue( - "node " + node + " at level " + level + " is itself out of bounds (live vectors: " - + liveVectorCount + ")", + "node " + + node + + " at level " + + level + + " is itself out of bounds (live vectors: " + + liveVectorCount + + ")", node >= 0 && node < liveVectorCount); graph.seek(level, node); - for (int neighbor = graph.nextNeighbor(); neighbor != NO_MORE_DOCS; + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; neighbor = graph.nextNeighbor()) { assertTrue( "node " @@ -171,12 +179,7 @@ private static HnswGraph graphOf(LeafReader leaf) throws Exception { * @return the number of documents deleted from this segment */ private static int addSegmentWithInterspersedDeletions( - IndexWriter writer, - int startId, - int count, - int dimension, - int deleteEveryNth, - Random random) + IndexWriter writer, int startId, int count, int dimension, int deleteEveryNth, Random random) throws Exception { float[][] dataset = generateDataset(random, count, dimension); for (int i = 0; i < count; i++) { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java index 1829e4db62..147b7aca24 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java @@ -144,7 +144,8 @@ public void testDeletedDocsAfterNativeFlatBufferedFlush() throws Exception { try (Directory dir = FSDirectory.open(indexDirPath); DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected the deletions to land in the same single segment", 1, reader.leaves().size()); + assertEquals( + "expected the deletions to land in the same single segment", 1, reader.leaves().size()); assertTrue("expected some deleted docs", reader.numDeletedDocs() > 0); IndexSearcher searcher = new IndexSearcher(reader); @@ -184,8 +185,7 @@ public void testOddGraphDegreeWithNativeFlatBuffering() throws Exception { try (Directory dir = FSDirectory.open(indexDirPath); DirectoryReader reader = DirectoryReader.open(dir)) { IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); assertEquals(5, results.scoreDocs.length); } }