From e5de014cd7502315d813abe2684b403f222a8dfa Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 14 Aug 2025 23:00:20 -0400 Subject: [PATCH 01/21] Refactor merge to call CAGRA's merge(), implement CAGRA prefiltering (#14) --- .../com/nvidia/cuvs/lucene/CuVSIndex.java | 4 +- .../cuvs/lucene/CuVSKnnFloatVectorQuery.java | 8 +- .../nvidia/cuvs/lucene/CuVSVectorsReader.java | 23 +- .../nvidia/cuvs/lucene/CuVSVectorsWriter.java | 257 +++- .../java/com/nvidia/cuvs/lucene/Utils.java | 51 + .../cuvs/lucene/TestCuVSDeletedDocuments.java | 342 +++++ .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 197 +++ ...va => TestCuVSRandomizedVectorSearch.java} | 84 +- .../cuvs/lucene/TestCuVSVectorsFormat.java | 2 + .../lucene/TestIndexOutputOutputStream.java | 11 +- .../com/nvidia/cuvs/lucene/TestMerge.java | 1169 +++++++++++++++++ .../com/nvidia/cuvs/lucene/TestUtils.java | 50 + 12 files changed, 2087 insertions(+), 111 deletions(-) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/Utils.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java rename src/test/java/com/nvidia/cuvs/lucene/{TestCuVS.java => TestCuVSRandomizedVectorSearch.java} (74%) create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestMerge.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestUtils.java diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java index 78c5dc14..e9a20147 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java @@ -15,8 +15,6 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.CuVSVectorsReader.handleThrowable; - import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.HnswIndex; @@ -112,7 +110,7 @@ private void destroyIndices() throws IOException { hnswIndex.destroyIndex(); } } catch (Throwable t) { - handleThrowable(t); + Utils.handleThrowable(t); } } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java index fa1e71d5..8caf30ae 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java @@ -19,6 +19,7 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.knn.KnnCollectorManager; import org.apache.lucene.util.Bits; @@ -29,8 +30,9 @@ public class CuVSKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; - public CuVSKnnFloatVectorQuery(String field, float[] target, int k, int iTopK, int searchWidth) { - super(field, target, k); + public CuVSKnnFloatVectorQuery( + String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { + super(field, target, k, filter); this.iTopK = iTopK; this.searchWidth = searchWidth; } @@ -46,7 +48,7 @@ protected TopDocs approximateSearch( PerLeafCuVSKnnCollector results = new PerLeafCuVSKnnCollector(k, iTopK, searchWidth); LeafReader reader = context.reader(); - reader.searchNearestVectors(field, this.getTargetCopy(), results, null); + reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs); return results.topDocs(); } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java index c770015d..4118a0aa 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java @@ -271,7 +271,7 @@ private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { } } } catch (Throwable t) { - handleThrowable(t); + Utils.handleThrowable(t); } return new CuVSIndex(cagraIndex, bruteForceIndex, hnswIndex); } @@ -367,7 +367,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits try { searchResult = cagraIndex.search(query).getResults(); } catch (Throwable t) { - handleThrowable(t); + Utils.handleThrowable(t); } // List expected to have only one entry because of single query "target". assert searchResult.size() == 1; @@ -385,7 +385,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits try { searchResult = bruteforceIndex.search(query).getResults(); } catch (Throwable t) { - handleThrowable(t); + Utils.handleThrowable(t); } assert searchResult.size() == 1; result = searchResult.getFirst(); @@ -472,12 +472,15 @@ static void checkVersion(int versionMeta, int versionVectorData, IndexInput in) } } - static void handleThrowable(Throwable t) throws IOException { - switch (t) { - case IOException ioe -> throw ioe; - case Error error -> throw error; - case RuntimeException re -> throw re; - case null, default -> throw new RuntimeException("UNEXPECTED: exception type", t); - } + public FieldInfos getFieldInfos() { + return fieldInfos; + } + + public IntObjectHashMap getCuvsIndexes() { + return cuvsIndices; + } + + public IntObjectHashMap getFieldEntries() { + return fields; } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java index 3e22b4da..0bc3eca8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java @@ -20,7 +20,6 @@ import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_EXT; import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_CURRENT; -import static com.nvidia.cuvs.lucene.CuVSVectorsReader.handleThrowable; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -32,6 +31,7 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSMatrix.DataType; import com.nvidia.cuvs.CuVSResources; import java.io.IOException; import java.io.OutputStream; @@ -41,14 +41,18 @@ import java.util.ArrayList; import java.util.List; import java.util.Objects; +import java.util.function.Supplier; import java.util.logging.Logger; +import java.util.stream.IntStream; import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.KnnVectorValues; @@ -57,6 +61,7 @@ import org.apache.lucene.index.Sorter; import org.apache.lucene.index.Sorter.DocMap; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntObjectHashMap; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; @@ -277,37 +282,36 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { } private void writeField(CuVSFieldWriter fieldData) throws IOException { - // TODO: Argh! https://github.com/rapidsai/cuvs/issues/698 + // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors + // from the flat writer to the CuVSMatrix? List vectors = fieldData.getVectors(); - CuVSMatrix.Builder builder = - CuVSMatrix.builder( - vectors.size(), fieldData.fieldInfo().getVectorDimension(), CuVSMatrix.DataType.FLOAT); - for (float[] vec : vectors) builder.addVector(vec); - writeFieldInternal(fieldData.fieldInfo(), builder.build()); + writeFieldInternal( + fieldData.fieldInfo(), + () -> Utils.createFloatMatrix(vectors, fieldData.fieldInfo().getVectorDimension()), + vectors.size()); } private void writeSortingField(CuVSFieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord - mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); - - float[][] oldVectors = fieldData.getVectors().toArray(float[][]::new); - CuVSMatrix.Builder builder = - CuVSMatrix.builder( - fieldData.getVectors().size(), - fieldData.fieldInfo().getVectorDimension(), - CuVSMatrix.DataType.FLOAT); - for (int i = 0; i < oldVectors.length; i++) { - float[] vec = oldVectors[new2OldOrd[i]]; - builder.addVector(vec); + // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors + // from the flat writer to the CuVSMatrix? + List sortedVectors = new ArrayList(); + for (int i = 0; i < fieldData.getVectors().size(); i++) { + sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); } - writeFieldInternal(fieldData.fieldInfo(), builder.build()); + writeFieldInternal( + fieldData.fieldInfo(), + () -> Utils.createFloatMatrix(sortedVectors, fieldData.fieldInfo().getVectorDimension()), + sortedVectors.size()); } - private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - if (dataset.size() == 0) { + private void writeFieldInternal( + FieldInfo fieldInfo, Supplier datasetSupplier, int datasetSize) + throws IOException { + if (datasetSize == 0) { writeEmpty(fieldInfo); return; } @@ -317,7 +321,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws // workaround for the minimum number of vectors for Cagra IndexType indexType = - this.indexType.cagra() && dataset.size() < MIN_CAGRA_INDEX_SIZE + this.indexType.cagra() && datasetSize < MIN_CAGRA_INDEX_SIZE ? IndexType.BRUTE_FORCE : this.indexType; @@ -326,7 +330,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws if (indexType.cagra()) { try { var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - writeCagraIndex(cagraIndexOutputStream, dataset); + writeCagraIndex(cagraIndexOutputStream, datasetSupplier.get()); } catch (Throwable t) { handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); // workaround for cuVS issue @@ -338,16 +342,16 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws bruteForceIndexOffset = cuvsIndex.getFilePointer(); if (indexType.bruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - writeBruteForceIndex(bruteForceIndexOutputStream, dataset); + writeBruteForceIndex(bruteForceIndexOutputStream, datasetSupplier.get()); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; } hnswIndexOffset = cuvsIndex.getFilePointer(); if (indexType.hnsw()) { var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - if (dataset.size() > MIN_CAGRA_INDEX_SIZE) { + if (datasetSize > MIN_CAGRA_INDEX_SIZE) { try { - writeHNSWIndex(hnswIndexOutputStream, dataset); + writeHNSWIndex(hnswIndexOutputStream, datasetSupplier.get()); } catch (Throwable t) { handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); } @@ -357,7 +361,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws writeMeta( fieldInfo, - (int) dataset.size(), + (int) datasetSize, cagraIndexOffset, cagraIndexLength, bruteForceIndexOffset, @@ -365,7 +369,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws hnswIndexOffset, hnswIndexLength); } catch (Throwable t) { - handleThrowable(t); + Utils.handleThrowable(t); } } @@ -418,46 +422,185 @@ static void handleThrowableWithIgnore(Throwable t, String msg) throws IOExceptio if (t.getMessage().contains(msg)) { return; } - handleThrowable(t); + Utils.handleThrowable(t); + } + + private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + + List cagraIndexes = new ArrayList<>(); + // We need this count so that the merged segment's meta information has the vector count. + int totalVectorCount = 0; + + for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { + KnnVectorsReader knnReader = mergeState.knnVectorsReaders[i]; + // Access the CAGRA index for this field from the reader + + if (knnReader != null) { + if (knnReader instanceof CuVSVectorsReader cvr) { + if (cvr != null) { + totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); + CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); + if (cagraIndex != null) { + cagraIndexes.add(cagraIndex); + } + } + } else { + // This should never happen + throw new RuntimeException( + "Reader is not of CuVSVectorsReader type. Instead it is: " + knnReader.getClass()); + } + } + } + assert cagraIndexes.size() > 1; + + CagraIndex mergedIndex = + CagraIndex.merge(cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()])); + writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + info("Successfully merged " + cagraIndexes.size() + " CAGRA indexes using native merge API"); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } } /** - * Copies the vector values into dst. Returns the actual number of vectors - * copied. + * Fallback method that rebuilds indexes from merged vectors. + * Used when native CAGRA merge() is not possible. Also used + * when non-CAGRA index types are used (for e.g. Brute Force index). */ - private static int getVectorData(FloatVectorValues floatVectorValues, CuVSMatrix.Builder builder) - throws IOException { - DocsWithFieldSet docsWithField = new DocsWithFieldSet(); - int count = 0; - KnnVectorValues.DocIndexIterator iter = floatVectorValues.iterator(); + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + if (fieldInfo.getVectorEncoding() != FLOAT32) { + throw new AssertionError("Only Float32 supported"); + } + try { + // We need to compute the size of the number of merged documents up-front so that we can + // compute the CuVSMatrix capacity. TODO: Find a way to do this without merging twice. + final int numMergedDocs = getMergedDocsCount(fieldInfo, mergeState); + + if (numMergedDocs != 0) { + writeFieldInternal( + fieldInfo, + () -> { + try { + return createMatrixFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues( + fieldInfo, mergeState), + numMergedDocs); + } catch (IOException e) { + throw new RuntimeException(e); + } + }, + numMergedDocs); + } else { + writeEmpty(fieldInfo); + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + private int getMergedDocsCount(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + KnnVectorValues.DocIndexIterator iter = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) + .iterator(); + int numMergedDocs = 0; + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + numMergedDocs++; + } + return numMergedDocs; + } + + /** + * Creates CuVSMatrix from merged vectors + * */ + private CuVSMatrix createMatrixFromMergedVectors( + FloatVectorValues mergedVectorValues, int numMergedDocs) throws IOException { + CuVSMatrix.Builder builder = + CuVSMatrix.builder(numMergedDocs, mergedVectorValues.dimension(), DataType.FLOAT); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - assert iter.index() == count; - builder.addVector(floatVectorValues.vectorValue(iter.index())); // is this correct? - docsWithField.add(docV); - count++; + int ordinal = iter.index(); + float[] vector = mergedVectorValues.vectorValue(ordinal); + builder.addVector(vector.clone()); + } + return builder.build(); + } + + /** + * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. + */ + private CagraIndex getCagraIndexFromReader(CuVSVectorsReader reader, String fieldName) { + try { + IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); + FieldInfos fieldInfos = reader.getFieldInfos(); + + FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); + + if (fieldInfo != null) { + CuVSIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); + if (cuvsIndex != null) { + return cuvsIndex.getCagraIndex(); + } + } + } catch (Exception e) { + e.printStackTrace(); + info("Failed to extract CAGRA index for field " + fieldName + ": " + e.getMessage()); + } + return null; + } + + /** + * Writes a pre-built merged CAGRA index to the output. + */ + private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, int vectorCount) + throws IOException { + try { + long cagraIndexOffset = cuvsIndex.getFilePointer(); + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + + // Serialize the merged index + Path tmpFile = Files.createTempFile(resources.tempDirectory(), "mergedindex", "cag"); + mergedIndex.serialize(cagraIndexOutputStream, tmpFile); + long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + + // Write metadata (assuming no brute force or HNSW indexes for merged result) + writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); + + // Clean up the merged index + mergedIndex.destroyIndex(); + } catch (Throwable t) { + Utils.handleThrowable(t); } - return docsWithField.cardinality(); } @Override public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { flatVectorsWriter.mergeOneField(fieldInfo, mergeState); - try { - final FloatVectorValues mergedVectorValues = - switch (fieldInfo.getVectorEncoding()) { - case BYTE -> throw new AssertionError("bytes not supported"); - case FLOAT32 -> - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - }; - - // Also will be replaced with the cuVS merge api - CuVSMatrix.Builder builder = - CuVSMatrix.builder( - mergedVectorValues.size(), mergedVectorValues.dimension(), CuVSMatrix.DataType.FLOAT); - getVectorData(mergedVectorValues, builder); - writeFieldInternal(fieldInfo, builder.build()); - } catch (Throwable t) { - handleThrowable(t); + + if (indexType.cagra() && !indexType.bruteForce()) { + // Since CAGRA merge does not support merging of indexes with purging of deletes, + // we fallback to vector-based re-indexing. Issue: + // https://github.com/rapidsai/cuvs/issues/1253 + boolean hasDeletions = + IntStream.range(0, mergeState.liveDocs.length) + .anyMatch( + i -> + mergeState.liveDocs[i] == null + || IntStream.range(0, mergeState.maxDocs[i]) + .anyMatch(j -> !mergeState.liveDocs[i].get(j))); + + if (mergeState.knnVectorsReaders.length > 1 && !hasDeletions) { + mergeCagraIndexes(fieldInfo, mergeState); + } else { + // CAGRA's merge API does not handle the trivial case of merging 1 index. + vectorBasedMerge(fieldInfo, mergeState); + } + + } else { + // If there is a Brute Force index then re-index using the vectors even if there is a CAGRA + // index. + vectorBasedMerge(fieldInfo, mergeState); } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java new file mode 100644 index 00000000..d3d3b48f --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -0,0 +1,51 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.List; + +public class Utils { + + static void handleThrowable(Throwable t) throws IOException { + switch (t) { + case IOException ioe -> throw ioe; + case Error error -> throw error; + case RuntimeException re -> throw re; + case null, default -> throw new RuntimeException("UNEXPECTED: exception type", t); + } + } + + /** + * A method to build a {@link CuVSMatrix} from a list of float vectors. + * + * Note: This could be a memory-intensive operation and should therefore be avoided. + * Consider using this {@link CuVSMatrix.Builder} instead for copying the vectors without loading them in heap. + * + * @param data The float vectors + * @param dimensions The number float elements in each vector + * @return an instance of {@link CuVSMatrix} + */ + static CuVSMatrix createFloatMatrix(List data, int dimensions) { + CuVSMatrix.Builder builder = + CuVSMatrix.builder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + for (float[] vector : data) { + builder.addVector(vector); + } + return builder.build(); + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java new file mode 100644 index 00000000..ef17136e --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -0,0 +1,342 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; +import java.util.logging.Logger; +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.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +public class TestCuVSDeletedDocuments extends LuceneTestCase { + + protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + private static Random random; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + random = random(); + } + + @Test + public void testVectorSearchWithDeletedDocuments() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200, 1000); // 200-1200 documents + int dimensions = random.nextInt(64, 256); // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set deletedDocs = new HashSet<>(); + + // Create index with all documents having vectors + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + + // Delete documents randomly based on probability + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Search and verify deleted documents are not returned + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + // Use a random vector for query + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we got results + assertTrue("Should have search results", hits.length > 0); + + // Verify no deleted documents in results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse( + "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); + log.info("Found non-deleted document: " + id + ", Score: " + hit.score); + } + + // Verify deleted documents are truly deleted + for (int deletedId : deletedDocs) { + TopDocs result = + searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + assertEquals( + "Deleted document " + deletedId + " should not be found", + 0, + result.totalHits.value()); + } + } + } + } + + @Test + public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + Set docsWithoutVectors = new HashSet<>(); + Set deletedDocs = new HashSet<>(); + + // Create index with mixed documents + try (RandomIndexWriter writer = createWriter(directory)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + // Randomly assign categories + String category = random.nextBoolean() ? "A" : "B"; + doc.add(new StringField("category", category, Field.Store.YES)); + + // Randomly decide whether to add vectors + if (random.nextFloat() < vectorProbability) { + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } else { + docsWithoutVectors.add(i); + } + writer.addDocument(doc); + } + + // Delete documents randomly + for (int i = 0; i < datasetSize; i++) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + deletedDocs.add(i); + } + } + writer.commit(); + } + + // Test vector search behavior + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify results + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertFalse("Deleted document should not appear", deletedDocs.contains(id)); + assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); + log.info("Found document with vector: " + id + ", Score: " + hit.score); + } + + // Test filtered search with deletions + Query filter = new TermQuery(new Term("category", "A")); + Query filteredQuery = + new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + for (ScoreDoc hit : filteredHits) { + Document doc = reader.storedFields().document(hit.doc); + String category = doc.get("category"); + assertEquals("Should only match category A", "A", category); + int id = Integer.parseInt(doc.get("id")); + assertFalse( + "Deleted document should not appear in filtered results", deletedDocs.contains(id)); + } + } + } + } + + @Test + public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test + int dimensions = random.nextInt(128) + 32; // 32-160 dimensions + int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + + // Create and delete all documents + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + + // Delete all documents + for (int i = 0; i < datasetSize; i++) { + writer.deleteDocuments(new Term("id", String.valueOf(i))); + } + writer.commit(); + writer.forceMerge(1); // Force merge to apply deletions + } + + // Verify search returns no results + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + TopDocs results = searcher.search(query, topK); + + assertEquals( + "Should return no results when all documents are deleted", + 0, + results.totalHits.value()); + } + } + } + + @Test + public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { + + try (Directory directory = newDirectory()) { + int datasetSize = random.nextInt(200) + 50; // 50-250 documents + int dimensions = random.nextInt(256) + 64; // 64-320 dimensions + int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results + float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + + float[][] dataset = generateDataset(random, datasetSize, dimensions); + List activeDocIds = new ArrayList<>(); + + // Initial indexing + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + for (int i = 0; i < initialDocs; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + + // Delete some documents randomly + List candidatesForDeletion = new ArrayList<>(activeDocIds); + for (int docId : candidatesForDeletion) { + if (random.nextFloat() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(docId))); + activeDocIds.remove(Integer.valueOf(docId)); + } + } + + // Add new documents with higher IDs + for (int i = initialDocs; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + activeDocIds.add(i); + } + writer.commit(); + } + + // Verify search behavior after deletions and additions + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = newSearcher(reader); + float[] queryVector = generateRandomVector(dimensions, random); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + Set resultIds = new HashSet<>(); + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + resultIds.add(id); + assertTrue("Result should be from active documents", activeDocIds.contains(id)); + } + + log.info( + "Search returned " + + hits.length + + " results from " + + activeDocIds.size() + + " active documents"); + } + } + } + + private RandomIndexWriter createWriter(Directory directory) throws IOException { + return new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + } + + private IndexWriterConfig createWriterConfig() { + return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy()); + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java new file mode 100644 index 00000000..6d4ab286 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Logger; +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.IndexReader; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +public class TestCuVSGaps extends LuceneTestCase { + + protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static IndexSearcher searcher; + static IndexReader reader; + static Directory directory; + static Random random; + + static int DATASET_SIZE_LIMIT = 1000; + static int DIMENSIONS_LIMIT = 2048; + static int NUM_QUERIES_LIMIT = 10; + static int TOP_K_LIMIT = 64; + + static int datasetSize; + static int dimension; + static float[][] dataset; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + directory = newDirectory(); + random = random(); + + RandomIndexWriter writer = + new RandomIndexWriter( + random(), + directory, + newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) + .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + + log.info("Merge Policy: " + writer.w.getConfig().getMergePolicy()); + + datasetSize = random.nextInt(100, DATASET_SIZE_LIMIT); + dimension = random.nextInt(8, DIMENSIONS_LIMIT); + dataset = generateDataset(random, datasetSize, dimension); + + // Create documents where only even-numbered documents have vectors + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); + + // Only add vectors to even-numbered documents + if (i % 2 == 0) { + doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.info("Test finished"); + } + + @Test + public void testVectorSearchWithAlternatingDocuments() throws IOException { + assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + + // Use the first vector (from document 0) as query + float[] queryVector = dataset[0]; + int topK = random.nextInt(5, TOP_K_LIMIT); + + Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we get exactly TOP_K results + assertEquals("Should return exactly " + topK + " results", topK, hits.length); + + // Verify all returned documents have vectors (even-numbered IDs) + for (ScoreDoc hit : hits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + int id = Integer.parseInt(docId); + assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); + log.info("Document ID: " + id + ", Score: " + hit.score); + } + + // Verify the results match expected top-k based on Euclidean distance + List expectedIds = calculateExpectedTopK(queryVector, topK, dataset); + for (int i = 0; i < hits.length; i++) { + String docId = reader.storedFields().document(hits[i].doc).get("id"); + int id = Integer.parseInt(docId); + assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); + } + + log.info("Alternating document test passed with " + hits.length + " results"); + } + + @Test + public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { + assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + + // Use the first vector (from document 0) as query + float[] queryVector = dataset[0]; + int topK = random.nextInt(5, TOP_K_LIMIT); + + // Create a filter that only matches documents with ID less than 10 + // This should further restrict our results to even numbers 0, 2, 4, 6, 8 + Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 + + Query filteredQuery = new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Should only get document 8 (the only one that matches the filter and has a vector) + assertEquals("Should return exactly 1 result", 1, filteredHits.length); + + String docId = reader.storedFields().document(filteredHits[0].doc).get("id"); + assertEquals("Should only return document 8", "8", docId); + + log.info("Filtered alternating document test passed with " + filteredHits.length + " results"); + } + + public static List calculateExpectedTopK(float[] query, int topK, float[][] dataset) { + Map distances = new TreeMap<>(); + + // Calculate distances only for documents that have vectors (even-numbered) + for (int i = 0; i < dataset.length; i += 2) { + double distance = 0; + for (int j = 0; j < dataset[0].length; j++) { + distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); + } + distances.put(i, distance); + } + + // Sort by distance and return top-k + return distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .limit(topK) + .toList(); + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVS.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java similarity index 74% rename from src/test/java/com/nvidia/cuvs/lucene/TestCuVS.java rename to src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 124c2f93..a0973599 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVS.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -15,6 +15,9 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; + import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; @@ -29,11 +32,13 @@ import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.document.StringField; import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.Term; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.analysis.MockAnalyzer; import org.apache.lucene.tests.analysis.MockTokenizer; @@ -47,9 +52,9 @@ import org.junit.Test; @SuppressSysoutChecks(bugUrl = "prints info from within cuvs") -public class TestCuVS extends LuceneTestCase { +public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVS.class.getName()); + protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); static IndexSearcher searcher; @@ -60,8 +65,7 @@ public class TestCuVS extends LuceneTestCase { static int DIMENSIONS_LIMIT = 2048; static int NUM_QUERIES_LIMIT = 10; static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 - - public static float[][] dataset; + static float[][] dataset; @BeforeClass public static void beforeClass() throws Exception { @@ -88,7 +92,8 @@ public static void beforeClass() throws Exception { doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); boolean skipVector = - random.nextInt(10) < 0; // disable testing with holes for now, there's some bug. + random.nextInt(10) + < 4; // some documents won't have vectors to test deleted/missing vectors if (!skipVector || datasetSize < 100) { // about 10th of the documents shouldn't have a single vector doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); @@ -146,28 +151,6 @@ public void testVectorSearch() throws IOException { } } - private static float[][] generateQueries(Random random, int dimensions, int numQueries) { - // Generate random query vectors - float[][] queries = new float[numQueries][dimensions]; - for (int i = 0; i < numQueries; i++) { - for (int j = 0; j < dimensions; j++) { - queries[i][j] = random.nextFloat() * 100; - } - } - return queries; - } - - private static float[][] generateDataset(Random random, int datasetSize, int dimensions) { - // Generate a random dataset - float[][] dataset = new float[datasetSize][dimensions]; - for (int i = 0; i < datasetSize; i++) { - for (int j = 0; j < dimensions; j++) { - dataset[i][j] = random.nextFloat() * 100; - } - } - return dataset; - } - private static List> generateExpectedResults( int topK, float[][] dataset, float[][] queries) { List> neighborsResult = new ArrayList<>(); @@ -192,15 +175,50 @@ private static List> generateExpectedResults( .sorted(Map.Entry.comparingByValue()) .map(Map.Entry::getKey) .toList(); - neighborsResult.add( - neighbors.subList( - 0, - Math.min( - topK * 3, - dataset.length))); // generate double the topK results in the expected array + neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); } log.info("Expected results generated successfully."); return neighborsResult; } + + @Test + public void testVectorSearchWithFilter() throws IOException { + assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + + Random random = random(); + int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); + + if (dataset.length < topK) topK = dataset.length; + + // Find a document that has a vector by doing a search first + Query unfiltered = new KnnFloatVectorQuery("vector", dataset[0], 1); + ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; + + // Skip test if no vectors found at all + assumeTrue( + "Need at least one document with vector for filtering test", unfilteredHits.length > 0); + + String targetDocId = reader.storedFields().document(unfilteredHits[0].doc).get("id"); + float[] queryVector = dataset[0]; + + // Create a filter that matches only the document we know has a vector + Query filter = new TermQuery(new Term("id", targetDocId)); + + // Test the new constructor with filter + Query filteredQuery = new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Ensure we got some results + assertTrue("Should have at least one result", filteredHits.length > 0); + + // Verify that all results match the filter + for (ScoreDoc hit : filteredHits) { + String docId = reader.storedFields().document(hit.doc).get("id"); + assertEquals("All results should match the filter", targetDocId, docId); + } + + log.info("Prefiltering test passed with " + filteredHits.length + " results"); + } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index d19180eb..84aa1a33 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -31,10 +31,12 @@ import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.tests.util.TestUtil; import org.junit.BeforeClass; import org.junit.Ignore; +@SuppressSysoutChecks(bugUrl = "") public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { @BeforeClass diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java b/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java index d5f46d73..5f6ff3aa 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java @@ -21,7 +21,9 @@ import java.util.Random; import org.apache.lucene.store.IOContext; import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +@SuppressSysoutChecks(bugUrl = "") public class TestIndexOutputOutputStream extends LuceneTestCase { public void testBasic() throws IOException { @@ -33,8 +35,8 @@ public void testBasic() throws IOException { out.close(); } - try (var indexIn = dir.openInput("test", IOContext.DEFAULT)) { - var in = new IndexInputInputStream(indexIn); + try (var indexIn = dir.openInput("test", IOContext.DEFAULT); + var in = new IndexInputInputStream(indexIn)) { byte[] ba = new byte[6]; assertEquals(6, in.read(ba)); assertArrayEquals(new byte[] {0x56, 0x10, 0x11, 0x12, 0x13, 0x14}, ba); @@ -76,9 +78,8 @@ public void testWithRandom() throws IOException { out.close(); } - try (var indexIn = dir.openInput("test", IOContext.DEFAULT)) { - // TODO: close this stream properly in a subsequent PR. - var in = new IndexInputInputStream(indexIn); + try (var indexIn = dir.openInput("test", IOContext.DEFAULT); + var in = new IndexInputInputStream(indexIn); ) { int i = 0; while (i < data.length) { if (random.nextBoolean()) { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java new file mode 100644 index 00000000..d5b15616 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -0,0 +1,1169 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; + +import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.logging.Logger; +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.SortedDocValuesField; +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.LeafReader; +import org.apache.lucene.index.SortedDocValues; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.Sort; +import org.apache.lucene.search.SortField; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +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.util.BytesRef; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * Comprehensive tests for merge functionality with CuVS indexes. + * Tests merge operations across different index types including brute force, + * CAGRA, and combined index configurations to ensure proper vector handling + * and search functionality after segment merging. + */ +@SuppressSysoutChecks(bugUrl = "CuVS native library produces verbose logging output") +public class TestMerge extends LuceneTestCase { + + private static final Logger log = Logger.getLogger(TestMerge.class.getName()); + + private static final int MIN_VECTOR_DIMENSION = 64; + private static final int MAX_VECTOR_DIMENSION = 256; + private static final int TOP_K_LIMIT = 64; + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS is not supported", CuVSVectorsFormat.supported()); + } + + private Directory directory; + private int vectorDimension; + + @Before + public void setUp() throws Exception { + super.setUp(); + directory = newDirectory(); + + // Randomize vector dimension for each test + vectorDimension = + MIN_VECTOR_DIMENSION + random().nextInt(MAX_VECTOR_DIMENSION - MIN_VECTOR_DIMENSION + 1); + // Ensure dimension is multiple of 4 for better performance + vectorDimension = (vectorDimension / 4) * 4; + + log.info("Using randomized vector dimension: " + vectorDimension); + } + + @After + public void tearDown() throws Exception { + if (directory != null) { + directory.close(); + } + super.tearDown(); + } + + /** + * Test merging many documents across multiple segments + **/ + @Test + public void testMergeManyDocumentsMultipleSegments() throws IOException { + log.info("Starting testMergeManyDocumentsMultipleSegments"); + + // Randomize configuration parameters + int maxBufferedDocs = 5 + random().nextInt(16); // 5-20 docs per buffer + int totalBatches = 8 + random().nextInt(8); // 8-15 batches + int docsPerBatch = 15 + random().nextInt(11); // 15-25 docs per batch + int totalDocuments = totalBatches * docsPerBatch; + + // Randomize vector presence probability (60-85%) + double vectorProbability = 0.6 + (random().nextDouble() * 0.25); + + log.info( + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", totalBatches=" + + totalBatches + + ", docsPerBatch=" + + docsPerBatch + + ", totalDocuments=" + + totalDocuments + + ", vectorProbability=" + + vectorProbability); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + List expectedVectors = new ArrayList<>(); + List expectedDocIds = new ArrayList<>(); + int documentsWithVectors = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Add documents in multiple batches to create many segments + for (int batch = 0; batch < totalBatches; batch++) { + for (int i = 0; i < docsPerBatch; i++) { + int docId = batch * docsPerBatch + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new NumericDocValuesField("batch", batch)); + + // Randomly decide if document has vector + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + expectedVectors.add(vector); + expectedDocIds.add(docId); + documentsWithVectors++; + } + + writer.addDocument(doc); + } + writer.commit(); // Create a new segment + } + + int documentsWithoutVectors = totalDocuments - documentsWithVectors; + log.info("Created " + totalDocuments + " documents in " + totalBatches + " segments"); + log.info("Documents with vectors: " + documentsWithVectors); + log.info("Documents without vectors: " + documentsWithoutVectors); + + // Force merge to trigger merge logic + writer.forceMerge(1); + log.info("Forced merge to single segment completed"); + } + + // Verify the merged index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Verify vector search works correctly after merge + if (documentsWithVectors > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Randomize search parameters + int searchK = + Math.min(5 + random().nextInt(10), Math.min(documentsWithVectors, TOP_K_LIMIT)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs results = searcher.search(query, searchK); + + assertTrue("Should find some results after merge", results.scoreDocs.length > 0); + assertTrue( + "Should find reasonable number of results", + results.scoreDocs.length <= documentsWithVectors); + + log.info( + "Vector search returned " + + results.scoreDocs.length + + " results out of " + + documentsWithVectors + + " documents with vectors"); + + // Verify all returned documents have valid IDs + for (ScoreDoc scoreDoc : results.scoreDocs) { + int docId = Integer.parseInt(searcher.storedFields().document(scoreDoc.doc).get("id")); + assertTrue("Document ID should be valid", docId >= 0 && docId < totalDocuments); + } + } else { + log.info("No documents with vectors - skipping vector search verification"); + } + + log.info("Merge verification completed successfully"); + } + } + + /** + * Test merging with index sorting enabled using text-based sorting and SortingMergePolicy + **/ + @Test + public void testMergeWithIndexSorting() throws IOException { + log.info("Starting testMergeWithIndexSorting with text-based sorting"); + + // Randomize sort field type + SortField.Type sortType = random().nextBoolean() ? SortField.Type.STRING : SortField.Type.LONG; + String sortFieldName = sortType == SortField.Type.STRING ? "text_sort_key" : "numeric_sort_key"; + + // Configure index sorting by a randomized field + Sort indexSort = new Sort(new SortField(sortFieldName, sortType)); + + // Randomize merge policy parameters + TieredMergePolicy mergePolicy = new TieredMergePolicy(); + mergePolicy.setMaxMergedSegmentMB(128 + random().nextInt(257)); // 128-384 MB + mergePolicy.setSegmentsPerTier(3 + random().nextInt(4)); // 3-6 segments per tier + + // Randomize writer configuration parameters + int maxBufferedDocs = 10 + random().nextInt(16); // 10-25 docs per buffer + int totalDocuments = 80 + random().nextInt(81); // 80-160 documents + int segmentSize = 15 + random().nextInt(11); // 15-25 docs per segment + double vectorProbability = 0.65 + (random().nextDouble() * 0.25); // 65-90% have vectors + + log.info( + "Randomized sorting parameters: sortType=" + sortType + ", sortFieldName=" + sortFieldName); + log.info( + "Randomized config: maxBufferedDocs=" + + maxBufferedDocs + + ", totalDocuments=" + + totalDocuments + + ", segmentSize=" + + segmentSize + + ", vectorProbability=" + + vectorProbability); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setIndexSort(indexSort) // This automatically enables sorting during merges + .setMergePolicy(mergePolicy) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + // List documents = new ArrayList<>(); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create documents with randomized sort keys + for (int i = 0; i < totalDocuments; i++) { + float[] vector = null; + + // Randomly decide if document has vector + if (random().nextDouble() < vectorProbability) { + vector = generateRandomVector(vectorDimension, random()); + } + + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(new StringField("original_order", String.valueOf(i), Field.Store.YES)); + + // Add sort field based on randomized type + if (sortType == SortField.Type.STRING) { + // Randomize text sort key length (4-12 characters) + int keyLength = 4 + random().nextInt(9); + String textSortKey = generateRandomText(random(), keyLength); + doc.add(new SortedDocValuesField(sortFieldName, new BytesRef(textSortKey))); + doc.add(new StringField(sortFieldName + "_stored", textSortKey, Field.Store.YES)); + } else { + // Use numeric sort key with wider range + long numericSortKey = random().nextLong() % 100000; // Can be negative for more variety + doc.add(new NumericDocValuesField(sortFieldName, numericSortKey)); + doc.add( + new StringField( + sortFieldName + "_stored", String.valueOf(numericSortKey), Field.Store.YES)); + } + + if (vector != null) { + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + } + + writer.addDocument(doc); + + // Commit based on randomized segment size + if ((i + 1) % segmentSize == 0) { + writer.commit(); + log.info( + "Committed segment " + + ((i + 1) / segmentSize) + + " with " + + (i + 1) + + " total documents"); + } + } + + log.info("Created " + totalDocuments + " documents with text-based index sorting"); + + // Force merge with sorting - this will use the sorting merge policy + writer.forceMerge(1); + log.info("Forced merge with text-based sorting completed"); + } + + // Verify the merged and sorted index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Verify documents are sorted correctly by the randomized sort field + log.info( + "Verifying document sorting order using sortType: " + + sortType + + ", field: " + + sortFieldName); + + if (sortType == SortField.Type.STRING) { + // Verify string-based sorting + String previousSortKey = ""; + SortedDocValues sortedValues = leafReader.getSortedDocValues(sortFieldName); + + for (int docId = 0; docId < leafReader.maxDoc(); docId++) { + String currentSortKey = ""; + if (sortedValues != null && sortedValues.advanceExact(docId)) { + currentSortKey = sortedValues.lookupOrd(sortedValues.ordValue()).utf8ToString(); + } + + assertTrue( + "Documents should be sorted by " + + sortFieldName + + ": '" + + previousSortKey + + "' should be <= '" + + currentSortKey + + "'", + previousSortKey.compareTo(currentSortKey) <= 0); + previousSortKey = currentSortKey; + + // Log first 10 documents to verify sorting + if (docId < 10) { + IndexSearcher searcher = new IndexSearcher(reader); + String originalOrder = searcher.storedFields().document(docId).get("original_order"); + log.info( + "DocId: " + + docId + + ", OriginalOrder: " + + originalOrder + + ", SortKey: '" + + currentSortKey + + "'"); + } + } + } else { + // Verify numeric-based sorting + long previousSortKey = Long.MIN_VALUE; + var numericValues = leafReader.getNumericDocValues(sortFieldName); + + for (int docId = 0; docId < leafReader.maxDoc(); docId++) { + long currentSortKey = Long.MIN_VALUE; + if (numericValues != null && numericValues.advanceExact(docId)) { + currentSortKey = numericValues.longValue(); + } + + assertTrue( + "Documents should be sorted by " + + sortFieldName + + ": " + + previousSortKey + + " should be <= " + + currentSortKey, + previousSortKey <= currentSortKey); + previousSortKey = currentSortKey; + + // Log first 10 documents to verify sorting + if (docId < 10) { + IndexSearcher searcher = new IndexSearcher(reader); + String originalOrder = searcher.storedFields().document(docId).get("original_order"); + log.info( + "DocId: " + + docId + + ", OriginalOrder: " + + originalOrder + + ", SortKey: " + + currentSortKey); + } + } + } + + // Count total vectors by checking if vector field exists and has values + var vectorValues = leafReader.getFloatVectorValues("vector"); + int documentsWithVectors = vectorValues != null ? vectorValues.size() : 0; + + log.info("Found " + documentsWithVectors + " documents with vectors after sorted merge"); + + // Test vector search on sorted index + if (documentsWithVectors > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + KnnFloatVectorQuery query = + new KnnFloatVectorQuery("vector", queryVector, Math.min(10, documentsWithVectors)); + TopDocs results = searcher.search(query, 10); + + assertTrue("Should find results in sorted index", results.scoreDocs.length > 0); + log.info("Vector search on sorted index returned " + results.scoreDocs.length + " results"); + + // Verify that returned documents maintain sort order if we check their sort keys + log.info("Verifying vector search results maintain sorting consistency..."); + for (int i = 0; i < Math.min(3, results.scoreDocs.length); i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + String originalOrder = + searcher.storedFields().document(scoreDoc.doc).get("original_order"); + String sortKey = + searcher.storedFields().document(scoreDoc.doc).get(sortFieldName + "_stored"); + log.info( + "Result " + + i + + ": DocId=" + + scoreDoc.doc + + ", OriginalOrder=" + + originalOrder + + ", SortKey='" + + sortKey + + "', Score=" + + scoreDoc.score); + } + } + + log.info("Text-based index sorting verification completed successfully"); + } + } + + /** + * Test merging segments with various patterns of missing vectors + **/ + @Test + public void testMergeWithMissingVectors() throws IOException { + log.info("Starting testMergeWithMissingVectors"); + + // Randomize configuration + int maxBufferedDocs = 10 + random().nextInt(11); // 10-20 docs per buffer + int numSegments = 3 + random().nextInt(3); // 3-5 segments + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + log.info( + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments); + + int totalExpectedVectors = 0; + int totalDocuments = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int seg = 0; seg < numSegments; seg++) { + // Randomize segment characteristics + int docsInSegment = 15 + random().nextInt(16); // 15-30 docs per segment + double vectorProbability = random().nextDouble(); // 0-100% vector probability + String segmentType = "seg_" + seg + "_prob_" + String.format("%.2f", vectorProbability); + + int segmentVectorCount = 0; + + for (int i = 0; i < docsInSegment; i++) { + Document doc = new Document(); + doc.add(new StringField("id", "seg" + seg + "_" + i, Field.Store.YES)); + doc.add(new StringField("segment", segmentType, Field.Store.YES)); + doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new NumericDocValuesField("doc_in_segment", i)); + + // Randomly add vector based on segment's probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + segmentVectorCount++; + } + + writer.addDocument(doc); + } + + writer.commit(); + totalDocuments += docsInSegment; + totalExpectedVectors += segmentVectorCount; + + log.info( + "Created segment " + + seg + + ": " + + docsInSegment + + " documents, " + + segmentVectorCount + + " with vectors (probability: " + + String.format("%.2f", vectorProbability) + + ")"); + } + + // Force merge all segments + writer.forceMerge(1); + log.info("Forced merge of " + numSegments + " segments completed"); + } + + // Verify the merged index handles missing vectors correctly + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + log.info( + "Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); + + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); + + // Test vector search if we have vectors + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Randomize search parameters + int searchK = Math.min(5 + random().nextInt(10), Math.min(actualVectorCount, TOP_K_LIMIT)); + + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + + assertTrue("Should find some vector results", vectorResults.scoreDocs.length > 0); + assertTrue( + "Should not find more vectors than exist", + vectorResults.scoreDocs.length <= actualVectorCount); + + log.info( + "Found " + + vectorResults.scoreDocs.length + + " vector results out of " + + actualVectorCount + + " available"); + } else { + log.info("No vectors in merged index - skipping vector search"); + } + + log.info("Missing vectors test completed successfully"); + } + } + + /** + * Test merge behavior with document deletions + **/ + @Test + public void testMergeWithDeletions() throws IOException { + log.info("Starting testMergeWithDeletions"); + + // Randomize configuration parameters + int maxBufferedDocs = 15 + random().nextInt(11); // 15-25 docs per buffer + int numSegments = 3 + random().nextInt(4); // 3-6 segments + int docsPerSegment = 20 + random().nextInt(21); // 20-40 docs per segment + double vectorProbability = 0.7 + (random().nextDouble() * 0.25); // 70-95% have vectors + double deletionProbability = 0.2 + (random().nextDouble() * 0.3); // 20-50% deletion rate + + log.info( + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments + + ", docsPerSegment=" + + docsPerSegment + + ", vectorProbability=" + + vectorProbability + + ", deletionProbability=" + + deletionProbability); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + List expectedRemainingDocs = new ArrayList<>(); + List deletedDocs = new ArrayList<>(); + int totalDocuments = numSegments * docsPerSegment; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create multiple segments with documents + for (int seg = 0; seg < numSegments; seg++) { + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new StringField("segment", "seg_" + seg, Field.Store.YES)); + doc.add(new NumericDocValuesField("doc_num", docId)); + doc.add(new NumericDocValuesField("segment_num", seg)); + + // Randomly add vectors + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + } + + writer.addDocument(doc); + } + writer.commit(); + } + + log.info( + "Created " + + numSegments + + " segments with " + + docsPerSegment + + " documents each (" + + totalDocuments + + " total)"); + + // Delete documents randomly and track which ones are deleted + int deletedCount = 0; + for (int docId = 0; docId < totalDocuments; docId++) { + if (random().nextDouble() < deletionProbability) { + writer.deleteDocuments(new Term("id", String.valueOf(docId))); + deletedDocs.add(docId); + deletedCount++; + } else { + expectedRemainingDocs.add(docId); + } + } + + log.info( + "Deleted " + + deletedCount + + " documents (" + + String.format("%.1f", (100.0 * deletedCount / totalDocuments)) + + "%), remaining: " + + expectedRemainingDocs.size()); + + writer.commit(); + + // Force merge to apply deletions + writer.forceMerge(1); + log.info("Forced merge with deletions completed"); + } + + // Verify the merged index correctly handles deletions + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + int expectedRemaining = expectedRemainingDocs.size(); + assertEquals( + "Should have correct number of documents after deletions", + expectedRemaining, + leafReader.maxDoc()); + + // Verify that deleted documents are not present + IndexSearcher searcher = new IndexSearcher(reader); + + // Test that we can find expected remaining documents + for (int i = 0; i < Math.min(10, expectedRemainingDocs.size()); i++) { + int docId = expectedRemainingDocs.get(i); + TopDocs result = searcher.search(new TermQuery(new Term("id", String.valueOf(docId))), 1); + assertEquals("Should find remaining document " + docId, 1, (int) result.totalHits.value()); + } + + // Test that actually deleted documents are not found + int deletedDocsToCheck = Math.min(10, deletedDocs.size()); // Check up to 10 deleted docs + for (int i = 0; i < deletedDocsToCheck; i++) { + int docId = deletedDocs.get(i); + TopDocs result = searcher.search(new TermQuery(new Term("id", String.valueOf(docId))), 1); + assertEquals( + "Should not find deleted document " + docId, 0, (int) result.totalHits.value()); + } + + // Test vector search works after deletions + float[] queryVector = generateRandomVector(vectorDimension, random()); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, 10); + TopDocs vectorResults = searcher.search(vectorQuery, 10); + + assertTrue( + "Should find some vector results after deletions", vectorResults.scoreDocs.length > 0); + + log.info("Found " + vectorResults.scoreDocs.length + " vector results after deletions"); + log.info("Deletion merge verification completed successfully"); + } + } + + /** + * Test merging segments for {@link IndexType#BRUTE_FORCE} + * */ + @Test + public void testMergeBruteForceIndex() throws IOException { + log.info("Starting testMergeBruteForceIndex"); + + // Randomize configuration parameters + int maxBufferedDocs = 8 + random().nextInt(8); // 8-15 docs per buffer + int numSegments = 3 + random().nextInt(3); // 3-5 segments + int docsPerSegment = 12 + random().nextInt(9); // 12-20 docs per segment + double vectorProbability = 0.8 + (random().nextDouble() * 0.2); // 80-100% have vectors + + log.info( + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments + + ", docsPerSegment=" + + docsPerSegment + + ", vectorProbability=" + + vectorProbability); + + // Configure with brute force index type + CuVSVectorsFormat bruteForceFormat = + new CuVSVectorsFormat( + 32, // writer threads + 128, // intermediate graph degree + 64, // graph degree + IndexType.BRUTE_FORCE); // Use brute force index + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(bruteForceFormat)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + int totalDocuments = numSegments * docsPerSegment; + int totalExpectedVectors = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create multiple segments with brute force index + for (int seg = 0; seg < numSegments; seg++) { + int segmentVectorCount = 0; + + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new StringField("segment", "seg_" + seg, Field.Store.YES)); + doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new NumericDocValuesField("doc_in_segment", i)); + + // Randomly add vectors based on probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + segmentVectorCount++; + } + + writer.addDocument(doc); + } + + writer.commit(); + totalExpectedVectors += segmentVectorCount; + + log.info( + "Created brute force segment " + + seg + + ": " + + docsPerSegment + + " documents, " + + segmentVectorCount + + " with vectors"); + } + + log.info( + "Created " + + numSegments + + " brute force segments with " + + totalDocuments + + " total documents and " + + totalExpectedVectors + + " vectors"); + + // Force merge all brute force segments + writer.forceMerge(1); + log.info("Forced merge of brute force segments completed"); + } + + // Verify the merged brute force index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + log.info( + "Brute force merge results: Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); + + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); + + // Test brute force vector search (exact search) + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Search for reasonable number of results + int searchK = Math.min(8 + random().nextInt(8), Math.min(actualVectorCount, TOP_K_LIMIT)); + + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + + assertTrue( + "Should find some vector results in brute force index", + vectorResults.scoreDocs.length > 0); + assertTrue( + "Should not find more vectors than exist", + vectorResults.scoreDocs.length <= actualVectorCount); + + log.info( + "Brute force search found " + + vectorResults.scoreDocs.length + + " results out of " + + actualVectorCount + + " available vectors"); + + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + String docId = searcher.storedFields().document(scoreDoc.doc).get("id"); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); + } + } else { + log.info("No vectors in brute force merged index - skipping vector search"); + } + + log.info("Brute force merge verification completed successfully"); + } + } + + /** + * Test merging segments for {@link IndexType#CAGRA_AND_BRUTE_FORCE} + * */ + @Test + public void testMergeCagraAndBruteForceIndex() throws IOException { + log.info("Starting testMergeCagraAndBruteForceIndex"); + + // Use moderate dataset size + int maxBufferedDocs = 15 + random().nextInt(10); // 15-24 docs per buffer + int numSegments = + 4; // Fixed 4 segments: alternating CAGRA vs small segments (brute force fallback) + int docsPerSegment = 20 + random().nextInt(11); // 20-30 docs per segment + double vectorProbability = 0.9 + (random().nextDouble() * 0.1); // 90-100% have vectors + + log.info( + "Randomized parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", numSegments=" + + numSegments + + ", docsPerSegment=" + + docsPerSegment + + ", vectorProbability=" + + vectorProbability); + + // Configure with CAGRA + brute force combined index type + CuVSVectorsFormat combinedFormat = + new CuVSVectorsFormat( + 32, // writer threads + 128, // intermediate graph degree + 64, // graph degree + IndexType.CAGRA_AND_BRUTE_FORCE); // Use combined CAGRA + brute force + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(combinedFormat)) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + int totalDocuments = numSegments * docsPerSegment; + int totalExpectedVectors = 0; + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Create segments that will result in mixed index types during merge + for (int seg = 0; seg < numSegments; seg++) { + int segmentVectorCount = 0; + + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new StringField("segment", "mixed_seg_" + seg, Field.Store.YES)); + doc.add(new StringField("index_type", "cagra_and_brute_force", Field.Store.YES)); + doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new NumericDocValuesField("doc_in_segment", i)); + + // Add vectors based on probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + segmentVectorCount++; + } + + writer.addDocument(doc); + } + + writer.commit(); + totalExpectedVectors += segmentVectorCount; + + log.info( + "Created CAGRA+brute force segment " + + seg + + ": " + + docsPerSegment + + " documents, " + + segmentVectorCount + + " with vectors"); + } + + log.info( + "Created " + + numSegments + + " CAGRA+brute force segments with " + + totalDocuments + + " total documents and " + + totalExpectedVectors + + " vectors"); + + // Force merge all CAGRA+brute force segments + writer.forceMerge(1); + log.info("Forced merge of CAGRA+brute force segments completed"); + } + + // Verify the merged CAGRA+brute force index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + log.info( + "CAGRA+brute force merge results: Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); + + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); + + // Test CAGRA+brute force index vector search + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Search for reasonable number of results + int searchK = Math.min(12 + random().nextInt(8), Math.min(actualVectorCount, TOP_K_LIMIT)); + + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + + assertTrue( + "Should find some vector results in CAGRA+brute force index", + vectorResults.scoreDocs.length > 0); + assertTrue( + "Should not find more vectors than exist", + vectorResults.scoreDocs.length <= actualVectorCount); + + log.info( + "CAGRA+brute force index search found " + + vectorResults.scoreDocs.length + + " results out of " + + actualVectorCount + + " available vectors"); + + // Verify all returned documents are valid and have expected metadata + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document resultDoc = searcher.storedFields().document(scoreDoc.doc); + String docId = resultDoc.get("id"); + String indexType = resultDoc.get("index_type"); + + assertNotNull("Document should have valid ID", docId); + assertEquals( + "Document should be marked as CAGRA+brute force index type", + "cagra_and_brute_force", + indexType); + assertTrue("Score should be positive", scoreDoc.score > 0); + } + + // Test that the CAGRA+brute force index handles both approximate and exact search + // consistently + for (int trial = 0; trial < 3; trial++) { + float[] trialQueryVector = generateRandomVector(vectorDimension, random()); + KnnFloatVectorQuery trialQuery = + new KnnFloatVectorQuery("vector", trialQueryVector, Math.min(5, actualVectorCount)); + TopDocs trialResults = searcher.search(trialQuery, Math.min(5, actualVectorCount)); + + assertTrue("Trial " + trial + " should find results", trialResults.scoreDocs.length > 0); + log.info("Trial " + trial + " found " + trialResults.scoreDocs.length + " results"); + } + } else { + log.info("No vectors in CAGRA+brute force merged index - skipping vector search"); + } + + log.info("CAGRA+brute force merge verification completed successfully"); + } + } + + /** + * Test large scale merge to stress test the system + **/ + @Test + public void testLargeScaleMerge() throws IOException { + assumeTrue( + "testLargeScaleMerge requires -DlargeScale=true", + Boolean.parseBoolean(System.getProperty("largeScale", "false"))); + + log.info("Starting testLargeScaleMerge"); + + // Randomize large scale parameters + int maxBufferedDocs = 40 + random().nextInt(21); // 40-60 docs per buffer + int segmentCount = 15 + random().nextInt(11); // 15-25 segments + int docsPerSegment = 30 + random().nextInt(21); // 30-50 docs per segment + int totalDocuments = segmentCount * docsPerSegment; + + log.info( + "Randomized large scale parameters: maxBufferedDocs=" + + maxBufferedDocs + + ", segmentCount=" + + segmentCount + + ", docsPerSegment=" + + docsPerSegment + + ", totalDocuments=" + + totalDocuments); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int seg = 0; seg < segmentCount; seg++) { + log.info("Creating segment " + (seg + 1) + "/" + segmentCount); + + // Randomize vector probability per segment + double vectorProbability = + 0.5 + (random().nextDouble() * 0.4); // 50-90% vectors per segment + + for (int i = 0; i < docsPerSegment; i++) { + int docId = seg * docsPerSegment + i; + Document doc = new Document(); + doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); + doc.add(new NumericDocValuesField("segment", seg)); + doc.add(new NumericDocValuesField("position", i)); + + // Add vector based on segment's randomized probability + if (random().nextDouble() < vectorProbability) { + float[] vector = generateRandomVector(vectorDimension, random()); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + } + + writer.addDocument(doc); + } + writer.commit(); + } + + log.info("Created " + segmentCount + " segments with " + totalDocuments + " total documents"); + + // Force merge all segments + long startTime = System.currentTimeMillis(); + writer.forceMerge(1); + long mergeTime = System.currentTimeMillis() - startTime; + + log.info("Large scale merge completed in " + mergeTime + "ms"); + } + + // Verify the large merged index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Test vector search performance + var vectorValues = leafReader.getFloatVectorValues("vector"); + int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; + + if (actualVectorCount > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateRandomVector(vectorDimension, random()); + + // Randomize search parameters for large scale test + int searchK = + Math.min(20 + random().nextInt(31), Math.min(actualVectorCount, TOP_K_LIMIT)); // 20-50 + + long searchStart = System.currentTimeMillis(); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); + TopDocs vectorResults = searcher.search(vectorQuery, searchK); + long searchTime = System.currentTimeMillis() - searchStart; + + assertTrue("Should find vector results in large index", vectorResults.scoreDocs.length > 0); + log.info( + "Vector search in large index returned " + + vectorResults.scoreDocs.length + + " results out of " + + actualVectorCount + + " vectors in " + + searchTime + + "ms"); + } else { + log.info("No vectors in large merged index - skipping vector search"); + } + + log.info("Large scale merge verification completed successfully"); + } + } + + /** Helper method to generate random vectors */ + private float[] generateRandomVector(int dimension, Random random) { + float[] vector = new float[dimension]; + for (int i = 0; i < dimension; i++) { + vector[i] = (float) random().nextGaussian(); + } + // Normalize the vector + float norm = 0.0f; + for (float v : vector) { + norm += v * v; + } + norm = (float) Math.sqrt(norm); + if (norm > 0) { + for (int i = 0; i < dimension; i++) { + vector[i] /= norm; + } + } + return vector; + } + + /** Helper method to generate random text strings for sorting */ + private String generateRandomText(Random random, int length) { + StringBuilder sb = new StringBuilder(length); + String chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; + for (int i = 0; i < length; i++) { + sb.append(chars.charAt(random().nextInt(chars.length()))); + } + return sb.toString(); + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java new file mode 100644 index 00000000..8bd8339f --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -0,0 +1,50 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import java.util.Random; + +public class TestUtils { + + public static float[][] generateDataset(Random random, int size, int dimensions) { + float[][] dataset = new float[size][dimensions]; + for (int i = 0; i < size; i++) { + for (int j = 0; j < dimensions; j++) { + dataset[i][j] = random.nextFloat() * 100; + } + } + return dataset; + } + + public static float[] generateRandomVector(int dimensions, Random random) { + float[] vector = new float[dimensions]; + for (int i = 0; i < dimensions; i++) { + vector[i] = random.nextFloat() * 100; + } + return vector; + } + + public static float[][] generateQueries(Random random, int dimensions, int numQueries) { + // Generate random query vectors + float[][] queries = new float[numQueries][dimensions]; + for (int i = 0; i < numQueries; i++) { + for (int j = 0; j < dimensions; j++) { + queries[i][j] = random.nextFloat() * 100; + } + } + return queries; + } +} From 87070484e2ddefa55c35a62ba3a286229df64185 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 14 Aug 2025 22:44:18 -0400 Subject: [PATCH 02/21] Cagra to HNSW serialization and searching feature > Co-authored-by: Ishan Chattopadhyaya > Co-authored-by: Puneet Ahuja --- .gitignore | 3 +- pom.xml | 21 +- .../cuvs/lucene/CuVSCPUSearchCodec.java | 72 ++ .../com/nvidia/cuvs/lucene/CuVSCodec.java | 5 +- .../nvidia/cuvs/lucene/CuVSVectorsFormat.java | 39 +- .../nvidia/cuvs/lucene/CuVSVectorsWriter.java | 803 ++++++++++++++---- .../nvidia/cuvs/lucene/OnHeapHnswGraph.java | 261 ++++++ .../java/com/nvidia/cuvs/lucene/Utils.java | 5 + .../services/org.apache.lucene.codecs.Codec | 2 + ...TestCagraToHnswSerializationAndSearch.java | 169 ++++ .../com/nvidia/cuvs/lucene/TestMerge.java | 2 + 11 files changed, 1214 insertions(+), 168 deletions(-) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java create mode 100644 src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java create mode 100644 src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java diff --git a/.gitignore b/.gitignore index 63d9e5e0..43a99397 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ target -**/.DS_Store \ No newline at end of file +**/.DS_Store +bin diff --git a/pom.xml b/pom.xml index c79f26ee..5c1969f7 100644 --- a/pom.xml +++ b/pom.xml @@ -68,7 +68,7 @@ com.nvidia.cuvs cuvs-java - 25.8.0-4f53f-SNAPSHOT + 25.8.0-c9599-SNAPSHOT @@ -103,6 +103,25 @@ + + org.apache.maven.plugins + maven-assembly-plugin + 3.6.0 + + + jar-with-dependencies + + + + + make-assembly + package + + single + + + + diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java new file mode 100644 index 00000000..f2a679f5 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java @@ -0,0 +1,72 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.LibraryException; +import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene101.Lucene101Codec; + +/** CuVS based codec for GPU based vector search */ +public class CuVSCPUSearchCodec extends FilterCodec { + + public CuVSCPUSearchCodec() { + this("CuVSCPUSearchCodec", new Lucene101Codec()); + } + + public CuVSCPUSearchCodec(String name, Codec delegate) { + super(name, delegate); + initializeFormat(); + } + + public CuVSCPUSearchCodec( + int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { + this("CuVSCPUSearchCodec", new Lucene101Codec()); + initializeFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); + } + + private void initializeFormat() { + initializeFormat(1, 128, 64, 1); // Default values + } + + private void initializeFormat( + int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { + KnnVectorsFormat format; + try { + format = + new CuVSVectorsFormat( + cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, IndexType.HNSW_LUCENE); + setKnnFormat(format); + } catch (LibraryException ex) { + Logger log = Logger.getLogger(CuVSCodec.class.getName()); + log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); + } + } + + KnnVectorsFormat knnFormat = null; + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return knnFormat; + } + + public void setKnnFormat(KnnVectorsFormat format) { + this.knnFormat = format; + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java index ec861765..8b99ec2a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java @@ -34,8 +34,9 @@ public CuVSCodec(String name, Codec delegate) { super(name, delegate); KnnVectorsFormat format; try { - // TODO: Remove this hard coded values. - format = new CuVSVectorsFormat(1, 128, 64, IndexType.CAGRA); + // TODO: The hard-coded values passed below should be configurable. + // To make relevant changes in a subsequent PR. + format = new CuVSVectorsFormat(1, 128, 64, 1, IndexType.CAGRA); setKnnFormat(format); } catch (LibraryException ex) { Logger log = Logger.getLogger(CuVSCodec.class.getName()); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java index f82198b1..6048e717 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java @@ -21,16 +21,18 @@ import java.io.IOException; import java.util.logging.Logger; import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; /** CuVS based KnnVectorsFormat for GPU acceleration */ public class CuVSVectorsFormat extends KnnVectorsFormat { - private static final Logger LOG = Logger.getLogger(CuVSVectorsFormat.class.getName()); + private static final Logger log = Logger.getLogger(CuVSVectorsFormat.class.getName()); // TODO: fix Lucene version in name, to the final targeted release, if any static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; @@ -45,6 +47,7 @@ public class CuVSVectorsFormat extends KnnVectorsFormat { public static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; public static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + public static final int HNSW_GRAPH_LAYERS = 1; static CuVSResources resources = cuVSResourcesOrNull(); @@ -56,6 +59,7 @@ public class CuVSVectorsFormat extends KnnVectorsFormat { final int cuvsWriterThreads; final int intGraphDegree; final int graphDegree; + final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion final CuVSVectorsWriter.IndexType indexType; // the index type to build, when writing /** @@ -68,6 +72,7 @@ public CuVSVectorsFormat() { DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, + HNSW_GRAPH_LAYERS, DEFAULT_INDEX_TYPE); } @@ -77,11 +82,16 @@ public CuVSVectorsFormat() { * @throws LibraryException if the native library fails to load */ public CuVSVectorsFormat( - int cuvsWriterThreads, int intGraphDegree, int graphDegree, IndexType indexType) { + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + IndexType indexType) { super("CuVSVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.hnswLayers = hnswLayers; this.indexType = indexType; } @@ -90,12 +100,12 @@ private static CuVSResources cuVSResourcesOrNull() { resources = CuVSResources.create(); return resources; } catch (UnsupportedOperationException uoe) { - LOG.warning("cuvs is not supported on this platform or java version: " + uoe.getMessage()); + log.warning("cuvs is not supported on this platform or java version: " + uoe.getMessage()); } catch (Throwable t) { if (t instanceof ExceptionInInitializerError ex) { t = ex.getCause(); } - LOG.warning("Exception occurred during creation of cuvs resources. " + t); + log.warning("Exception occurred during creation of cuvs resources. " + t); } return null; } @@ -116,14 +126,28 @@ public CuVSVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio checkSupported(); var flatWriter = flatVectorsFormat.fieldsWriter(state); return new CuVSVectorsWriter( - state, cuvsWriterThreads, intGraphDegree, graphDegree, indexType, resources, flatWriter); + state, + cuvsWriterThreads, + intGraphDegree, + graphDegree, + hnswLayers, + indexType, + resources, + flatWriter); } @Override - public CuVSVectorsReader fieldsReader(SegmentReadState state) throws IOException { + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { checkSupported(); var flatReader = flatVectorsFormat.fieldsReader(state); - return new CuVSVectorsReader(state, resources, flatReader); + if (this.indexType == IndexType.HNSW_LUCENE) { + log.info("Using Reader: Lucene99HnswVectorsReader"); + return new Lucene99HnswVectorsReader(state, flatReader); + } else { + checkSupported(); + log.info("Using Reader: CuVSVectorsReader"); + return new CuVSVectorsReader(state, resources, flatReader); + } } @Override @@ -137,6 +161,7 @@ public String toString() { sb.append("cuvsWriterThreads=").append(cuvsWriterThreads); sb.append("intGraphDegree=").append(intGraphDegree); sb.append("graphDegree=").append(graphDegree); + sb.append("hnswLayers=").append(hnswLayers); sb.append("resources=").append(resources); sb.append(")"); return sb.toString(); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java index 0bc3eca8..47627c4f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java @@ -31,17 +31,15 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSMatrix.DataType; import com.nvidia.cuvs.CuVSResources; import java.io.IOException; import java.io.OutputStream; import java.nio.file.Files; import java.nio.file.Path; -import java.time.Duration; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Objects; -import java.util.function.Supplier; import java.util.logging.Logger; import java.util.stream.IntStream; import org.apache.lucene.codecs.CodecUtil; @@ -50,6 +48,7 @@ import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; @@ -65,6 +64,10 @@ import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.apache.lucene.util.packed.DirectMonotonicWriter; /** * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into @@ -87,32 +90,39 @@ public class CuVSVectorsWriter extends KnnVectorsWriter { private final int cuvsWriterThreads; private final int intGraphDegree; private final int graphDegree; + private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion private final CuVSResources resources; private final IndexType indexType; private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors private final List fields = new ArrayList<>(); - private final IndexOutput meta, cuvsIndex; + private IndexOutput meta = null, cuvsIndex = null; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; private final InfoStream infoStream; private boolean finished; + private String vemFileName; + private String vexFileName; /** The CuVS index Type. */ public enum IndexType { /** Builds a Cagra index. */ - CAGRA(true, false, false), + CAGRA(true, false, false, false), /** Builds a Brute Force index. */ - BRUTE_FORCE(false, true, false), + BRUTE_FORCE(false, true, false, false), /** Builds an HSNW index - suitable for searching on CPU. */ - HNSW(false, false, true), + HNSW(false, false, true, false), /** Builds a Cagra and a Brute Force index. */ - CAGRA_AND_BRUTE_FORCE(true, true, false); - private final boolean cagra, bruteForce, hnsw; + CAGRA_AND_BRUTE_FORCE(true, true, false, false), + /** Builds a Lucene HNSW index via CAGRA. */ + HNSW_LUCENE(false, false, false, true); + private final boolean cagra, bruteForce, hnsw, hnswLucene; - IndexType(boolean cagra, boolean bruteForce, boolean hnsw) { + IndexType(boolean cagra, boolean bruteForce, boolean hnsw, boolean hnswLucene) { this.cagra = cagra; this.bruteForce = bruteForce; this.hnsw = hnsw; + this.hnswLucene = hnswLucene; } public boolean cagra() { @@ -126,6 +136,10 @@ public boolean bruteForce() { public boolean hnsw() { return hnsw; } + + public boolean hnswLucene() { + return hnswLucene; + } } public CuVSVectorsWriter( @@ -133,6 +147,7 @@ public CuVSVectorsWriter( int cuvsWriterThreads, int intGraphDegree, int graphDegree, + int hnswLayers, IndexType indexType, CuVSResources resources, FlatVectorsWriter flatVectorsWriter) @@ -142,10 +157,17 @@ public CuVSVectorsWriter( this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.hnswLayers = hnswLayers; this.resources = resources; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + vemFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "vem"); + + vexFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "vex"); + String metaFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, CUVS_META_CODEC_EXT); @@ -154,20 +176,42 @@ public CuVSVectorsWriter( boolean success = false; try { - meta = state.directory.createOutput(metaFileName, state.context); - cuvsIndex = state.directory.createOutput(cagraFileName, state.context); - CodecUtil.writeIndexHeader( - meta, - CUVS_META_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - CodecUtil.writeIndexHeader( - cuvsIndex, - CUVS_INDEX_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); + + // Only create CAGRA files if not in HNSW_LUCENE mode + if (indexType == IndexType.HNSW_LUCENE) { + + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + + CodecUtil.writeIndexHeader( + hnswMeta, + "Lucene99HnswVectorsFormatMeta", + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + "Lucene99HnswVectorsFormatIndex", + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + } else { + // Only create CAGRA files if not in HNSW_LUCENE mode + meta = state.directory.createOutput(metaFileName, state.context); + cuvsIndex = state.directory.createOutput(cagraFileName, state.context); + CodecUtil.writeIndexHeader( + meta, + CUVS_META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + cuvsIndex, + CUVS_INDEX_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + } success = true; } finally { if (success == false) { @@ -214,25 +258,554 @@ private CagraIndexParams cagraIndexParams(int size) { .build(); } - static long nanosToMillis(long nanos) { - return Duration.ofNanos(nanos).toMillis(); - } - private void info(String msg) { if (infoStream.isEnabled(CUVS_COMPONENT)) { infoStream.message(CUVS_COMPONENT, msg); } } + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo); + return; + } + long cagraIndexOffset, cagraIndexLength = 0L; + long bruteForceIndexOffset, bruteForceIndexLength = 0L; + long hnswIndexOffset, hnswIndexLength = 0L; + + // workaround for the minimum number of vectors for Cagra + IndexType indexType = + this.indexType.cagra() && vectors.size() < MIN_CAGRA_INDEX_SIZE + ? IndexType.BRUTE_FORCE + : this.indexType; + + info("=== INDEX TYPE DEBUG: original=" + this.indexType + ", effective=" + indexType + " ==="); + + try { + if (indexType.hnswLucene()) { + info("=== ENTERED HNSW_LUCENE BLOCK (HNSW-only mode) ==="); + info("Entered the writeFieldInternal's HNSW LUCENE block - writing only HNSW files"); + try { + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeHnswOnlyIndex(dataset, fieldInfo, vectors); + } catch (Throwable t) { + info("=== ERROR IN HNSW_LUCENE: " + t.getMessage() + " ==="); + handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + // workaround for cuVS issue + indexType = IndexType.BRUTE_FORCE; + } + // For HNSW_LUCENE, we don't write any CAGRA data, so set lengths to 0 + cagraIndexLength = 0L; + cagraIndexOffset = 0L; + bruteForceIndexOffset = 0L; + bruteForceIndexLength = 0L; + hnswIndexOffset = 0L; + hnswIndexLength = 0L; + } else { + cagraIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.cagra()) { + try { + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeCagraIndex(cagraIndexOutputStream, dataset); + } catch (Throwable t) { + handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + // workaround for cuVS issue + indexType = IndexType.BRUTE_FORCE; + } + cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + } + + bruteForceIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.bruteForce()) { + var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeBruteForceIndex(bruteForceIndexOutputStream, dataset); + bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; + } + + hnswIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.hnsw()) { + var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + if (vectors.size() > MIN_CAGRA_INDEX_SIZE) { + try { + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeHNSWIndex(hnswIndexOutputStream, dataset); + } catch (Throwable t) { + handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + } + } + hnswIndexLength = cuvsIndex.getFilePointer() - hnswIndexOffset; + } + } + + // Only write meta for non-HNSW_LUCENE modes + if (indexType != IndexType.HNSW_LUCENE) { + writeMeta( + fieldInfo, + vectors.size(), + cagraIndexOffset, + cagraIndexLength, + bruteForceIndexOffset, + bruteForceIndexLength, + hnswIndexOffset, + hnswIndexLength); + } + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + private void writeHnswOnlyIndex( + CuVSMatrix dataset, FieldInfo fieldInfo, List originalVectors) throws Throwable { + if (dataset.size() < 2) { + throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); + } + CagraIndexParams params = cagraIndexParams((int) dataset.size()); + long startTime = System.nanoTime(); + CagraIndex index = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); + + // Get the adjacency list from CAGRA index + int[][] adjacencyList; + try { + adjacencyList = index.getGraph(); + info("=== SUCCESS: Got adjacency list from CAGRA index ==="); + info("Successfully got adjacency list from CAGRA index"); + } catch (Exception e) { + info("=== FAILED: getGraph() method failed: " + e.getMessage() + " ==="); + info("getGraph() method failed or doesn't exist: " + e.getMessage()); + // Create a mock adjacency list for testing + int size = (int) dataset.size(); + adjacencyList = new int[size][]; + for (int i = 0; i < size; i++) { + // Create connections to next few nodes (circular) + int degree = Math.min(10, size - 1); // up to 10 connections + adjacencyList[i] = new int[degree]; + for (int j = 0; j < degree; j++) { + adjacencyList[i][j] = (i + j + 1) % size; + } + } + info( + "=== CREATED MOCK ADJACENCY LIST: " + + size + + " nodes, degree=" + + (adjacencyList.length > 0 ? adjacencyList[0].length : 0) + + " ==="); + info( + "Created mock adjacency list with " + + size + + " nodes, degree=" + + (adjacencyList.length > 0 ? adjacencyList[0].length : 0)); + } + + int size = (int) dataset.size(); + int dimensions = fieldInfo.getVectorDimension(); + + // Debug: Check if we got valid adjacency data + info( + "Adjacency list info: " + + (adjacencyList == null + ? "null" + : "length=" + + adjacencyList.length + + ", first row=" + + (adjacencyList.length > 0 && adjacencyList[0] != null + ? adjacencyList[0].length + : "null"))); + + // Create HNSW graph from CAGRA - multi-layer if original vectors available + OnHeapHnswGraph hnswGraph; + if (originalVectors != null && originalVectors.size() > 0) { + info("=== Creating 3-layer HNSW graph using original vectors ==="); + hnswGraph = + createMultiLayerHnswGraph(fieldInfo, size, dimensions, adjacencyList, originalVectors); + } else { + info("=== Creating single-layer HNSW graph (no original vectors) ==="); + // Create single layer graph + List singleLayerNodes = new ArrayList<>(); + List singleLayerAdjacencies = new ArrayList<>(); + singleLayerAdjacencies.add(adjacencyList); + hnswGraph = new OnHeapHnswGraph(size, dimensions, singleLayerNodes, singleLayerAdjacencies); + } + + // Remember the vector index offset before writing + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + + // Calculate the length of written data + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets); + + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); + info("HNSW-only graph created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); + + // Don't serialize CAGRA index - destroy it immediately + index.destroyIndex(); + } + + /** + * 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 + */ + private OnHeapHnswGraph createMultiLayerHnswGraph( + FieldInfo fieldInfo, + int size, + int dimensions, + int[][] adjacencyList, + List originalVectors) + throws Throwable { + + // Calculate M as cagraGraphDegree/2 + int M = graphDegree / 2; + info( + "=== Creating " + + hnswLayers + + "-layer HNSW graph with M=" + + M + + " (cagraGraphDegree/2) ==="); + info("Creating " + hnswLayers + "-layer HNSW graph with size=" + size + ", M=" + M); + + // Store all layers data + java.util.List layerNodes = new java.util.ArrayList<>(); + java.util.List layerAdjacencies = new java.util.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(adjacencyList); + + // Build higher layers - create exactly hnswLayers-1 additional layers (layer 0 is already + // added) + int currentLayerSize = size; + int layerIndex = 1; + java.util.Random random = new java.util.Random(42); // Fixed seed for reproducibility + + while (layerIndex < hnswLayers && currentLayerSize > 1) { + // Calculate size for next layer (1/M of current layer) + int nextLayerSize = Math.max(1, currentLayerSize / M); + + info( + "=== Layer " + + layerIndex + + " will have " + + nextLayerSize + + " nodes out of " + + currentLayerSize + + " (previous layer) ==="); + info("Layer " + layerIndex + " will have " + nextLayerSize + " nodes"); + + // Select nodes for this layer + java.util.Set selectedNodesSet = new java.util.HashSet<>(); + + 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); + + info( + "=== Selected Layer " + + layerIndex + + " nodes: " + + java.util.Arrays.toString( + java.util.Arrays.copyOf(selectedNodes, Math.min(10, selectedNodes.length))) + + (selectedNodes.length > 10 ? "..." : "") + + " ==="); + + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + int nodeId = selectedNodes[i]; + if (nodeId < originalVectors.size()) { + selectedVectors[i] = originalVectors.get(nodeId); + } else { + selectedVectors[i] = createRandomVector(dimensions, nodeId); + } + } + + // Build CAGRA graph for this layer + int[][] layerAdjacency = buildCagraGraphForSubset(selectedVectors, selectedNodes); + layerAdjacencies.add(layerAdjacency); + + // Update for next iteration + currentLayerSize = nextLayerSize; + layerIndex++; + + // Use different seed for each layer + random = new java.util.Random(42 + layerIndex); + } + + int numLayers = layerAdjacencies.size(); + info("=== Total layers created: " + numLayers + " ==="); + info("Created " + numLayers + " layers total"); + + // Create the multi-layer graph with all layers + return new OnHeapHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } + + /** + * Creates a random vector for fallback purposes + */ + private float[] createRandomVector(int dimensions, int seed) { + float[] vector = new float[dimensions]; + java.util.Random random = new java.util.Random(seed); + for (int i = 0; i < dimensions; i++) { + vector[i] = random.nextFloat(); + } + return vector; + } + + /** + * Builds a CAGRA graph for a subset of vectors + */ + private int[][] buildCagraGraphForSubset(float[][] vectors, int[] originalNodeIds) + throws Throwable { + if (vectors.length < 2) { + // Can't build CAGRA graph with less than 2 vectors + return new int[vectors.length][0]; + } + + try { + // Create CuVSMatrix from the subset vectors + CuVSMatrix subsetDataset = CuVSMatrix.ofArray(vectors); + + // Build CAGRA index for the subset + CagraIndexParams params = cagraIndexParams(vectors.length); + CagraIndex subsetIndex = + CagraIndex.newBuilder(resources) + .withDataset(subsetDataset) + .withIndexParams(params) + .build(); + + // Get adjacency list from subset CAGRA index + int[][] subsetAdjacency; + try { + subsetAdjacency = subsetIndex.getGraph(); + info("=== SUCCESS: Got adjacency list from Layer 1 CAGRA index ==="); + info("Successfully got adjacency list from Layer 1 CAGRA index"); + } catch (Exception e) { + info("=== FAILED: getGraph() method failed for Layer 1: " + e.getMessage() + " ==="); + info("getGraph() method failed for Layer 1: " + e.getMessage()); + // Create mock adjacency list + subsetAdjacency = new int[vectors.length][]; + for (int i = 0; i < vectors.length; i++) { + int degree = Math.min(5, vectors.length - 1); + subsetAdjacency[i] = new int[degree]; + for (int j = 0; j < degree; j++) { + subsetAdjacency[i][j] = (i + j + 1) % vectors.length; + } + } + } + + // Convert subset adjacency to use original node IDs + int[][] layer1Adjacency = new int[vectors.length][]; + for (int i = 0; i < vectors.length; i++) { + if (subsetAdjacency[i] != null) { + layer1Adjacency[i] = new int[subsetAdjacency[i].length]; + for (int j = 0; j < subsetAdjacency[i].length; j++) { + // Map subset index back to original node ID + int subsetNeighborId = subsetAdjacency[i][j]; + layer1Adjacency[i][j] = originalNodeIds[subsetNeighborId]; + } + } else { + layer1Adjacency[i] = new int[0]; + } + } + + subsetIndex.destroyIndex(); + return layer1Adjacency; + + } catch (Exception e) { + info("=== FAILED to build CAGRA graph for subset: " + e.getMessage() + " ==="); + info("Failed to build CAGRA graph for subset: " + e.getMessage()); + + // Fallback: create simple connections between Layer 1 nodes + int[][] fallbackAdjacency = new int[vectors.length][]; + for (int i = 0; i < vectors.length; i++) { + int degree = Math.min(3, vectors.length - 1); + fallbackAdjacency[i] = new int[degree]; + for (int j = 0; j < degree; j++) { + int neighborIdx = (i + j + 1) % vectors.length; + fallbackAdjacency[i][j] = originalNodeIds[neighborIdx]; + } + } + return fallbackAdjacency; + } + } + + private void writeMeta( + IndexOutput vectorIndex, + IndexOutput meta, + FieldInfo field, + long vectorIndexOffset, + long vectorIndexLength, + int count, + HnswGraph graph, + int[][] graphLevelNodeOffsets) + throws IOException { + info( + "=== writeMeta: Writing field " + + field.name + + " with count=" + + count + + ", dimensions=" + + field.getVectorDimension() + + " ==="); + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); + meta.writeVLong(vectorIndexOffset); + meta.writeVLong(vectorIndexLength); + meta.writeVInt(field.getVectorDimension()); + meta.writeInt(count); + // Use M = cagraGraphDegree/2 + int M = graphDegree / 2; + info("=== writeMeta: Writing M=" + M + " (cagraGraphDegree/2) ==="); + meta.writeVInt(M); // M = cagraGraphDegree/2 + // write graph nodes on each level + if (graph == null) { + meta.writeVInt(0); + } else { + meta.writeVInt(graph.numLevels()); + long valueCount = 0; + for (int level = 0; level < graph.numLevels(); level++) { + NodesIterator nodesOnLevel = graph.getNodesOnLevel(level); + valueCount += nodesOnLevel.size(); + if (level > 0) { + int[] nol = new int[nodesOnLevel.size()]; + int numberConsumed = nodesOnLevel.consume(nol); + Arrays.sort(nol); + assert numberConsumed == nodesOnLevel.size(); + meta.writeVInt(nol.length); // number of nodes on a level + for (int i = nodesOnLevel.size() - 1; i > 0; --i) { + nol[i] -= nol[i - 1]; + } + for (int n : nol) { + assert n >= 0 : "delta encoding for nodes failed; expected nodes to be sorted"; + meta.writeVInt(n); + } + } else { + assert nodesOnLevel.size() == count : "Level 0 expects to have all nodes"; + } + } + long start = vectorIndex.getFilePointer(); + meta.writeLong(start); + meta.writeVInt(16); // DIRECT_MONOTONIC_BLOCK_SHIFT); + final DirectMonotonicWriter memoryOffsetsWriter = + DirectMonotonicWriter.getInstance( + meta, vectorIndex, valueCount, 16); // DIRECT_MONOTONIC_BLOCK_SHIFT); + long cumulativeOffsetSum = 0; + int totalOffsetsWritten = 0; + for (int[] levelOffsets : graphLevelNodeOffsets) { + info( + "=== writeMeta: Writing offsets for level with " + + levelOffsets.length + + " entries ==="); + for (int v : levelOffsets) { + memoryOffsetsWriter.add(cumulativeOffsetSum); + cumulativeOffsetSum += v; + totalOffsetsWritten++; + } + } + info( + "=== writeMeta: Total offsets written: " + + totalOffsetsWritten + + ", expected: " + + valueCount + + " ==="); + memoryOffsetsWriter.finish(); + meta.writeLong(vectorIndex.getFilePointer() - start); + } + } + + private int[][] writeGraph(OnHeapHnswGraph graph, IndexOutput vectorIndex) throws IOException { + if (graph == null) return new int[0][0]; + // write vectors' neighbors on each level into the vectorIndex file + 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[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); + offsets[level] = new int[sortedNodes.length]; + int nodeOffsetId = 0; + // Debug: print the actual number of nodes being processed + info( + "=== writeGraph: Level " + + level + + " has " + + sortedNodes.length + + " nodes, expected " + + (level == 0 ? countOnLevel0 : "unknown") + + " ==="); + for (int node : sortedNodes) { + NeighborArray neighbors = graph.getNeighbors(level, node); + int size = neighbors.size(); + // Write size in VInt as the neighbors list is typically small + long offsetStart = vectorIndex.getFilePointer(); + int[] nnodes = neighbors.nodes(); + 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; + } + 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]; + } + // Write the size after duplicates are removed + vectorIndex.writeVInt(actualSize); + for (int i = 0; i < actualSize; i++) { + vectorIndex.writeVInt(scratch[i]); + } + offsets[level][nodeOffsetId++] = + Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); + } + } + return offsets; + } + private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { if (dataset.size() < 2) { throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); } CagraIndexParams params = cagraIndexParams((int) dataset.size()); long startTime = System.nanoTime(); - var index = + CagraIndex index = CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); - long elapsedMillis = nanosToMillis(System.nanoTime() - startTime); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("Cagra index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); Path tmpFile = Files.createTempFile(resources.tempDirectory(), "tmpindex", "cag"); index.serialize(os, tmpFile); @@ -248,7 +821,7 @@ private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Th long startTime = System.nanoTime(); var index = BruteForceIndex.newBuilder(resources).withIndexParams(params).withDataset(dataset).build(); - long elapsedMillis = nanosToMillis(System.nanoTime() - startTime); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("bf index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); index.serialize(os); index.destroyIndex(); @@ -260,9 +833,9 @@ private void writeHNSWIndex(OutputStream os, CuVSMatrix dataset) throws Throwabl } CagraIndexParams indexParams = cagraIndexParams((int) dataset.size()); long startTime = System.nanoTime(); - var index = + CagraIndex index = CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(indexParams).build(); - long elapsedMillis = nanosToMillis(System.nanoTime() - startTime); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("HNSW index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); Path tmpFile = Files.createTempFile("tmpindex", "hnsw"); index.serializeToHNSW(os, tmpFile); @@ -282,95 +855,25 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { } private void writeField(CuVSFieldWriter fieldData) throws IOException { - // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors - // from the flat writer to the CuVSMatrix? - List vectors = fieldData.getVectors(); - writeFieldInternal( - fieldData.fieldInfo(), - () -> Utils.createFloatMatrix(vectors, fieldData.fieldInfo().getVectorDimension()), - vectors.size()); + writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); } private void writeSortingField(CuVSFieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors // from the flat writer to the CuVSMatrix? + + // TODO: This is slightly different.... List sortedVectors = new ArrayList(); for (int i = 0; i < fieldData.getVectors().size(); i++) { sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); } - writeFieldInternal( - fieldData.fieldInfo(), - () -> Utils.createFloatMatrix(sortedVectors, fieldData.fieldInfo().getVectorDimension()), - sortedVectors.size()); - } - - private void writeFieldInternal( - FieldInfo fieldInfo, Supplier datasetSupplier, int datasetSize) - throws IOException { - if (datasetSize == 0) { - writeEmpty(fieldInfo); - return; - } - long cagraIndexOffset, cagraIndexLength = 0L; - long bruteForceIndexOffset, bruteForceIndexLength = 0L; - long hnswIndexOffset, hnswIndexLength = 0L; - - // workaround for the minimum number of vectors for Cagra - IndexType indexType = - this.indexType.cagra() && datasetSize < MIN_CAGRA_INDEX_SIZE - ? IndexType.BRUTE_FORCE - : this.indexType; - try { - cagraIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.cagra()) { - try { - var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - writeCagraIndex(cagraIndexOutputStream, datasetSupplier.get()); - } catch (Throwable t) { - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - // workaround for cuVS issue - indexType = IndexType.BRUTE_FORCE; - } - cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; - } - - bruteForceIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.bruteForce()) { - var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - writeBruteForceIndex(bruteForceIndexOutputStream, datasetSupplier.get()); - bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; - } - - hnswIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.hnsw()) { - var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - if (datasetSize > MIN_CAGRA_INDEX_SIZE) { - try { - writeHNSWIndex(hnswIndexOutputStream, datasetSupplier.get()); - } catch (Throwable t) { - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - } - } - hnswIndexLength = cuvsIndex.getFilePointer() - hnswIndexOffset; - } - - writeMeta( - fieldInfo, - (int) datasetSize, - cagraIndexOffset, - cagraIndexLength, - bruteForceIndexOffset, - bruteForceIndexLength, - hnswIndexOffset, - hnswIndexLength); - } catch (Throwable t) { - Utils.handleThrowable(t); - } + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); } private void writeEmpty(FieldInfo fieldInfo) throws IOException { @@ -464,6 +967,21 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw } } + /** + * Creates List from merged vectors + * */ + private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List res = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + int ordinal = iter.index(); + float[] vector = mergedVectorValues.vectorValue(ordinal); + res.add(vector); + } + return res; + } + /** * Fallback method that rebuilds indexes from merged vectors. * Used when native CAGRA merge() is not possible. Also used @@ -474,59 +992,16 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws throw new AssertionError("Only Float32 supported"); } try { - // We need to compute the size of the number of merged documents up-front so that we can - // compute the CuVSMatrix capacity. TODO: Find a way to do this without merging twice. - final int numMergedDocs = getMergedDocsCount(fieldInfo, mergeState); - if (numMergedDocs != 0) { - writeFieldInternal( - fieldInfo, - () -> { - try { - return createMatrixFromMergedVectors( - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues( - fieldInfo, mergeState), - numMergedDocs); - } catch (IOException e) { - throw new RuntimeException(e); - } - }, - numMergedDocs); - } else { - writeEmpty(fieldInfo); - } + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); } catch (Throwable t) { Utils.handleThrowable(t); } } - private int getMergedDocsCount(FieldInfo fieldInfo, MergeState mergeState) throws IOException { - KnnVectorValues.DocIndexIterator iter = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) - .iterator(); - int numMergedDocs = 0; - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - numMergedDocs++; - } - return numMergedDocs; - } - - /** - * Creates CuVSMatrix from merged vectors - * */ - private CuVSMatrix createMatrixFromMergedVectors( - FloatVectorValues mergedVectorValues, int numMergedDocs) throws IOException { - CuVSMatrix.Builder builder = - CuVSMatrix.builder(numMergedDocs, mergedVectorValues.dimension(), DataType.FLOAT); - KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - int ordinal = iter.index(); - float[] vector = mergedVectorValues.vectorValue(ordinal); - builder.addVector(vector.clone()); - } - return builder.build(); - } - /** * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. */ @@ -565,7 +1040,10 @@ private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; // Write metadata (assuming no brute force or HNSW indexes for merged result) - writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); + // Only write meta for non-HNSW_LUCENE modes + if (indexType != IndexType.HNSW_LUCENE) { + writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); + } // Clean up the merged index mergedIndex.destroyIndex(); @@ -620,11 +1098,22 @@ public void finish() throws IOException { if (cuvsIndex != null) { CodecUtil.writeFooter(cuvsIndex); } + + { + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } } @Override public void close() throws IOException { - IOUtils.close(meta, cuvsIndex, flatVectorsWriter); + IOUtils.close(meta, cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java b/src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java new file mode 100644 index 00000000..b204eaf9 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java @@ -0,0 +1,261 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import java.util.ArrayList; +import java.util.List; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.NeighborArray; + +public class OnHeapHnswGraph extends HnswGraph { + + private final int size; + private final int dimensions; + private final int numLevels; + + // Store layers data - each layer has its own nodes and adjacency lists + private final List layerNodes; + private final List layerNeighbors; + + // Layer 0 is special - it contains all nodes + private final NeighborArray[] layer0Neighbors; + + // Multi-layer constructor that supports arbitrary number of layers + public OnHeapHnswGraph( + int size, int dimensions, List layerNodes, List layerAdjacencies) { + + this.size = size; + this.dimensions = dimensions; + this.numLevels = layerAdjacencies.size(); + this.layerNodes = new ArrayList<>(); + this.layerNeighbors = new ArrayList<>(); + + // Process Layer 0 (base layer with all nodes) + int[][] layer0Adjacency = layerAdjacencies.get(0); + this.layer0Neighbors = new NeighborArray[size]; + + for (int i = 0; i < size; i++) { + if (layer0Adjacency[i] != null && layer0Adjacency[i].length > 0) { + layer0Neighbors[i] = new NeighborArray(layer0Adjacency[i].length, true); + for (int j = 0; j < layer0Adjacency[i].length; j++) { + layer0Neighbors[i].addInOrder(layer0Adjacency[i][j], 1.0f - (j * 0.001f)); + } + } else { + layer0Neighbors[i] = new NeighborArray(0, true); + } + } + + // Process higher layers (1 to numLevels-1) + for (int level = 1; level < numLevels; level++) { + int[] nodes = layerNodes.get(level); + int[][] adjacency = layerAdjacencies.get(level); + + this.layerNodes.add(nodes); + NeighborArray[] neighbors = new NeighborArray[nodes.length]; + + for (int i = 0; i < nodes.length; i++) { + if (adjacency[i] != null && adjacency[i].length > 0) { + neighbors[i] = new NeighborArray(adjacency[i].length, true); + for (int j = 0; j < adjacency[i].length; j++) { + neighbors[i].addInOrder(adjacency[i][j], 1.0f - (j * 0.001f)); + } + } else { + neighbors[i] = new NeighborArray(0, true); + } + } + + this.layerNeighbors.add(neighbors); + } + } + + public int size() { + return size; + } + + public int numLevels() { + return numLevels; + } + + public NodesIterator getNodesOnLevel(int level) { + if (level == 0) { + return new ArrayNodesIterator(size); + } else if (level > 0 && level < numLevels) { + int[] nodes = layerNodes.get(level - 1); + return new SpecificNodesIterator(nodes); + } else { + return new ArrayNodesIterator(0); + } + } + + public NeighborArray getNeighbors(int level, int node) { + if (level == 0 && node < size) { + return layer0Neighbors[node]; + } else if (level > 0 && level < numLevels) { + int[] nodes = layerNodes.get(level - 1); + NeighborArray[] neighbors = layerNeighbors.get(level - 1); + + // Find the index of this node in the layer + for (int i = 0; i < nodes.length; i++) { + if (nodes[i] == node) { + return neighbors[i]; + } + } + } + return null; + } + + // Implementation of abstract methods from HnswGraph + private int currentNode = -1; + private int currentLevel = -1; + private int neighborIndex = -1; + + @Override + public void seek(int level, int target) { + currentLevel = level; + currentNode = target; + neighborIndex = -1; + } + + @Override + public int nextNeighbor() { + if (currentLevel == 0 + && currentNode >= 0 + && currentNode < size + && layer0Neighbors[currentNode] != null) { + neighborIndex++; + if (neighborIndex < layer0Neighbors[currentNode].size()) { + int neighborNode = layer0Neighbors[currentNode].nodes()[neighborIndex]; + if (neighborNode >= 0 && neighborNode < size) { + return neighborNode; + } else { + return nextNeighbor(); // Skip invalid neighbor + } + } + } else if (currentLevel > 0 && currentLevel < numLevels) { + // Handle higher layers + NeighborArray neighbors = getNeighbors(currentLevel, currentNode); + if (neighbors != null) { + neighborIndex++; + if (neighborIndex < neighbors.size()) { + return neighbors.nodes()[neighborIndex]; + } + } + } + return NO_MORE_DOCS; + } + + @Override + public int entryNode() { + // Entry node should be from the highest layer + if (numLevels > 1) { + int topLevel = numLevels - 1; + int[] topLayerNodes = layerNodes.get(topLevel - 1); + if (topLayerNodes != null && topLayerNodes.length > 0) { + // Use random node from top layer with fixed seed for reproducibility + java.util.Random random = new java.util.Random(44); + int randomIndex = random.nextInt(topLayerNodes.length); + return topLayerNodes[randomIndex]; + } + } + return 0; // Default to node 0 for single-layer graphs + } + + @Override + public int maxConn() { + // Return the maximum degree across all nodes in layer 0 + int max = 0; + for (NeighborArray neighbor : layer0Neighbors) { + if (neighbor != null) { + max = Math.max(max, neighbor.size()); + } + } + return max; + } + + @Override + public int neighborCount() { + if (currentLevel == 0 + && currentNode >= 0 + && currentNode < size + && layer0Neighbors[currentNode] != null) { + return layer0Neighbors[currentNode].size(); + } else if (currentLevel > 0 && currentLevel < numLevels) { + NeighborArray neighbors = getNeighbors(currentLevel, currentNode); + return neighbors != null ? neighbors.size() : 0; + } + return 0; + } + + // Simple implementation of NodesIterator for level 0 + private static class ArrayNodesIterator extends NodesIterator { + private int current = -1; + + ArrayNodesIterator(int size) { + super(size); + } + + @Override + public boolean hasNext() { + return current + 1 < size; + } + + @Override + public int nextInt() { + return ++current; + } + + @Override + public int consume(int[] dest) { + int numToCopy = Math.min(dest.length, size - (current + 1)); + for (int i = 0; i < numToCopy; i++) { + dest[i] = ++current; + } + return numToCopy; + } + } + + // NodesIterator for specific nodes in higher layers + private static class SpecificNodesIterator extends NodesIterator { + private final int[] nodeIds; + private int current = -1; + + SpecificNodesIterator(int[] nodeIds) { + super(nodeIds.length); + this.nodeIds = nodeIds; + } + + @Override + public boolean hasNext() { + return current + 1 < nodeIds.length; + } + + @Override + public int nextInt() { + return nodeIds[++current]; + } + + @Override + public int consume(int[] dest) { + int numToCopy = Math.min(dest.length, nodeIds.length - (current + 1)); + for (int i = 0; i < numToCopy; i++) { + dest[i] = nodeIds[++current]; + } + return numToCopy; + } + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index d3d3b48f..c8d207bb 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -17,6 +17,7 @@ import com.nvidia.cuvs.CuVSMatrix; import java.io.IOException; +import java.time.Duration; import java.util.List; public class Utils { @@ -48,4 +49,8 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions) { } return builder.build(); } + + static long nanosToMillis(long nanos) { + return Duration.ofNanos(nanos).toMillis(); + } } diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec new file mode 100644 index 00000000..7c0af61f --- /dev/null +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -0,0 +1,2 @@ +com.nvidia.cuvs.lucene.CuVSCodec +com.nvidia.cuvs.lucene.CuVSCPUSearchCodec diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java new file mode 100644 index 00000000..06a3a716 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -0,0 +1,169 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +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.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +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; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { + + protected static Logger log = + Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); + private static Random random; + private static Path indexDirPath; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", CuVSVectorsFormat.supported()); + random = new Random(); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @Test + public void testCagraToHnswSerializationAndSearch() throws IOException { + + Codec codec = new CuVSCPUSearchCodec(); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + int numDocs = random.nextInt(100, 1000); + int dimension = random.nextInt(8, 1024); + int topK = random.nextInt(5, 60); + final int COMMIT_FREQ = Math.min(numDocs, random.nextInt(100, 1000)); + int count = COMMIT_FREQ; + final String VECTOR_FIELD = "knn1"; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField("id", Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath)) { + try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.info("Successfully opened index"); + + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues("knn1"); + assertNotNull(knnValues); + log.info( + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + } + assertTrue("Dataset size mismatch", vectorCount == numDocs); + + log.info("\n2. Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.info("Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.info("\nknn1 search results (" + results.totalHits + " total hits):"); + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + log.info( + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + doc.get("id") + + "), score=" + + scoreDoc.score); + } + + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + // TODO: make this test a bit more meaningful like checking the quality of search results. + + } catch (Exception e) { + e.printStackTrace(); + } + } + } + + @AfterClass + public static void afterClass() throws Exception { + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + private static float[][] generateDataset(Random random, int datasetSize, int dimensions) { + float[][] dataset = new float[datasetSize][dimensions]; + for (int i = 0; i < datasetSize; i++) { + for (int j = 0; j < dimensions; j++) { + dataset[i][j] = random.nextFloat() * 100; + } + } + return dataset; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index d5b15616..af2027cc 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -735,6 +735,7 @@ public void testMergeBruteForceIndex() throws IOException { 32, // writer threads 128, // intermediate graph degree 64, // graph degree + 1, IndexType.BRUTE_FORCE); // Use brute force index IndexWriterConfig config = @@ -886,6 +887,7 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { 32, // writer threads 128, // intermediate graph degree 64, // graph degree + 1, IndexType.CAGRA_AND_BRUTE_FORCE); // Use combined CAGRA + brute force IndexWriterConfig config = From c69a568587d479ee88d7fa51c60d82ab58c4298a Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Mon, 18 Aug 2025 16:53:19 -0400 Subject: [PATCH 03/21] Codec, format, reader, and writer separation --- .../nvidia/cuvs/lucene/CuVSSegmentFile.java | 57 -- ...pHnswGraph.java => GPUBuiltHnswGraph.java} | 4 +- ...VSFieldWriter.java => GPUFieldWriter.java} | 12 +- .../lucene/{CuVSIndex.java => GPUIndex.java} | 6 +- ...Query.java => GPUKnnFloatVectorQuery.java} | 9 +- ...r.java => GPUPerLeafCuVSKnnCollector.java} | 4 +- .../{CuVSCodec.java => GPUSearchCodec.java} | 39 +- .../nvidia/cuvs/lucene/GPUVectorsFormat.java | 145 +++++ ...ctorsReader.java => GPUVectorsReader.java} | 32 +- .../nvidia/cuvs/lucene/GPUVectorsWriter.java | 607 ++++++++++++++++++ ...USearchCodec.java => HNSWSearchCodec.java} | 45 +- ...torsFormat.java => HNSWVectorsFormat.java} | 59 +- ...torsWriter.java => HNSWVectorsWriter.java} | 497 ++++---------- .../java/com/nvidia/cuvs/lucene/Utils.java | 24 + .../services/org.apache.lucene.codecs.Codec | 4 +- .../org.apache.lucene.codecs.KnnVectorsFormat | 3 +- ...TestCagraToHnswSerializationAndSearch.java | 21 +- .../cuvs/lucene/TestCuVSDeletedDocuments.java | 6 +- .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 10 +- .../TestCuVSRandomizedVectorSearch.java | 8 +- .../cuvs/lucene/TestCuVSVectorsFormat.java | 4 +- .../com/nvidia/cuvs/lucene/TestMerge.java | 22 +- 22 files changed, 1067 insertions(+), 551 deletions(-) delete mode 100644 src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java rename src/main/java/com/nvidia/cuvs/lucene/{OnHeapHnswGraph.java => GPUBuiltHnswGraph.java} (98%) rename src/main/java/com/nvidia/cuvs/lucene/{CuVSFieldWriter.java => GPUFieldWriter.java} (84%) rename src/main/java/com/nvidia/cuvs/lucene/{CuVSIndex.java => GPUIndex.java} (94%) rename src/main/java/com/nvidia/cuvs/lucene/{CuVSKnnFloatVectorQuery.java => GPUKnnFloatVectorQuery.java} (86%) rename src/main/java/com/nvidia/cuvs/lucene/{PerLeafCuVSKnnCollector.java => GPUPerLeafCuVSKnnCollector.java} (93%) rename src/main/java/com/nvidia/cuvs/lucene/{CuVSCodec.java => GPUSearchCodec.java} (56%) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java rename src/main/java/com/nvidia/cuvs/lucene/{CuVSVectorsReader.java => GPUVectorsReader.java} (94%) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java rename src/main/java/com/nvidia/cuvs/lucene/{CuVSCPUSearchCodec.java => HNSWSearchCodec.java} (60%) rename src/main/java/com/nvidia/cuvs/lucene/{CuVSVectorsFormat.java => HNSWVectorsFormat.java} (70%) rename src/main/java/com/nvidia/cuvs/lucene/{CuVSVectorsWriter.java => HNSWVectorsWriter.java} (63%) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java deleted file mode 100644 index 8a601b7e..00000000 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs.lucene; - -import java.io.IOException; -import java.io.OutputStream; -import java.util.Collections; -import java.util.HashSet; -import java.util.Set; -import java.util.logging.Logger; -import java.util.zip.Deflater; -import java.util.zip.ZipEntry; -import java.util.zip.ZipOutputStream; - -/** Methods to deal with a CuVS composite file inside a segment */ -/*package-private*/ class CuVSSegmentFile implements AutoCloseable { - private final ZipOutputStream zos; - - private Set filesAdded = new HashSet(); - - public CuVSSegmentFile(OutputStream out) { - zos = new ZipOutputStream(out); - zos.setLevel(Deflater.NO_COMPRESSION); - } - - protected Logger log = Logger.getLogger(getClass().getName()); - - public void addFile(String name, byte[] bytes) throws IOException { - ZipEntry indexFileZipEntry = new ZipEntry(name); - zos.putNextEntry(indexFileZipEntry); - zos.write(bytes, 0, bytes.length); - zos.closeEntry(); - filesAdded.add(name); - } - - public Set getFilesAdded() { - return Collections.unmodifiableSet(filesAdded); - } - - @Override - public void close() throws IOException { - zos.close(); - } -} diff --git a/src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java similarity index 98% rename from src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index b204eaf9..e5ecd617 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/OnHeapHnswGraph.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -22,7 +22,7 @@ import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; -public class OnHeapHnswGraph extends HnswGraph { +public class GPUBuiltHnswGraph extends HnswGraph { private final int size; private final int dimensions; @@ -36,7 +36,7 @@ public class OnHeapHnswGraph extends HnswGraph { private final NeighborArray[] layer0Neighbors; // Multi-layer constructor that supports arbitrary number of layers - public OnHeapHnswGraph( + public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies) { this.size = size; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSFieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java similarity index 84% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSFieldWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java index acd15181..483e18d3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSFieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java @@ -24,16 +24,16 @@ import org.apache.lucene.util.RamUsageEstimator; /** CuVS based fields writer */ -/*package-private*/ class CuVSFieldWriter extends KnnFieldVectorsWriter { +/*package-private*/ class GPUFieldWriter extends KnnFieldVectorsWriter { private static final long SHALLOW_SIZE = - RamUsageEstimator.shallowSizeOfInstance(CuVSFieldWriter.class); + RamUsageEstimator.shallowSizeOfInstance(GPUFieldWriter.class); private final FieldInfo fieldInfo; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; - public CuVSFieldWriter( + public GPUFieldWriter( FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { this.fieldInfo = fieldInfo; this.flatFieldVectorsWriter = flatFieldVectorsWriter; @@ -74,6 +74,10 @@ public long ramBytesUsed() { @Override public String toString() { - return "CuVSFieldWriter[field name=" + fieldInfo.name + ", number=" + fieldInfo.number + "]"; + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(field name=").append(fieldInfo.name); + sb.append("number=").append(fieldInfo.number); + sb.append(")"); + return sb.toString(); } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java similarity index 94% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java index e9a20147..651354d6 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java @@ -23,7 +23,7 @@ import java.util.Objects; /** This class holds references to the actual CuVS Index (Cagra, Brute force, etc.) */ -public class CuVSIndex implements Closeable { +public class GPUIndex implements Closeable { private final CagraIndex cagraIndex; private final BruteForceIndex bruteforceIndex; private final HnswIndex hnswIndex; @@ -33,7 +33,7 @@ public class CuVSIndex implements Closeable { private String segmentName; private volatile boolean closed; - public CuVSIndex( + public GPUIndex( String segmentName, String fieldName, CagraIndex cagraIndex, @@ -50,7 +50,7 @@ public CuVSIndex( this.hnswIndex = null; // TODO: } - public CuVSIndex(CagraIndex cagraIndex, BruteForceIndex bruteforceIndex, HnswIndex hnswIndex) { + public GPUIndex(CagraIndex cagraIndex, BruteForceIndex bruteforceIndex, HnswIndex hnswIndex) { this.cagraIndex = cagraIndex; this.bruteforceIndex = bruteforceIndex; this.hnswIndex = hnswIndex; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java similarity index 86% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 8caf30ae..64f70811 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -18,19 +18,20 @@ import java.io.IOException; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.KnnCollector; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.knn.KnnCollectorManager; import org.apache.lucene.util.Bits; -/** Query for CuVS */ -public class CuVSKnnFloatVectorQuery extends KnnFloatVectorQuery { +/** Query on GPU only */ +public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; - public CuVSKnnFloatVectorQuery( + public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { super(field, target, k, filter); this.iTopK = iTopK; @@ -45,7 +46,7 @@ protected TopDocs approximateSearch( KnnCollectorManager knnCollectorManager) throws IOException { - PerLeafCuVSKnnCollector results = new PerLeafCuVSKnnCollector(k, iTopK, searchWidth); + KnnCollector results = new GPUPerLeafCuVSKnnCollector(k, iTopK, searchWidth); LeafReader reader = context.reader(); reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs); diff --git a/src/main/java/com/nvidia/cuvs/lucene/PerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java similarity index 93% rename from src/main/java/com/nvidia/cuvs/lucene/PerLeafCuVSKnnCollector.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 8e005570..ca30d3d3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/PerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -24,7 +24,7 @@ import org.apache.lucene.search.knn.KnnSearchStrategy; /** KnnCollector for CuVS */ -/*package-private*/ class PerLeafCuVSKnnCollector implements KnnCollector { +/*package-private*/ class GPUPerLeafCuVSKnnCollector implements KnnCollector { public List scoreDocs; public int topK = 0; @@ -32,7 +32,7 @@ public int searchWidth = 1; // TODO getter, no setter public int results = 0; - public PerLeafCuVSKnnCollector(int topK, int iTopK, int searchWidth) { + public GPUPerLeafCuVSKnnCollector(int topK, int iTopK, int searchWidth) { super(); this.topK = topK; this.iTopK = iTopK; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java similarity index 56% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java index 8b99ec2a..97ca2621 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java @@ -16,7 +16,7 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.GPUVectorsWriter.IndexType; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.FilterCodec; @@ -24,34 +24,45 @@ import org.apache.lucene.codecs.lucene101.Lucene101Codec; /** CuVS based codec for GPU based vector search */ -public class CuVSCodec extends FilterCodec { +public class GPUSearchCodec extends FilterCodec { - public CuVSCodec() { - this("CuVSCodec", new Lucene101Codec()); + private static final Logger log = Logger.getLogger(GPUSearchCodec.class.getName()); + private static final String CLASS_NAME = "GPUSearchCodec"; + + private static final int DEFAULT_CUVS_WRITER_THREADS = 1; + private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + private static final int DEFAULT_GRAPH_DEGREE = 64; + private static final int DEFAULT_HNSW_LAYERS = 1; + private static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + + private KnnVectorsFormat format; + + public GPUSearchCodec() { + this(CLASS_NAME, new Lucene101Codec()); } - public CuVSCodec(String name, Codec delegate) { + public GPUSearchCodec(String name, Codec delegate) { super(name, delegate); - KnnVectorsFormat format; try { - // TODO: The hard-coded values passed below should be configurable. - // To make relevant changes in a subsequent PR. - format = new CuVSVectorsFormat(1, 128, 64, 1, IndexType.CAGRA); + format = + new GPUVectorsFormat( + DEFAULT_CUVS_WRITER_THREADS, + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + DEFAULT_HNSW_LAYERS, + DEFAULT_INDEX_TYPE); setKnnFormat(format); } catch (LibraryException ex) { - Logger log = Logger.getLogger(CuVSCodec.class.getName()); log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); } } - KnnVectorsFormat knnFormat = null; - @Override public KnnVectorsFormat knnVectorsFormat() { - return knnFormat; + return format; } public void setKnnFormat(KnnVectorsFormat format) { - this.knnFormat = format; + this.format = format; } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java new file mode 100644 index 00000000..2fb058d7 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java @@ -0,0 +1,145 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.LibraryException; +import com.nvidia.cuvs.lucene.GPUVectorsWriter.IndexType; +import java.io.IOException; +import java.util.logging.Logger; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** CuVS based KnnVectorsFormat for GPU acceleration */ +public class GPUVectorsFormat extends KnnVectorsFormat { + + static final Logger log = Logger.getLogger(GPUVectorsFormat.class.getName()); + + // TODO: fix Lucene version in name, to the final targeted release, if any + static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; + static final String CUVS_META_CODEC_EXT = "vemc"; // ""cagmf"; + static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex"; + static final String CUVS_INDEX_EXT = "vcag"; + + static final int VERSION_START = 0; + static final int VERSION_CURRENT = VERSION_START; + + static final int DEFAULT_WRITER_THREADS = 32; + static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + static final int DEFAULT_GRAPH_DEGREE = 64; + static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + static final int HNSW_GRAPH_LAYERS = 1; + + static CuVSResources resources = Utils.cuVSResourcesOrNull(); + + /** The format for storing, reading, and merging raw vectors on disk. */ + private static final FlatVectorsFormat flatVectorsFormat = + new Lucene99FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE); + + final int maxDimensions = 4096; + final int cuvsWriterThreads; + final int intGraphDegree; + final int graphDegree; + final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion + final GPUVectorsWriter.IndexType indexType; // the index type to build, when writing + + /** + * Creates a CuVSVectorsFormat, with default values. + * + * @throws LibraryException if the native library fails to load + */ + public GPUVectorsFormat() { + this( + DEFAULT_WRITER_THREADS, + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + HNSW_GRAPH_LAYERS, + DEFAULT_INDEX_TYPE); + } + + /** + * Creates a CuVSVectorsFormat, with the given threads, graph degree, etc. + * + * @throws LibraryException if the native library fails to load + */ + public GPUVectorsFormat( + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + IndexType indexType) { + super("CuVSVectorsFormat"); + this.cuvsWriterThreads = cuvsWriterThreads; + this.intGraphDegree = intGraphDegree; + this.graphDegree = graphDegree; + this.hnswLayers = hnswLayers; + this.indexType = indexType; + } + + /** Tells whether the platform supports cuvs. */ + public static boolean supported() { + return resources != null; + } + + private static void checkSupported() { + if (!supported()) { + throw new UnsupportedOperationException(); + } + } + + @Override + public GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + checkSupported(); + var flatWriter = flatVectorsFormat.fieldsWriter(state); + return new GPUVectorsWriter( + state, + cuvsWriterThreads, + intGraphDegree, + graphDegree, + hnswLayers, + indexType, + resources, + flatWriter); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + checkSupported(); + return new GPUVectorsReader(state, resources, flatVectorsFormat.fieldsReader(state)); + } + + @Override + public int getMaxDimensions(String fieldName) { + return maxDimensions; + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(cuvsWriterThreads=").append(cuvsWriterThreads); + sb.append("intGraphDegree=").append(intGraphDegree); + sb.append("graphDegree=").append(graphDegree); + sb.append("hnswLayers=").append(hnswLayers); + sb.append("resources=").append(resources); + sb.append(")"); + return sb.toString(); + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsReader.java similarity index 94% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUVectorsReader.java index 4118a0aa..541044cd 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsReader.java @@ -15,12 +15,12 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_CURRENT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_START; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.VERSION_START; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.BruteForceQuery; @@ -62,19 +62,19 @@ import org.apache.lucene.util.hnsw.IntToIntFunction; /** KnnVectorsReader instance associated with CuVS format */ -public class CuVSVectorsReader extends KnnVectorsReader { +public class GPUVectorsReader extends KnnVectorsReader { @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(CuVSVectorsReader.class.getName()); + private static final Logger log = Logger.getLogger(GPUVectorsReader.class.getName()); private final CuVSResources resources; private final FlatVectorsReader flatVectorsReader; // for reading the raw vectors private final FieldInfos fieldInfos; private final IntObjectHashMap fields; - private final IntObjectHashMap cuvsIndices; + private final IntObjectHashMap cuvsIndices; private final IndexInput cuvsIndexInput; - public CuVSVectorsReader( + public GPUVectorsReader( SegmentReadState state, CuVSResources resources, FlatVectorsReader flatReader) throws IOException { this.resources = resources; @@ -226,8 +226,8 @@ private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) return fieldEntry; } - private IntObjectHashMap loadCuVSIndices() throws IOException { - var indices = new IntObjectHashMap(); + private IntObjectHashMap loadCuVSIndices() throws IOException { + var indices = new IntObjectHashMap(); for (var e : fields) { var fieldEntry = e.value; int fieldNumber = e.key; @@ -237,7 +237,7 @@ private IntObjectHashMap loadCuVSIndices() throws IOException { return indices; } - private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { + private GPUIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { CagraIndex cagraIndex = null; BruteForceIndex bruteForceIndex = null; HnswIndex hnswIndex = null; @@ -273,7 +273,7 @@ private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { } catch (Throwable t) { Utils.handleThrowable(t); } - return new CuVSIndex(cagraIndex, bruteForceIndex, hnswIndex); + return new GPUIndex(cagraIndex, bruteForceIndex, hnswIndex); } @Override @@ -335,7 +335,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits var fieldNumber = fieldInfos.fieldInfo(field).number; - CuVSIndex cuvsIndex = cuvsIndices.get(fieldNumber); + GPUIndex cuvsIndex = cuvsIndices.get(fieldNumber); if (cuvsIndex == null) { throw new IllegalStateException("not index found for field:" + field); } @@ -476,7 +476,7 @@ public FieldInfos getFieldInfos() { return fieldInfos; } - public IntObjectHashMap getCuvsIndexes() { + public IntObjectHashMap getCuvsIndexes() { return cuvsIndices; } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java new file mode 100644 index 00000000..5bad8534 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java @@ -0,0 +1,607 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.GPUVectorsFormat.VERSION_CURRENT; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +import com.nvidia.cuvs.BruteForceIndex; +import com.nvidia.cuvs.BruteForceIndexParams; +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.logging.Logger; +import java.util.stream.IntStream; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +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.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntObjectHashMap; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; + +/** + * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into + * GPU + */ +public class GPUVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(GPUVectorsWriter.class); + + @SuppressWarnings("unused") + private static final Logger log = Logger.getLogger(GPUVectorsWriter.class.getName()); + + /** The name of the CUVS component for the info-stream * */ + private static final String CUVS_COMPONENT = "CUVS"; + + // The minimum number of vectors in the dataset required before + // we attempt to build a Cagra index + static final int MIN_CAGRA_INDEX_SIZE = 2; + + private final int cuvsWriterThreads; + private final int intGraphDegree; + private final int graphDegree; + + private final CuVSResources resources; + private final IndexType indexType; + + private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors + private final List fields = new ArrayList<>(); + private IndexOutput meta = null, cuvsIndex = null; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; + private final InfoStream infoStream; + private boolean finished; + + /** The CuVS index Type. */ + public enum IndexType { + + /** Builds a Cagra index. */ + CAGRA(true, false, false), + + /** Builds a Brute Force index. */ + BRUTE_FORCE(false, true, false), + + /** Builds an HSNW index - suitable for searching on CPU. */ + HNSW(false, false, true), + + /** Builds a Cagra and a Brute Force index. */ + CAGRA_AND_BRUTE_FORCE(true, true, false); + private final boolean cagra, bruteForce, hnsw; + + IndexType(boolean cagra, boolean bruteForce, boolean hnsw) { + this.cagra = cagra; + this.bruteForce = bruteForce; + this.hnsw = hnsw; + } + + public boolean cagra() { + return cagra; + } + + public boolean bruteForce() { + return bruteForce; + } + + public boolean hnsw() { + return hnsw; + } + } + + public GPUVectorsWriter( + SegmentWriteState state, + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + IndexType indexType, + CuVSResources resources, + FlatVectorsWriter flatVectorsWriter) + throws IOException { + super(); + this.indexType = indexType; + this.cuvsWriterThreads = cuvsWriterThreads; + this.intGraphDegree = intGraphDegree; + this.graphDegree = graphDegree; + this.resources = resources; + this.flatVectorsWriter = flatVectorsWriter; + this.infoStream = state.infoStream; + + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, CUVS_META_CODEC_EXT); + String cagraFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, CUVS_INDEX_EXT); + + boolean success = false; + try { + + meta = state.directory.createOutput(metaFileName, state.context); + cuvsIndex = state.directory.createOutput(cagraFileName, state.context); + CodecUtil.writeIndexHeader( + meta, + CUVS_META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + cuvsIndex, + CUVS_INDEX_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + var encoding = fieldInfo.getVectorEncoding(); + if (encoding != FLOAT32) { + throw new IllegalArgumentException("expected float32, got:" + encoding); + } + var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); + @SuppressWarnings("unchecked") + var flatWriter = (FlatFieldVectorsWriter) writer; + var cuvsFieldWriter = new GPUFieldWriter(fieldInfo, flatWriter); + fields.add(cuvsFieldWriter); + return writer; + } + + static String indexMsg(int size, int... args) { + StringBuilder sb = new StringBuilder("cagra index params"); + sb.append(": size=").append(size); + sb.append(", intGraphDegree=").append(args[0]); + sb.append(", actualIntGraphDegree=").append(args[1]); + sb.append(", graphDegree=").append(args[2]); + sb.append(", actualGraphDegree=").append(args[3]); + return sb.toString(); + } + + private CagraIndexParams cagraIndexParams(int size) { + if (size < 2) { + // https://github.com/rapidsai/cuvs/issues/666 + throw new IllegalArgumentException("cagra index must be greater than 2"); + } + + return new CagraIndexParams.Builder() + .withNumWriterThreads(cuvsWriterThreads) + .withIntermediateGraphDegree(intGraphDegree) + .withGraphDegree(graphDegree) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .build(); + } + + private void info(String msg) { + if (infoStream.isEnabled(CUVS_COMPONENT)) { + infoStream.message(CUVS_COMPONENT, msg); + } + } + + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo); + return; + } + long cagraIndexOffset, cagraIndexLength = 0L; + long bruteForceIndexOffset, bruteForceIndexLength = 0L; + long hnswIndexOffset, hnswIndexLength = 0L; + + // workaround for the minimum number of vectors for Cagra + IndexType indexType = + this.indexType.cagra() && vectors.size() < MIN_CAGRA_INDEX_SIZE + ? IndexType.BRUTE_FORCE + : this.indexType; + + try { + + cagraIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.cagra()) { + try { + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeCagraIndex(cagraIndexOutputStream, dataset); + } catch (Throwable t) { + Utils.handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + // workaround for cuVS issue + indexType = IndexType.BRUTE_FORCE; + } + cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + } + + bruteForceIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.bruteForce()) { + var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeBruteForceIndex(bruteForceIndexOutputStream, dataset); + bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; + } + + hnswIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.hnsw()) { + var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + if (vectors.size() > MIN_CAGRA_INDEX_SIZE) { + try { + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeHNSWIndex(hnswIndexOutputStream, dataset); + } catch (Throwable t) { + Utils.handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + } + } + hnswIndexLength = cuvsIndex.getFilePointer() - hnswIndexOffset; + } + + // Only write meta for non-HNSW_LUCENE modes + writeMeta( + fieldInfo, + vectors.size(), + cagraIndexOffset, + cagraIndexLength, + bruteForceIndexOffset, + bruteForceIndexLength, + hnswIndexOffset, + hnswIndexLength); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + if (dataset.size() < 2) { + throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); + } + CagraIndexParams params = cagraIndexParams((int) dataset.size()); + long startTime = System.nanoTime(); + CagraIndex index = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); + info("Cagra index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); + Path tmpFile = Files.createTempFile(resources.tempDirectory(), "tmpindex", "cag"); + index.serialize(os, tmpFile); + index.destroyIndex(); + } + + private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + BruteForceIndexParams params = + new BruteForceIndexParams.Builder() + .withNumWriterThreads(32) // TODO: Make this configurable. + .build(); + long startTime = System.nanoTime(); + var index = + BruteForceIndex.newBuilder(resources).withIndexParams(params).withDataset(dataset).build(); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); + info("bf index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); + index.serialize(os); + index.destroyIndex(); + } + + private void writeHNSWIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + if (dataset.size() < 2) { + throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); + } + CagraIndexParams indexParams = cagraIndexParams((int) dataset.size()); + long startTime = System.nanoTime(); + CagraIndex index = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(indexParams).build(); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); + info("HNSW index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); + Path tmpFile = Files.createTempFile("tmpindex", "hnsw"); + index.serializeToHNSW(os, tmpFile); + index.destroyIndex(); + } + + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + flatVectorsWriter.flush(maxDoc, sortMap); + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } + } + + private void writeField(GPUFieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); + } + + private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) + throws IOException { + + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord + mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); + + List sortedVectors = new ArrayList(); + for (int i = 0; i < fieldData.getVectors().size(); i++) { + sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); + } + + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); + } + + private void writeEmpty(FieldInfo fieldInfo) throws IOException { + writeMeta(fieldInfo, 0, 0L, 0L, 0L, 0L, 0L, 0L); + } + + private void writeMeta( + FieldInfo field, + int count, + long cagraIndexOffset, + long cagraIndexLength, + long bruteForceIndexOffset, + long bruteForceIndexLength, + long hnswIndexOffset, + long hnswIndexLength) + throws IOException { + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); + meta.writeInt(field.getVectorDimension()); + meta.writeInt(count); + meta.writeVLong(cagraIndexOffset); + meta.writeVLong(cagraIndexLength); + meta.writeVLong(bruteForceIndexOffset); + meta.writeVLong(bruteForceIndexLength); + meta.writeVLong(hnswIndexOffset); + meta.writeVLong(hnswIndexLength); + } + + static int distFuncToOrd(VectorSimilarityFunction func) { + for (int i = 0; i < SIMILARITY_FUNCTIONS.size(); i++) { + if (SIMILARITY_FUNCTIONS.get(i).equals(func)) { + return (byte) i; + } + } + throw new IllegalArgumentException("invalid distance function: " + func); + } + + // We currently ignore this, until cuVS supports tiered indices + private static final String CANNOT_GENERATE_CAGRA = + """ + Could not generate an intermediate CAGRA graph because the initial \ + kNN graph contains too many invalid or duplicated neighbor nodes. \ + This error can occur, for example, if too many overflows occur \ + during the norm computation between the dataset vectors\ + """; + + private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + + List cagraIndexes = new ArrayList<>(); + // We need this count so that the merged segment's meta information has the vector count. + int totalVectorCount = 0; + + for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { + KnnVectorsReader knnReader = mergeState.knnVectorsReaders[i]; + // Access the CAGRA index for this field from the reader + + if (knnReader != null) { + if (knnReader instanceof GPUVectorsReader cvr) { + if (cvr != null) { + totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); + CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); + if (cagraIndex != null) { + cagraIndexes.add(cagraIndex); + } + } + } else { + // This should never happen + throw new RuntimeException( + "Reader is not of CuVSVectorsReader type. Instead it is: " + knnReader.getClass()); + } + } + } + assert cagraIndexes.size() > 1; + + CagraIndex mergedIndex = + CagraIndex.merge(cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()])); + writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + info("Successfully merged " + cagraIndexes.size() + " CAGRA indexes using native merge API"); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Creates List from merged vectors + * */ + private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List res = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + int ordinal = iter.index(); + float[] vector = mergedVectorValues.vectorValue(ordinal); + res.add(vector); + } + return res; + } + + /** + * Fallback method that rebuilds indexes from merged vectors. + * Used when native CAGRA merge() is not possible. Also used + * when non-CAGRA index types are used (for e.g. Brute Force index). + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + if (fieldInfo.getVectorEncoding() != FLOAT32) { + throw new AssertionError("Only Float32 supported"); + } + try { + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. + */ + private CagraIndex getCagraIndexFromReader(GPUVectorsReader reader, String fieldName) { + try { + IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); + FieldInfos fieldInfos = reader.getFieldInfos(); + + FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); + + if (fieldInfo != null) { + GPUIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); + if (cuvsIndex != null) { + return cuvsIndex.getCagraIndex(); + } + } + } catch (Exception e) { + e.printStackTrace(); + info("Failed to extract CAGRA index for field " + fieldName + ": " + e.getMessage()); + } + return null; + } + + /** + * Writes a pre-built merged CAGRA index to the output. + */ + private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, int vectorCount) + throws IOException { + try { + long cagraIndexOffset = cuvsIndex.getFilePointer(); + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + + // Serialize the merged index + Path tmpFile = Files.createTempFile(resources.tempDirectory(), "mergedindex", "cag"); + mergedIndex.serialize(cagraIndexOutputStream, tmpFile); + long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + + writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); + + // Clean up the merged index + mergedIndex.destroyIndex(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + flatVectorsWriter.mergeOneField(fieldInfo, mergeState); + + if (indexType.cagra() && !indexType.bruteForce()) { + // Since CAGRA merge does not support merging of indexes with purging of deletes, + // we fallback to vector-based re-indexing. Issue: + // https://github.com/rapidsai/cuvs/issues/1253 + boolean hasDeletions = + IntStream.range(0, mergeState.liveDocs.length) + .anyMatch( + i -> + mergeState.liveDocs[i] == null + || IntStream.range(0, mergeState.maxDocs[i]) + .anyMatch(j -> !mergeState.liveDocs[i].get(j))); + + if (mergeState.knnVectorsReaders.length > 1 && !hasDeletions) { + mergeCagraIndexes(fieldInfo, mergeState); + } else { + // CAGRA's merge API does not handle the trivial case of merging 1 index. + vectorBasedMerge(fieldInfo, mergeState); + } + + } else { + // If there is a Brute Force index then re-index using the vectors even if there is a CAGRA + // index. + vectorBasedMerge(fieldInfo, mergeState); + } + } + + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + flatVectorsWriter.finish(); + + if (meta != null) { + // write end of fields marker + meta.writeInt(-1); + CodecUtil.writeFooter(meta); + } + if (cuvsIndex != null) { + CodecUtil.writeFooter(cuvsIndex); + } + + { + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } + } + + @Override + public void close() throws IOException { + IOUtils.close(meta, cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); + } + + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (var field : fields) { + total += field.ramBytesUsed(); + } + return total; + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java similarity index 60% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java index f2a679f5..df94f24c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSCPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java @@ -16,7 +16,6 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.FilterCodec; @@ -24,49 +23,57 @@ import org.apache.lucene.codecs.lucene101.Lucene101Codec; /** CuVS based codec for GPU based vector search */ -public class CuVSCPUSearchCodec extends FilterCodec { +public class HNSWSearchCodec extends FilterCodec { - public CuVSCPUSearchCodec() { - this("CuVSCPUSearchCodec", new Lucene101Codec()); + private static final Logger log = Logger.getLogger(HNSWSearchCodec.class.getName()); + + private static final int DEFAULT_CUVS_WRITER_THREADS = 1; + private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + private static final int DEFAULT_GRAPH_DEGREE = 64; + private static final int DEFAULT_HNSW_LAYERS = 1; + private static final String CLASS_NAME = "HNSWSearchCodec"; + + private KnnVectorsFormat format; + + public HNSWSearchCodec() { + this(CLASS_NAME, new Lucene101Codec()); } - public CuVSCPUSearchCodec(String name, Codec delegate) { + public HNSWSearchCodec(String name, Codec delegate) { super(name, delegate); - initializeFormat(); + initializeFormatDefaultValues(); } - public CuVSCPUSearchCodec( + public HNSWSearchCodec( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { - this("CuVSCPUSearchCodec", new Lucene101Codec()); + this(CLASS_NAME, new Lucene101Codec()); initializeFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); } - private void initializeFormat() { - initializeFormat(1, 128, 64, 1); // Default values + private void initializeFormatDefaultValues() { + initializeFormat( + DEFAULT_CUVS_WRITER_THREADS, + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + DEFAULT_HNSW_LAYERS); } private void initializeFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { - KnnVectorsFormat format; try { - format = - new CuVSVectorsFormat( - cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, IndexType.HNSW_LUCENE); + format = new HNSWVectorsFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); setKnnFormat(format); } catch (LibraryException ex) { - Logger log = Logger.getLogger(CuVSCodec.class.getName()); log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); } } - KnnVectorsFormat knnFormat = null; - @Override public KnnVectorsFormat knnVectorsFormat() { - return knnFormat; + return format; } public void setKnnFormat(KnnVectorsFormat format) { - this.knnFormat = format; + this.format = format; } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java similarity index 70% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java rename to src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java index 6048e717..4cf91971 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java @@ -17,7 +17,6 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; import java.io.IOException; import java.util.logging.Logger; import org.apache.lucene.codecs.KnnVectorsFormat; @@ -30,24 +29,23 @@ import org.apache.lucene.index.SegmentWriteState; /** CuVS based KnnVectorsFormat for GPU acceleration */ -public class CuVSVectorsFormat extends KnnVectorsFormat { +public class HNSWVectorsFormat extends KnnVectorsFormat { - private static final Logger log = Logger.getLogger(CuVSVectorsFormat.class.getName()); + private static final Logger log = Logger.getLogger(HNSWVectorsFormat.class.getName()); // TODO: fix Lucene version in name, to the final targeted release, if any static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; - static final String CUVS_META_CODEC_EXT = "vemc"; // ""cagmf"; + static final String CUVS_META_CODEC_EXT = "vemc"; static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex"; static final String CUVS_INDEX_EXT = "vcag"; static final int VERSION_START = 0; static final int VERSION_CURRENT = VERSION_START; - public static final int DEFAULT_WRITER_THREADS = 32; - public static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; - public static final int DEFAULT_GRAPH_DEGREE = 64; - public static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; - public static final int HNSW_GRAPH_LAYERS = 1; + static final int DEFAULT_WRITER_THREADS = 32; + static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + static final int DEFAULT_GRAPH_DEGREE = 64; + static final int HNSW_GRAPH_LAYERS = 1; static CuVSResources resources = cuVSResourcesOrNull(); @@ -60,20 +58,18 @@ public class CuVSVectorsFormat extends KnnVectorsFormat { final int intGraphDegree; final int graphDegree; final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion - final CuVSVectorsWriter.IndexType indexType; // the index type to build, when writing /** * Creates a CuVSVectorsFormat, with default values. * * @throws LibraryException if the native library fails to load */ - public CuVSVectorsFormat() { + public HNSWVectorsFormat() { this( DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, - HNSW_GRAPH_LAYERS, - DEFAULT_INDEX_TYPE); + HNSW_GRAPH_LAYERS); } /** @@ -81,18 +77,13 @@ public CuVSVectorsFormat() { * * @throws LibraryException if the native library fails to load */ - public CuVSVectorsFormat( - int cuvsWriterThreads, - int intGraphDegree, - int graphDegree, - int hnswLayers, - IndexType indexType) { + public HNSWVectorsFormat( + int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { super("CuVSVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; this.hnswLayers = hnswLayers; - this.indexType = indexType; } private static CuVSResources cuVSResourcesOrNull() { @@ -122,32 +113,16 @@ private static void checkSupported() { } @Override - public CuVSVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + public HNSWVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { checkSupported(); var flatWriter = flatVectorsFormat.fieldsWriter(state); - return new CuVSVectorsWriter( - state, - cuvsWriterThreads, - intGraphDegree, - graphDegree, - hnswLayers, - indexType, - resources, - flatWriter); + return new HNSWVectorsWriter( + state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); } @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - checkSupported(); - var flatReader = flatVectorsFormat.fieldsReader(state); - if (this.indexType == IndexType.HNSW_LUCENE) { - log.info("Using Reader: Lucene99HnswVectorsReader"); - return new Lucene99HnswVectorsReader(state, flatReader); - } else { - checkSupported(); - log.info("Using Reader: CuVSVectorsReader"); - return new CuVSVectorsReader(state, resources, flatReader); - } + return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state)); } @Override @@ -157,8 +132,8 @@ public int getMaxDimensions(String fieldName) { @Override public String toString() { - StringBuilder sb = new StringBuilder("CuVSVectorsFormat("); - sb.append("cuvsWriterThreads=").append(cuvsWriterThreads); + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(cuvsWriterThreads=").append(cuvsWriterThreads); sb.append("intGraphDegree=").append(intGraphDegree); sb.append("graphDegree=").append(graphDegree); sb.append("hnswLayers=").append(hnswLayers); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java similarity index 63% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java index 47627c4f..160bf3f5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java @@ -15,25 +15,17 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_CURRENT; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; -import com.nvidia.cuvs.BruteForceIndex; -import com.nvidia.cuvs.BruteForceIndexParams; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; import java.io.IOException; -import java.io.OutputStream; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; @@ -73,12 +65,12 @@ * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into * GPU */ -public class CuVSVectorsWriter extends KnnVectorsWriter { +public class HNSWVectorsWriter extends KnnVectorsWriter { - private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(CuVSVectorsWriter.class); + private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(HNSWVectorsWriter.class); @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(CuVSVectorsWriter.class.getName()); + private static final Logger log = Logger.getLogger(HNSWVectorsWriter.class.getName()); /** The name of the CUVS component for the info-stream * */ public static final String CUVS_COMPONENT = "CUVS"; @@ -93,10 +85,9 @@ public class CuVSVectorsWriter extends KnnVectorsWriter { private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion private final CuVSResources resources; - private final IndexType indexType; private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors - private final List fields = new ArrayList<>(); + private final List fields = new ArrayList<>(); private IndexOutput meta = null, cuvsIndex = null; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private final InfoStream infoStream; @@ -104,56 +95,16 @@ public class CuVSVectorsWriter extends KnnVectorsWriter { private String vemFileName; private String vexFileName; - /** The CuVS index Type. */ - public enum IndexType { - /** Builds a Cagra index. */ - CAGRA(true, false, false, false), - /** Builds a Brute Force index. */ - BRUTE_FORCE(false, true, false, false), - /** Builds an HSNW index - suitable for searching on CPU. */ - HNSW(false, false, true, false), - /** Builds a Cagra and a Brute Force index. */ - CAGRA_AND_BRUTE_FORCE(true, true, false, false), - /** Builds a Lucene HNSW index via CAGRA. */ - HNSW_LUCENE(false, false, false, true); - private final boolean cagra, bruteForce, hnsw, hnswLucene; - - IndexType(boolean cagra, boolean bruteForce, boolean hnsw, boolean hnswLucene) { - this.cagra = cagra; - this.bruteForce = bruteForce; - this.hnsw = hnsw; - this.hnswLucene = hnswLucene; - } - - public boolean cagra() { - return cagra; - } - - public boolean bruteForce() { - return bruteForce; - } - - public boolean hnsw() { - return hnsw; - } - - public boolean hnswLucene() { - return hnswLucene; - } - } - - public CuVSVectorsWriter( + public HNSWVectorsWriter( SegmentWriteState state, int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers, - IndexType indexType, CuVSResources resources, FlatVectorsWriter flatVectorsWriter) throws IOException { super(); - this.indexType = indexType; this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; @@ -168,50 +119,25 @@ public CuVSVectorsWriter( vexFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "vex"); - String metaFileName = - IndexFileNames.segmentFileName( - state.segmentInfo.name, state.segmentSuffix, CUVS_META_CODEC_EXT); - String cagraFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, CUVS_INDEX_EXT); - boolean success = false; try { - // Only create CAGRA files if not in HNSW_LUCENE mode - if (indexType == IndexType.HNSW_LUCENE) { - - hnswMeta = state.directory.createOutput(vemFileName, state.context); - hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); - - CodecUtil.writeIndexHeader( - hnswMeta, - "Lucene99HnswVectorsFormatMeta", - Lucene99HnswVectorsFormat.VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - CodecUtil.writeIndexHeader( - hnswVectorIndex, - "Lucene99HnswVectorsFormatIndex", - Lucene99HnswVectorsFormat.VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - } else { - // Only create CAGRA files if not in HNSW_LUCENE mode - meta = state.directory.createOutput(metaFileName, state.context); - cuvsIndex = state.directory.createOutput(cagraFileName, state.context); - CodecUtil.writeIndexHeader( - meta, - CUVS_META_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - CodecUtil.writeIndexHeader( - cuvsIndex, - CUVS_INDEX_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - } + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + + CodecUtil.writeIndexHeader( + hnswMeta, + "Lucene99HnswVectorsFormatMeta", + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + "Lucene99HnswVectorsFormatIndex", + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; } finally { if (success == false) { @@ -229,7 +155,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); @SuppressWarnings("unchecked") var flatWriter = (FlatFieldVectorsWriter) writer; - var cuvsFieldWriter = new CuVSFieldWriter(fieldInfo, flatWriter); + var cuvsFieldWriter = new GPUFieldWriter(fieldInfo, flatWriter); fields.add(cuvsFieldWriter); return writer; } @@ -269,191 +195,117 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeEmpty(fieldInfo); return; } - long cagraIndexOffset, cagraIndexLength = 0L; - long bruteForceIndexOffset, bruteForceIndexLength = 0L; - long hnswIndexOffset, hnswIndexLength = 0L; - - // workaround for the minimum number of vectors for Cagra - IndexType indexType = - this.indexType.cagra() && vectors.size() < MIN_CAGRA_INDEX_SIZE - ? IndexType.BRUTE_FORCE - : this.indexType; - - info("=== INDEX TYPE DEBUG: original=" + this.indexType + ", effective=" + indexType + " ==="); try { - if (indexType.hnswLucene()) { - info("=== ENTERED HNSW_LUCENE BLOCK (HNSW-only mode) ==="); - info("Entered the writeFieldInternal's HNSW LUCENE block - writing only HNSW files"); - try { - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeHnswOnlyIndex(dataset, fieldInfo, vectors); - } catch (Throwable t) { - info("=== ERROR IN HNSW_LUCENE: " + t.getMessage() + " ==="); - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - // workaround for cuVS issue - indexType = IndexType.BRUTE_FORCE; - } - // For HNSW_LUCENE, we don't write any CAGRA data, so set lengths to 0 - cagraIndexLength = 0L; - cagraIndexOffset = 0L; - bruteForceIndexOffset = 0L; - bruteForceIndexLength = 0L; - hnswIndexOffset = 0L; - hnswIndexLength = 0L; - } else { - cagraIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.cagra()) { - try { - var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeCagraIndex(cagraIndexOutputStream, dataset); - } catch (Throwable t) { - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - // workaround for cuVS issue - indexType = IndexType.BRUTE_FORCE; - } - cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; - } + info("=== ENTERED HNSW_LUCENE BLOCK (HNSW-only mode) ==="); + info("Entered the writeFieldInternal's HNSW LUCENE block - writing only HNSW files"); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - bruteForceIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.bruteForce()) { - var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeBruteForceIndex(bruteForceIndexOutputStream, dataset); - bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; - } + if (dataset.size() < 2) { + throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); + } + CagraIndexParams params = cagraIndexParams((int) dataset.size()); + long startTime = System.nanoTime(); + CagraIndex index = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); - hnswIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.hnsw()) { - var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - if (vectors.size() > MIN_CAGRA_INDEX_SIZE) { - try { - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeHNSWIndex(hnswIndexOutputStream, dataset); - } catch (Throwable t) { - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - } + // Get the adjacency list from CAGRA index + int[][] adjacencyList; + try { + adjacencyList = index.getGraph(); + info("=== SUCCESS: Got adjacency list from CAGRA index ==="); + info("Successfully got adjacency list from CAGRA index"); + } catch (Exception e) { + info("=== FAILED: getGraph() method failed: " + e.getMessage() + " ==="); + info("getGraph() method failed or doesn't exist: " + e.getMessage()); + // Create a mock adjacency list for testing + int size = (int) dataset.size(); + adjacencyList = new int[size][]; + for (int i = 0; i < size; i++) { + // Create connections to next few nodes (circular) + int degree = Math.min(10, size - 1); // up to 10 connections + adjacencyList[i] = new int[degree]; + for (int j = 0; j < degree; j++) { + adjacencyList[i][j] = (i + j + 1) % size; } - hnswIndexLength = cuvsIndex.getFilePointer() - hnswIndexOffset; } + info( + "=== CREATED MOCK ADJACENCY LIST: " + + size + + " nodes, degree=" + + (adjacencyList.length > 0 ? adjacencyList[0].length : 0) + + " ==="); + info( + "Created mock adjacency list with " + + size + + " nodes, degree=" + + (adjacencyList.length > 0 ? adjacencyList[0].length : 0)); } - // Only write meta for non-HNSW_LUCENE modes - if (indexType != IndexType.HNSW_LUCENE) { - writeMeta( - fieldInfo, - vectors.size(), - cagraIndexOffset, - cagraIndexLength, - bruteForceIndexOffset, - bruteForceIndexLength, - hnswIndexOffset, - hnswIndexLength); - } - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - - private void writeHnswOnlyIndex( - CuVSMatrix dataset, FieldInfo fieldInfo, List originalVectors) throws Throwable { - if (dataset.size() < 2) { - throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); - } - CagraIndexParams params = cagraIndexParams((int) dataset.size()); - long startTime = System.nanoTime(); - CagraIndex index = - CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); - - // Get the adjacency list from CAGRA index - int[][] adjacencyList; - try { - adjacencyList = index.getGraph(); - info("=== SUCCESS: Got adjacency list from CAGRA index ==="); - info("Successfully got adjacency list from CAGRA index"); - } catch (Exception e) { - info("=== FAILED: getGraph() method failed: " + e.getMessage() + " ==="); - info("getGraph() method failed or doesn't exist: " + e.getMessage()); - // Create a mock adjacency list for testing int size = (int) dataset.size(); - adjacencyList = new int[size][]; - for (int i = 0; i < size; i++) { - // Create connections to next few nodes (circular) - int degree = Math.min(10, size - 1); // up to 10 connections - adjacencyList[i] = new int[degree]; - for (int j = 0; j < degree; j++) { - adjacencyList[i][j] = (i + j + 1) % size; - } - } - info( - "=== CREATED MOCK ADJACENCY LIST: " - + size - + " nodes, degree=" - + (adjacencyList.length > 0 ? adjacencyList[0].length : 0) - + " ==="); + int dimensions = fieldInfo.getVectorDimension(); + + // Debug: Check if we got valid adjacency data info( - "Created mock adjacency list with " - + size - + " nodes, degree=" - + (adjacencyList.length > 0 ? adjacencyList[0].length : 0)); - } + "Adjacency list info: " + + (adjacencyList == null + ? "null" + : "length=" + + adjacencyList.length + + ", first row=" + + (adjacencyList.length > 0 && adjacencyList[0] != null + ? adjacencyList[0].length + : "null"))); + + // Create HNSW graph from CAGRA - multi-layer if original vectors available + GPUBuiltHnswGraph hnswGraph; + if (vectors != null && vectors.size() > 0) { + info("=== Creating 3-layer HNSW graph using original vectors ==="); + hnswGraph = createMultiLayerHnswGraph(fieldInfo, size, dimensions, adjacencyList, vectors); + } else { + info("=== Creating single-layer HNSW graph (no original vectors) ==="); + // Create single layer graph + List singleLayerNodes = new ArrayList<>(); + List singleLayerAdjacencies = new ArrayList<>(); + singleLayerAdjacencies.add(adjacencyList); + hnswGraph = + new GPUBuiltHnswGraph(size, dimensions, singleLayerNodes, singleLayerAdjacencies); + } - int size = (int) dataset.size(); - int dimensions = fieldInfo.getVectorDimension(); + // Remember the vector index offset before writing + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Debug: Check if we got valid adjacency data - info( - "Adjacency list info: " - + (adjacencyList == null - ? "null" - : "length=" - + adjacencyList.length - + ", first row=" - + (adjacencyList.length > 0 && adjacencyList[0] != null - ? adjacencyList[0].length - : "null"))); - - // Create HNSW graph from CAGRA - multi-layer if original vectors available - OnHeapHnswGraph hnswGraph; - if (originalVectors != null && originalVectors.size() > 0) { - info("=== Creating 3-layer HNSW graph using original vectors ==="); - hnswGraph = - createMultiLayerHnswGraph(fieldInfo, size, dimensions, adjacencyList, originalVectors); - } else { - info("=== Creating single-layer HNSW graph (no original vectors) ==="); - // Create single layer graph - List singleLayerNodes = new ArrayList<>(); - List singleLayerAdjacencies = new ArrayList<>(); - singleLayerAdjacencies.add(adjacencyList); - hnswGraph = new OnHeapHnswGraph(size, dimensions, singleLayerNodes, singleLayerAdjacencies); - } + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - // Remember the vector index offset before writing - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + // Calculate the length of written data + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets); - // Calculate the length of written data - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); + info( + "HNSW-only graph created in " + + elapsedMillis + + "ms, with " + + dataset.size() + + " vectors"); - long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); - info("HNSW-only graph created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); + // Don't serialize CAGRA index - destroy it immediately + index.destroyIndex(); - // Don't serialize CAGRA index - destroy it immediately - index.destroyIndex(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } } /** @@ -462,7 +314,7 @@ private void writeHnswOnlyIndex( * Each layer contains 1/M nodes from the previous layer * Creates layers until the highest layer has ≤ M nodes */ - private OnHeapHnswGraph createMultiLayerHnswGraph( + private GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int size, int dimensions, @@ -567,7 +419,7 @@ private OnHeapHnswGraph createMultiLayerHnswGraph( info("Created " + numLayers + " layers total"); // Create the multi-layer graph with all layers - return new OnHeapHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } /** @@ -745,7 +597,7 @@ private void writeMeta( } } - private int[][] writeGraph(OnHeapHnswGraph graph, IndexOutput vectorIndex) throws IOException { + private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { if (graph == null) return new int[0][0]; // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); @@ -797,51 +649,6 @@ private int[][] writeGraph(OnHeapHnswGraph graph, IndexOutput vectorIndex) throw return offsets; } - private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { - if (dataset.size() < 2) { - throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); - } - CagraIndexParams params = cagraIndexParams((int) dataset.size()); - long startTime = System.nanoTime(); - CagraIndex index = - CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); - long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); - info("Cagra index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); - Path tmpFile = Files.createTempFile(resources.tempDirectory(), "tmpindex", "cag"); - index.serialize(os, tmpFile); - index.destroyIndex(); - } - - private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { - BruteForceIndexParams params = - new BruteForceIndexParams.Builder() - .withNumWriterThreads(32) // TODO: Make this - // configurable later. - .build(); - long startTime = System.nanoTime(); - var index = - BruteForceIndex.newBuilder(resources).withIndexParams(params).withDataset(dataset).build(); - long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); - info("bf index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); - index.serialize(os); - index.destroyIndex(); - } - - private void writeHNSWIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { - if (dataset.size() < 2) { - throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); - } - CagraIndexParams indexParams = cagraIndexParams((int) dataset.size()); - long startTime = System.nanoTime(); - CagraIndex index = - CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(indexParams).build(); - long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); - info("HNSW index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); - Path tmpFile = Files.createTempFile("tmpindex", "hnsw"); - index.serializeToHNSW(os, tmpFile); - index.destroyIndex(); - } - @Override public void flush(int maxDoc, DocMap sortMap) throws IOException { flatVectorsWriter.flush(maxDoc, sortMap); @@ -854,18 +661,16 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { } } - private void writeField(CuVSFieldWriter fieldData) throws IOException { + private void writeField(GPUFieldWriter fieldData) throws IOException { writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); } - private void writeSortingField(CuVSFieldWriter fieldData, Sorter.DocMap sortMap) + private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); - // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors - // from the flat writer to the CuVSMatrix? // TODO: This is slightly different.... List sortedVectors = new ArrayList(); @@ -912,15 +717,6 @@ static int distFuncToOrd(VectorSimilarityFunction func) { throw new IllegalArgumentException("invalid distance function: " + func); } - // We currently ignore this, until cuVS supports tiered indices - private static final String CANNOT_GENERATE_CAGRA = - """ - Could not generate an intermediate CAGRA graph because the initial \ - kNN graph contains too many invalid or duplicated neighbor nodes. \ - This error can occur, for example, if too many overflows occur \ - during the norm computation between the dataset vectors\ - """; - static void handleThrowableWithIgnore(Throwable t, String msg) throws IOException { if (t.getMessage().contains(msg)) { return; @@ -940,7 +736,7 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw // Access the CAGRA index for this field from the reader if (knnReader != null) { - if (knnReader instanceof CuVSVectorsReader cvr) { + if (knnReader instanceof GPUVectorsReader cvr) { if (cvr != null) { totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); @@ -1005,15 +801,15 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws /** * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. */ - private CagraIndex getCagraIndexFromReader(CuVSVectorsReader reader, String fieldName) { + private CagraIndex getCagraIndexFromReader(GPUVectorsReader reader, String fieldName) { try { - IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); + IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); FieldInfos fieldInfos = reader.getFieldInfos(); FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); if (fieldInfo != null) { - CuVSIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); + GPUIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); if (cuvsIndex != null) { return cuvsIndex.getCagraIndex(); } @@ -1039,11 +835,9 @@ private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, mergedIndex.serialize(cagraIndexOutputStream, tmpFile); long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; - // Write metadata (assuming no brute force or HNSW indexes for merged result) - // Only write meta for non-HNSW_LUCENE modes - if (indexType != IndexType.HNSW_LUCENE) { - writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); - } + writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); + + // TODO: Path to writeFieldInternal missing. Fix this. // Clean up the merged index mergedIndex.destroyIndex(); @@ -1056,28 +850,21 @@ private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { flatVectorsWriter.mergeOneField(fieldInfo, mergeState); - if (indexType.cagra() && !indexType.bruteForce()) { - // Since CAGRA merge does not support merging of indexes with purging of deletes, - // we fallback to vector-based re-indexing. Issue: - // https://github.com/rapidsai/cuvs/issues/1253 - boolean hasDeletions = - IntStream.range(0, mergeState.liveDocs.length) - .anyMatch( - i -> - mergeState.liveDocs[i] == null - || IntStream.range(0, mergeState.maxDocs[i]) - .anyMatch(j -> !mergeState.liveDocs[i].get(j))); - - if (mergeState.knnVectorsReaders.length > 1 && !hasDeletions) { - mergeCagraIndexes(fieldInfo, mergeState); - } else { - // CAGRA's merge API does not handle the trivial case of merging 1 index. - vectorBasedMerge(fieldInfo, mergeState); - } - + // Since CAGRA merge does not support merging of indexes with purging of deletes, + // we fallback to vector-based re-indexing. Issue: + // https://github.com/rapidsai/cuvs/issues/1253 + boolean hasDeletions = + IntStream.range(0, mergeState.liveDocs.length) + .anyMatch( + i -> + mergeState.liveDocs[i] == null + || IntStream.range(0, mergeState.maxDocs[i]) + .anyMatch(j -> !mergeState.liveDocs[i].get(j))); + + if (mergeState.knnVectorsReaders.length > 1 && !hasDeletions) { + mergeCagraIndexes(fieldInfo, mergeState); } else { - // If there is a Brute Force index then re-index using the vectors even if there is a CAGRA - // index. + // CAGRA's merge API does not handle the trivial case of merging 1 index. vectorBasedMerge(fieldInfo, mergeState); } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index c8d207bb..a025fbcc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -16,6 +16,7 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; import java.io.IOException; import java.time.Duration; import java.util.List; @@ -53,4 +54,27 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions) { static long nanosToMillis(long nanos) { return Duration.ofNanos(nanos).toMillis(); } + + static CuVSResources cuVSResourcesOrNull() { + try { + GPUVectorsFormat.resources = CuVSResources.create(); + return GPUVectorsFormat.resources; + } catch (UnsupportedOperationException uoe) { + GPUVectorsFormat.log.warning( + "cuvs is not supported on this platform or java version: " + uoe.getMessage()); + } catch (Throwable t) { + if (t instanceof ExceptionInInitializerError ex) { + t = ex.getCause(); + } + GPUVectorsFormat.log.warning("Exception occurred during creation of cuvs resources. " + t); + } + return null; + } + + static void handleThrowableWithIgnore(Throwable t, String msg) throws IOException { + if (t.getMessage().contains(msg)) { + return; + } + handleThrowable(t); + } } diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec index 7c0af61f..0f25e335 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -1,2 +1,2 @@ -com.nvidia.cuvs.lucene.CuVSCodec -com.nvidia.cuvs.lucene.CuVSCPUSearchCodec +com.nvidia.cuvs.lucene.HNSWSearchCodec +com.nvidia.cuvs.lucene.GPUSearchCodec diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 6e486ee1..8c50f8cb 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -15,4 +15,5 @@ org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat -com.nvidia.cuvs.lucene.CuVSVectorsFormat +com.nvidia.cuvs.lucene.GPUVectorsFormat +com.nvidia.cuvs.lucene.HNSWVectorsFormat diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 06a3a716..1fb06296 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -21,7 +21,9 @@ import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.ArrayList; import java.util.Arrays; +import java.util.List; import java.util.Random; import java.util.UUID; import java.util.logging.Logger; @@ -59,20 +61,22 @@ public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", GPUVectorsFormat.supported()); random = new Random(); + // Fixed seed so that we can validate against the same result. + random.setSeed(222); indexDirPath = Paths.get(UUID.randomUUID().toString()); } @Test public void testCagraToHnswSerializationAndSearch() throws IOException { - Codec codec = new CuVSCPUSearchCodec(); + Codec codec = new HNSWSearchCodec(); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); - int numDocs = random.nextInt(100, 1000); - int dimension = random.nextInt(8, 1024); - int topK = random.nextInt(5, 60); + int numDocs = 2000; // random.nextInt(100, 1000); + int dimension = 32; // random.nextInt(8, 1024); + int topK = 100; // random.nextInt(5, 60); final int COMMIT_FREQ = Math.min(numDocs, random.nextInt(100, 1000)); int count = COMMIT_FREQ; final String VECTOR_FIELD = "knn1"; @@ -126,6 +130,9 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { TopDocs results = searcher.search(query, topK); log.info("\nknn1 search results (" + results.totalHits + " total hits):"); + int[] expected = {1803, 1869, 554, 1824, 1982, 1302, 320, 351, 707, 549}; + List res = new ArrayList(); + for (int i = 0; i < results.scoreDocs.length; i++) { ScoreDoc scoreDoc = results.scoreDocs[i]; Document doc = searcher.storedFields().document(scoreDoc.doc); @@ -138,10 +145,14 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { + doc.get("id") + "), score=" + scoreDoc.score); + res.add(Integer.valueOf(doc.get("id"))); } assertTrue("TopK results not returned", results.scoreDocs.length == topK); // TODO: make this test a bit more meaningful like checking the quality of search results. + for (int i : expected) { + assertTrue("Expected doc id is missing:" + i, res.contains(i)); + } } catch (Exception e) { e.printStackTrace(); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index ef17136e..14c1128d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -56,12 +56,12 @@ public class TestCuVSDeletedDocuments extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); private static Random random; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); random = random(); } @@ -194,7 +194,7 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Test filtered search with deletions Query filter = new TermQuery(new Term("category", "A")); Query filteredQuery = - new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; for (ScoreDoc hit : filteredHits) { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index 6d4ab286..2054d9e9 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -53,7 +53,7 @@ public class TestCuVSGaps extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); static IndexSearcher searcher; static IndexReader reader; static Directory directory; @@ -70,7 +70,7 @@ public class TestCuVSGaps extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); directory = newDirectory(); random = random(); @@ -120,7 +120,7 @@ public static void afterClass() throws Exception { @Test public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; @@ -153,7 +153,7 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { @Test public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; @@ -163,7 +163,7 @@ public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOExcepti // This should further restrict our results to even numbers 0, 2, 4, 6, 8 Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 - Query filteredQuery = new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; // Should only get document 8 (the only one that matches the filter and has a vector) diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index a0973599..1e83f66f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -56,7 +56,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); static IndexSearcher searcher; static IndexReader reader; static Directory directory; @@ -69,7 +69,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); directory = newDirectory(); RandomIndexWriter writer = @@ -184,7 +184,7 @@ private static List> generateExpectedResults( @Test public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); Random random = random(); int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); @@ -206,7 +206,7 @@ public void testVectorSearchWithFilter() throws IOException { Query filter = new TermQuery(new Term("id", targetDocId)); // Test the new constructor with filter - Query filteredQuery = new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 84aa1a33..3299ae22 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -41,12 +41,12 @@ public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { @BeforeClass public static void beforeClass() { - assumeTrue("cuvs is not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuvs is not supported", GPUVectorsFormat.supported()); } @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + return TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index af2027cc..a65209eb 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -17,7 +17,7 @@ import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -70,7 +70,7 @@ public class TestMerge extends LuceneTestCase { @BeforeClass public static void beforeClass() { - assumeTrue("cuVS is not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS is not supported", GPUVectorsFormat.supported()); } private Directory directory; @@ -128,7 +128,7 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -252,7 +252,7 @@ public void testMergeWithIndexSorting() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) .setIndexSort(indexSort) // This automatically enables sorting during merges .setMergePolicy(mergePolicy) .setMaxBufferedDocs(maxBufferedDocs) @@ -458,7 +458,7 @@ public void testMergeWithMissingVectors() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -597,7 +597,7 @@ public void testMergeWithDeletions() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -730,8 +730,8 @@ public void testMergeBruteForceIndex() throws IOException { + vectorProbability); // Configure with brute force index type - CuVSVectorsFormat bruteForceFormat = - new CuVSVectorsFormat( + GPUVectorsFormat bruteForceFormat = + new GPUVectorsFormat( 32, // writer threads 128, // intermediate graph degree 64, // graph degree @@ -882,8 +882,8 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { + vectorProbability); // Configure with CAGRA + brute force combined index type - CuVSVectorsFormat combinedFormat = - new CuVSVectorsFormat( + GPUVectorsFormat combinedFormat = + new GPUVectorsFormat( 32, // writer threads 128, // intermediate graph degree 64, // graph degree @@ -1058,7 +1058,7 @@ public void testLargeScaleMerge() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); From 62297696d36e0854941586cd0965ed0fb5650ac3 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Mon, 25 Aug 2025 13:05:14 -0400 Subject: [PATCH 04/21] Code refactoring - refinement, simplification, and cleanup. --- pom.xml | 2 +- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 85 ++-- .../java/com/nvidia/cuvs/lucene/GPUIndex.java | 8 +- .../nvidia/cuvs/lucene/GPUSearchCodec.java | 4 +- .../nvidia/cuvs/lucene/GPUVectorsWriter.java | 8 +- .../nvidia/cuvs/lucene/HNSWSearchCodec.java | 6 +- .../nvidia/cuvs/lucene/HNSWVectorsFormat.java | 14 +- .../nvidia/cuvs/lucene/HNSWVectorsWriter.java | 396 ++++-------------- .../cuvs/lucene/IndexOutputOutputStream.java | 14 +- 9 files changed, 148 insertions(+), 389 deletions(-) diff --git a/pom.xml b/pom.xml index 5c1969f7..cd9ee5bf 100644 --- a/pom.xml +++ b/pom.xml @@ -68,7 +68,7 @@ com.nvidia.cuvs cuvs-java - 25.8.0-c9599-SNAPSHOT + 25.10.0-d0a83-SNAPSHOT diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index e5ecd617..7235ec97 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -17,6 +17,8 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.RowView; import java.util.ArrayList; import java.util.List; import org.apache.lucene.util.hnsw.HnswGraph; @@ -37,7 +39,7 @@ public class GPUBuiltHnswGraph extends HnswGraph { // Multi-layer constructor that supports arbitrary number of layers public GPUBuiltHnswGraph( - int size, int dimensions, List layerNodes, List layerAdjacencies) { + int size, int dimensions, List layerNodes, List layerAdjacencies) { this.size = size; this.dimensions = dimensions; @@ -46,59 +48,42 @@ public GPUBuiltHnswGraph( this.layerNeighbors = new ArrayList<>(); // Process Layer 0 (base layer with all nodes) - int[][] layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = new NeighborArray[size]; - - for (int i = 0; i < size; i++) { - if (layer0Adjacency[i] != null && layer0Adjacency[i].length > 0) { - layer0Neighbors[i] = new NeighborArray(layer0Adjacency[i].length, true); - for (int j = 0; j < layer0Adjacency[i].length; j++) { - layer0Neighbors[i].addInOrder(layer0Adjacency[i][j], 1.0f - (j * 0.001f)); - } - } else { - layer0Neighbors[i] = new NeighborArray(0, true); - } - } + CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { int[] nodes = layerNodes.get(level); - int[][] adjacency = layerAdjacencies.get(level); - + CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); - NeighborArray[] neighbors = new NeighborArray[nodes.length]; + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + } + } - for (int i = 0; i < nodes.length; i++) { - if (adjacency[i] != null && adjacency[i].length > 0) { - neighbors[i] = new NeighborArray(adjacency[i].length, true); - for (int j = 0; j < adjacency[i].length; j++) { - neighbors[i].addInOrder(adjacency[i][j], 1.0f - (j * 0.001f)); - } - } else { - neighbors[i] = new NeighborArray(0, true); + private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + NeighborArray[] neighbors = new NeighborArray[size]; + for (int i = 0; i < size; i++) { + RowView rv = adjacency.getRow(i); + if (rv != null && rv.size() > 0) { + neighbors[i] = 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)); } + } else { + neighbors[i] = new NeighborArray(0, true); } - - this.layerNeighbors.add(neighbors); } - } - - public int size() { - return size; - } - - public int numLevels() { - return numLevels; + return neighbors; } public NodesIterator getNodesOnLevel(int level) { if (level == 0) { - return new ArrayNodesIterator(size); + return new Level0NodesIterator(size); } else if (level > 0 && level < numLevels) { int[] nodes = layerNodes.get(level - 1); - return new SpecificNodesIterator(nodes); + return new HigherLevelNodesIterator(nodes); } else { - return new ArrayNodesIterator(0); + return new Level0NodesIterator(0); } } @@ -201,11 +186,11 @@ public int neighborCount() { return 0; } - // Simple implementation of NodesIterator for level 0 - private static class ArrayNodesIterator extends NodesIterator { + // NodesIterator for level 0 + private static class Level0NodesIterator extends NodesIterator { private int current = -1; - ArrayNodesIterator(int size) { + Level0NodesIterator(int size) { super(size); } @@ -229,12 +214,12 @@ public int consume(int[] dest) { } } - // NodesIterator for specific nodes in higher layers - private static class SpecificNodesIterator extends NodesIterator { + // NodesIterator for higher layers + private static class HigherLevelNodesIterator extends NodesIterator { private final int[] nodeIds; private int current = -1; - SpecificNodesIterator(int[] nodeIds) { + HigherLevelNodesIterator(int[] nodeIds) { super(nodeIds.length); this.nodeIds = nodeIds; } @@ -258,4 +243,16 @@ public int consume(int[] dest) { return numToCopy; } } + + public int size() { + return size; + } + + public int numLevels() { + return numLevels; + } + + public int dimensions() { + return dimensions; + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java index 651354d6..ea33c842 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java @@ -47,7 +47,7 @@ public GPUIndex( throw new IllegalArgumentException("negative maxDocs:" + maxDocs); } this.maxDocs = maxDocs; - this.hnswIndex = null; // TODO: + this.hnswIndex = null; // TODO: remove hnswlib logic in a subsequent PR } public GPUIndex(CagraIndex cagraIndex, BruteForceIndex bruteforceIndex, HnswIndex hnswIndex) { @@ -101,13 +101,13 @@ public void close() throws IOException { private void destroyIndices() throws IOException { try { if (cagraIndex != null) { - cagraIndex.destroyIndex(); + cagraIndex.close(); } if (bruteforceIndex != null) { - bruteforceIndex.destroyIndex(); + bruteforceIndex.close(); } if (hnswIndex != null) { - hnswIndex.destroyIndex(); + hnswIndex.close(); } } catch (Throwable t) { Utils.handleThrowable(t); diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java index 97ca2621..08e6a1ea 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java @@ -27,7 +27,7 @@ public class GPUSearchCodec extends FilterCodec { private static final Logger log = Logger.getLogger(GPUSearchCodec.class.getName()); - private static final String CLASS_NAME = "GPUSearchCodec"; + private static final String NAME = "GPUSearchCodec"; private static final int DEFAULT_CUVS_WRITER_THREADS = 1; private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; @@ -38,7 +38,7 @@ public class GPUSearchCodec extends FilterCodec { private KnnVectorsFormat format; public GPUSearchCodec() { - this(CLASS_NAME, new Lucene101Codec()); + this(NAME, new Lucene101Codec()); } public GPUSearchCodec(String name, Codec delegate) { diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java index 5bad8534..547a875d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java @@ -305,7 +305,7 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab info("Cagra index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); Path tmpFile = Files.createTempFile(resources.tempDirectory(), "tmpindex", "cag"); index.serialize(os, tmpFile); - index.destroyIndex(); + index.close(); } private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { @@ -319,7 +319,7 @@ private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Th long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("bf index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); index.serialize(os); - index.destroyIndex(); + index.close(); } private void writeHNSWIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { @@ -334,7 +334,7 @@ private void writeHNSWIndex(OutputStream os, CuVSMatrix dataset) throws Throwabl info("HNSW index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); Path tmpFile = Files.createTempFile("tmpindex", "hnsw"); index.serializeToHNSW(os, tmpFile); - index.destroyIndex(); + index.close(); } @Override @@ -526,7 +526,7 @@ private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); // Clean up the merged index - mergedIndex.destroyIndex(); + mergedIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java index df94f24c..32627141 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java @@ -31,12 +31,12 @@ public class HNSWSearchCodec extends FilterCodec { private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; private static final int DEFAULT_GRAPH_DEGREE = 64; private static final int DEFAULT_HNSW_LAYERS = 1; - private static final String CLASS_NAME = "HNSWSearchCodec"; + private static final String NAME = "HNSWSearchCodec"; private KnnVectorsFormat format; public HNSWSearchCodec() { - this(CLASS_NAME, new Lucene101Codec()); + this(NAME, new Lucene101Codec()); } public HNSWSearchCodec(String name, Codec delegate) { @@ -46,7 +46,7 @@ public HNSWSearchCodec(String name, Codec delegate) { public HNSWSearchCodec( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { - this(CLASS_NAME, new Lucene101Codec()); + this(NAME, new Lucene101Codec()); initializeFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java index 4cf91971..07b8a412 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java @@ -33,20 +33,16 @@ public class HNSWVectorsFormat extends KnnVectorsFormat { private static final Logger log = Logger.getLogger(HNSWVectorsFormat.class.getName()); - // TODO: fix Lucene version in name, to the final targeted release, if any - static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; - static final String CUVS_META_CODEC_EXT = "vemc"; - static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex"; - static final String CUVS_INDEX_EXT = "vcag"; - - static final int VERSION_START = 0; - static final int VERSION_CURRENT = VERSION_START; - static final int DEFAULT_WRITER_THREADS = 32; static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; static final int DEFAULT_GRAPH_DEGREE = 64; static final int HNSW_GRAPH_LAYERS = 1; + static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; + static final String HNSW_META_CODEC_EXT = "vem"; + static final String HNSW_INDEX_CODEC_NAME = "Lucene99HnswVectorsFormatIndex"; + static final String HNSW_INDEX_EXT = "vex"; + static CuVSResources resources = cuVSResourcesOrNull(); /** The format for storing, reading, and merging raw vectors on disk. */ diff --git a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java index 160bf3f5..456b81ee 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java @@ -15,6 +15,10 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_META_CODEC_NAME; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -32,6 +36,9 @@ import java.util.Arrays; import java.util.List; import java.util.Objects; +import java.util.Random; +import java.util.SortedSet; +import java.util.TreeSet; import java.util.logging.Logger; import java.util.stream.IntStream; import org.apache.lucene.codecs.CodecUtil; @@ -73,24 +80,18 @@ public class HNSWVectorsWriter extends KnnVectorsWriter { private static final Logger log = Logger.getLogger(HNSWVectorsWriter.class.getName()); /** The name of the CUVS component for the info-stream * */ - public static final String CUVS_COMPONENT = "CUVS"; - - // The minimum number of vectors in the dataset required before - // we attempt to build a Cagra index - static final int MIN_CAGRA_INDEX_SIZE = 2; + private static final String CUVS_COMPONENT = "CUVS"; private final int cuvsWriterThreads; private final int intGraphDegree; private final int graphDegree; private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion - private final CuVSResources resources; - private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors private final List fields = new ArrayList<>(); - private IndexOutput meta = null, cuvsIndex = null; - private IndexOutput hnswMeta = null, hnswVectorIndex = null; private final InfoStream infoStream; + private IndexOutput cuvsIndex = null; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; private String vemFileName; private String vexFileName; @@ -114,10 +115,11 @@ public HNSWVectorsWriter( this.infoStream = state.infoStream; vemFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "vem"); + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); vexFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "vex"); + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); boolean success = false; try { @@ -127,13 +129,13 @@ public HNSWVectorsWriter( CodecUtil.writeIndexHeader( hnswMeta, - "Lucene99HnswVectorsFormatMeta", + HNSW_META_CODEC_NAME, Lucene99HnswVectorsFormat.VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); CodecUtil.writeIndexHeader( hnswVectorIndex, - "Lucene99HnswVectorsFormatIndex", + HNSW_INDEX_CODEC_NAME, Lucene99HnswVectorsFormat.VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); @@ -170,12 +172,7 @@ static String indexMsg(int size, int... args) { return sb.toString(); } - private CagraIndexParams cagraIndexParams(int size) { - if (size < 2) { - // https://github.com/rapidsai/cuvs/issues/666 - throw new IllegalArgumentException("cagra index must be greater than 2"); - } - + private CagraIndexParams cagraIndexParams() { return new CagraIndexParams.Builder() .withNumWriterThreads(cuvsWriterThreads) .withIntermediateGraphDegree(intGraphDegree) @@ -191,94 +188,40 @@ private void info(String msg) { } private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { writeEmpty(fieldInfo); return; } try { - info("=== ENTERED HNSW_LUCENE BLOCK (HNSW-only mode) ==="); - info("Entered the writeFieldInternal's HNSW LUCENE block - writing only HNSW files"); CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); if (dataset.size() < 2) { throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); } - CagraIndexParams params = cagraIndexParams((int) dataset.size()); + long startTime = System.nanoTime(); - CagraIndex index = + CagraIndexParams params = cagraIndexParams(); + CagraIndex cagraIndex = CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); // Get the adjacency list from CAGRA index - int[][] adjacencyList; - try { - adjacencyList = index.getGraph(); - info("=== SUCCESS: Got adjacency list from CAGRA index ==="); - info("Successfully got adjacency list from CAGRA index"); - } catch (Exception e) { - info("=== FAILED: getGraph() method failed: " + e.getMessage() + " ==="); - info("getGraph() method failed or doesn't exist: " + e.getMessage()); - // Create a mock adjacency list for testing - int size = (int) dataset.size(); - adjacencyList = new int[size][]; - for (int i = 0; i < size; i++) { - // Create connections to next few nodes (circular) - int degree = Math.min(10, size - 1); // up to 10 connections - adjacencyList[i] = new int[degree]; - for (int j = 0; j < degree; j++) { - adjacencyList[i][j] = (i + j + 1) % size; - } - } - info( - "=== CREATED MOCK ADJACENCY LIST: " - + size - + " nodes, degree=" - + (adjacencyList.length > 0 ? adjacencyList[0].length : 0) - + " ==="); - info( - "Created mock adjacency list with " - + size - + " nodes, degree=" - + (adjacencyList.length > 0 ? adjacencyList[0].length : 0)); - } + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); int size = (int) dataset.size(); int dimensions = fieldInfo.getVectorDimension(); - // Debug: Check if we got valid adjacency data - info( - "Adjacency list info: " - + (adjacencyList == null - ? "null" - : "length=" - + adjacencyList.length - + ", first row=" - + (adjacencyList.length > 0 && adjacencyList[0] != null - ? adjacencyList[0].length - : "null"))); - - // Create HNSW graph from CAGRA - multi-layer if original vectors available - GPUBuiltHnswGraph hnswGraph; - if (vectors != null && vectors.size() > 0) { - info("=== Creating 3-layer HNSW graph using original vectors ==="); - hnswGraph = createMultiLayerHnswGraph(fieldInfo, size, dimensions, adjacencyList, vectors); - } else { - info("=== Creating single-layer HNSW graph (no original vectors) ==="); - // Create single layer graph - List singleLayerNodes = new ArrayList<>(); - List singleLayerAdjacencies = new ArrayList<>(); - singleLayerAdjacencies.add(adjacencyList); - hnswGraph = - new GPUBuiltHnswGraph(size, dimensions, singleLayerNodes, singleLayerAdjacencies); - } + // Create multi-layer HNSW graph from CAGRA + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, size, dimensions, adjacencyListMatrix, vectors, hnswLayers); - // Remember the vector index offset before writing long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - // Calculate the length of written data long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -293,15 +236,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro graphLevelNodeOffsets); long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); - info( - "HNSW-only graph created in " - + elapsedMillis - + "ms, with " - + dataset.size() - + " vectors"); + info("HNSW graph created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); - // Don't serialize CAGRA index - destroy it immediately - index.destroyIndex(); + cagraIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); @@ -318,50 +255,32 @@ private GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int size, int dimensions, - int[][] adjacencyList, - List originalVectors) + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers) throws Throwable { // Calculate M as cagraGraphDegree/2 int M = graphDegree / 2; - info( - "=== Creating " - + hnswLayers - + "-layer HNSW graph with M=" - + M - + " (cagraGraphDegree/2) ==="); - info("Creating " + hnswLayers + "-layer HNSW graph with size=" + size + ", M=" + M); // Store all layers data - java.util.List layerNodes = new java.util.ArrayList<>(); - java.util.List layerAdjacencies = new java.util.ArrayList<>(); + 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(adjacencyList); + layerAdjacencies.add(adjacencyListMatrix); - // Build higher layers - create exactly hnswLayers-1 additional layers (layer 0 is already - // added) int currentLayerSize = size; int layerIndex = 1; - java.util.Random random = new java.util.Random(42); // Fixed seed for reproducibility + Random random = new Random(); while (layerIndex < hnswLayers && currentLayerSize > 1) { // Calculate size for next layer (1/M of current layer) int nextLayerSize = Math.max(1, currentLayerSize / M); - info( - "=== Layer " - + layerIndex - + " will have " - + nextLayerSize - + " nodes out of " - + currentLayerSize - + " (previous layer) ==="); - info("Layer " + layerIndex + " will have " + nextLayerSize + " nodes"); - // Select nodes for this layer - java.util.Set selectedNodesSet = new java.util.HashSet<>(); + SortedSet selectedNodesSet = new TreeSet<>(); if (layerIndex == 1) { // Select from all nodes (Layer 0) @@ -380,136 +299,48 @@ private GPUBuiltHnswGraph createMultiLayerHnswGraph( // Convert to sorted array int[] selectedNodes = selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - layerNodes.add(selectedNodes); - info( - "=== Selected Layer " - + layerIndex - + " nodes: " - + java.util.Arrays.toString( - java.util.Arrays.copyOf(selectedNodes, Math.min(10, selectedNodes.length))) - + (selectedNodes.length > 10 ? "..." : "") - + " ==="); + layerNodes.add(selectedNodes); // Extract vectors for selected nodes float[][] selectedVectors = new float[nextLayerSize][]; for (int i = 0; i < nextLayerSize; i++) { - int nodeId = selectedNodes[i]; - if (nodeId < originalVectors.size()) { - selectedVectors[i] = originalVectors.get(nodeId); - } else { - selectedVectors[i] = createRandomVector(dimensions, nodeId); - } + selectedVectors[i] = vectors.get(selectedNodes[i]); } // Build CAGRA graph for this layer - int[][] layerAdjacency = buildCagraGraphForSubset(selectedVectors, selectedNodes); - layerAdjacencies.add(layerAdjacency); + layerAdjacencies.add(buildCagraGraphForSubset(selectedVectors)); // Update for next iteration currentLayerSize = nextLayerSize; layerIndex++; // Use different seed for each layer - random = new java.util.Random(42 + layerIndex); + random = new Random(new Random().nextLong()); } - int numLayers = layerAdjacencies.size(); - info("=== Total layers created: " + numLayers + " ==="); - info("Created " + numLayers + " layers total"); - // Create the multi-layer graph with all layers return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } - /** - * Creates a random vector for fallback purposes - */ - private float[] createRandomVector(int dimensions, int seed) { - float[] vector = new float[dimensions]; - java.util.Random random = new java.util.Random(seed); - for (int i = 0; i < dimensions; i++) { - vector[i] = random.nextFloat(); - } - return vector; - } - /** * Builds a CAGRA graph for a subset of vectors */ - private int[][] buildCagraGraphForSubset(float[][] vectors, int[] originalNodeIds) - throws Throwable { - if (vectors.length < 2) { - // Can't build CAGRA graph with less than 2 vectors - return new int[vectors.length][0]; - } + private CuVSMatrix buildCagraGraphForSubset(float[][] vectors) throws Throwable { - try { - // Create CuVSMatrix from the subset vectors - CuVSMatrix subsetDataset = CuVSMatrix.ofArray(vectors); - - // Build CAGRA index for the subset - CagraIndexParams params = cagraIndexParams(vectors.length); - CagraIndex subsetIndex = - CagraIndex.newBuilder(resources) - .withDataset(subsetDataset) - .withIndexParams(params) - .build(); - - // Get adjacency list from subset CAGRA index - int[][] subsetAdjacency; - try { - subsetAdjacency = subsetIndex.getGraph(); - info("=== SUCCESS: Got adjacency list from Layer 1 CAGRA index ==="); - info("Successfully got adjacency list from Layer 1 CAGRA index"); - } catch (Exception e) { - info("=== FAILED: getGraph() method failed for Layer 1: " + e.getMessage() + " ==="); - info("getGraph() method failed for Layer 1: " + e.getMessage()); - // Create mock adjacency list - subsetAdjacency = new int[vectors.length][]; - for (int i = 0; i < vectors.length; i++) { - int degree = Math.min(5, vectors.length - 1); - subsetAdjacency[i] = new int[degree]; - for (int j = 0; j < degree; j++) { - subsetAdjacency[i][j] = (i + j + 1) % vectors.length; - } - } - } + // Create CuVSMatrix from the subset vectors + CuVSMatrix subsetDataset = CuVSMatrix.ofArray(vectors); - // Convert subset adjacency to use original node IDs - int[][] layer1Adjacency = new int[vectors.length][]; - for (int i = 0; i < vectors.length; i++) { - if (subsetAdjacency[i] != null) { - layer1Adjacency[i] = new int[subsetAdjacency[i].length]; - for (int j = 0; j < subsetAdjacency[i].length; j++) { - // Map subset index back to original node ID - int subsetNeighborId = subsetAdjacency[i][j]; - layer1Adjacency[i][j] = originalNodeIds[subsetNeighborId]; - } - } else { - layer1Adjacency[i] = new int[0]; - } - } + // Build CAGRA index for the subset + CagraIndexParams params = cagraIndexParams(); + CagraIndex subsetIndex = + CagraIndex.newBuilder(resources).withDataset(subsetDataset).withIndexParams(params).build(); - subsetIndex.destroyIndex(); - return layer1Adjacency; + // Get adjacency list from subset CAGRA index + CuVSMatrix cagraGraph = subsetIndex.getGraph(); - } catch (Exception e) { - info("=== FAILED to build CAGRA graph for subset: " + e.getMessage() + " ==="); - info("Failed to build CAGRA graph for subset: " + e.getMessage()); - - // Fallback: create simple connections between Layer 1 nodes - int[][] fallbackAdjacency = new int[vectors.length][]; - for (int i = 0; i < vectors.length; i++) { - int degree = Math.min(3, vectors.length - 1); - fallbackAdjacency[i] = new int[degree]; - for (int j = 0; j < degree; j++) { - int neighborIdx = (i + j + 1) % vectors.length; - fallbackAdjacency[i][j] = originalNodeIds[neighborIdx]; - } - } - return fallbackAdjacency; - } + subsetIndex.close(); + return cagraGraph; } private void writeMeta( @@ -522,14 +353,7 @@ private void writeMeta( HnswGraph graph, int[][] graphLevelNodeOffsets) throws IOException { - info( - "=== writeMeta: Writing field " - + field.name - + " with count=" - + count - + ", dimensions=" - + field.getVectorDimension() - + " ==="); + meta.writeInt(field.number); meta.writeInt(field.getVectorEncoding().ordinal()); meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); @@ -537,10 +361,8 @@ private void writeMeta( meta.writeVLong(vectorIndexLength); meta.writeVInt(field.getVectorDimension()); meta.writeInt(count); - // Use M = cagraGraphDegree/2 - int M = graphDegree / 2; - info("=== writeMeta: Writing M=" + M + " (cagraGraphDegree/2) ==="); - meta.writeVInt(M); // M = cagraGraphDegree/2 + meta.writeVInt(graphDegree / 2); // M = cagraGraphDegree/2 + // write graph nodes on each level if (graph == null) { meta.writeVInt(0); @@ -560,45 +382,34 @@ private void writeMeta( nol[i] -= nol[i - 1]; } for (int n : nol) { - assert n >= 0 : "delta encoding for nodes failed; expected nodes to be sorted"; meta.writeVInt(n); } } else { assert nodesOnLevel.size() == count : "Level 0 expects to have all nodes"; } } + long start = vectorIndex.getFilePointer(); meta.writeLong(start); meta.writeVInt(16); // DIRECT_MONOTONIC_BLOCK_SHIFT); + final DirectMonotonicWriter memoryOffsetsWriter = - DirectMonotonicWriter.getInstance( - meta, vectorIndex, valueCount, 16); // DIRECT_MONOTONIC_BLOCK_SHIFT); + DirectMonotonicWriter.getInstance(meta, vectorIndex, valueCount, 16); long cumulativeOffsetSum = 0; - int totalOffsetsWritten = 0; for (int[] levelOffsets : graphLevelNodeOffsets) { - info( - "=== writeMeta: Writing offsets for level with " - + levelOffsets.length - + " entries ==="); for (int v : levelOffsets) { memoryOffsetsWriter.add(cumulativeOffsetSum); cumulativeOffsetSum += v; - totalOffsetsWritten++; } } - info( - "=== writeMeta: Total offsets written: " - + totalOffsetsWritten - + ", expected: " - + valueCount - + " ==="); + memoryOffsetsWriter.finish(); + meta.writeLong(vectorIndex.getFilePointer() - start); } } private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { - if (graph == null) return new int[0][0]; // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); int[][] offsets = new int[graph.numLevels()][]; @@ -607,21 +418,17 @@ private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) thr int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; int nodeOffsetId = 0; - // Debug: print the actual number of nodes being processed - info( - "=== writeGraph: Level " - + level - + " has " - + sortedNodes.length - + " nodes, expected " - + (level == 0 ? countOnLevel0 : "unknown") - + " ==="); + 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 @@ -630,8 +437,10 @@ private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) thr 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]) { continue; } @@ -639,6 +448,7 @@ private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) thr } // 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]); } @@ -646,6 +456,7 @@ private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) thr Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } } + // Return offsets (information written while writing the meta info) return offsets; } @@ -672,7 +483,6 @@ private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); - // TODO: This is slightly different.... List sortedVectors = new ArrayList(); for (int i = 0; i < fieldData.getVectors().size(); i++) { sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); @@ -682,30 +492,7 @@ private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) } private void writeEmpty(FieldInfo fieldInfo) throws IOException { - writeMeta(fieldInfo, 0, 0L, 0L, 0L, 0L, 0L, 0L); - } - - private void writeMeta( - FieldInfo field, - int count, - long cagraIndexOffset, - long cagraIndexLength, - long bruteForceIndexOffset, - long bruteForceIndexLength, - long hnswIndexOffset, - long hnswIndexLength) - throws IOException { - meta.writeInt(field.number); - meta.writeInt(field.getVectorEncoding().ordinal()); - meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); - meta.writeInt(field.getVectorDimension()); - meta.writeInt(count); - meta.writeVLong(cagraIndexOffset); - meta.writeVLong(cagraIndexLength); - meta.writeVLong(bruteForceIndexOffset); - meta.writeVLong(bruteForceIndexLength); - meta.writeVLong(hnswIndexOffset); - meta.writeVLong(hnswIndexLength); + writeMeta(null, hnswMeta, fieldInfo, 0, 0, 0, null, null); } static int distFuncToOrd(VectorSimilarityFunction func) { @@ -717,13 +504,6 @@ static int distFuncToOrd(VectorSimilarityFunction func) { throw new IllegalArgumentException("invalid distance function: " + func); } - static void handleThrowableWithIgnore(Throwable t, String msg) throws IOException { - if (t.getMessage().contains(msg)) { - return; - } - Utils.handleThrowable(t); - } - private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { @@ -768,14 +548,13 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw * */ private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) throws IOException { - List res = new ArrayList(); + List vectors = new ArrayList(); KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - int ordinal = iter.index(); - float[] vector = mergedVectorValues.vectorValue(ordinal); - res.add(vector); + float[] vector = mergedVectorValues.vectorValue(iter.index()); + vectors.add(vector); } - return res; + return vectors; } /** @@ -788,7 +567,6 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws throw new AssertionError("Only Float32 supported"); } try { - List dataset = createListFromMergedVectors( KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); @@ -805,7 +583,6 @@ private CagraIndex getCagraIndexFromReader(GPUVectorsReader reader, String field try { IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); FieldInfos fieldInfos = reader.getFieldInfos(); - FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); if (fieldInfo != null) { @@ -827,20 +604,16 @@ private CagraIndex getCagraIndexFromReader(GPUVectorsReader reader, String field private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, int vectorCount) throws IOException { try { - long cagraIndexOffset = cuvsIndex.getFilePointer(); var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); // Serialize the merged index Path tmpFile = Files.createTempFile(resources.tempDirectory(), "mergedindex", "cag"); mergedIndex.serialize(cagraIndexOutputStream, tmpFile); - long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; - - writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); // TODO: Path to writeFieldInternal missing. Fix this. // Clean up the merged index - mergedIndex.destroyIndex(); + mergedIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); } @@ -877,30 +650,23 @@ public void finish() throws IOException { finished = true; flatVectorsWriter.finish(); - if (meta != null) { - // write end of fields marker - meta.writeInt(-1); - CodecUtil.writeFooter(meta); - } if (cuvsIndex != null) { CodecUtil.writeFooter(cuvsIndex); } - { - if (hnswMeta != null) { - // write end of fields marker - hnswMeta.writeInt(-1); - CodecUtil.writeFooter(hnswMeta); - } - if (hnswVectorIndex != null) { - CodecUtil.writeFooter(hnswVectorIndex); - } + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); } } @Override public void close() throws IOException { - IOUtils.close(meta, cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); + IOUtils.close(cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java index b7866363..53283ebf 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java +++ b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java @@ -27,7 +27,7 @@ final class IndexOutputOutputStream extends OutputStream { final IndexOutput out; final int bufferSize; final byte[] buffer; - int idx; + int pos; IndexOutputOutputStream(IndexOutput out) { this(out, DEFAULT_BUFFER_SIZE); @@ -41,16 +41,16 @@ final class IndexOutputOutputStream extends OutputStream { @Override public void write(int b) throws IOException { - buffer[idx] = (byte) b; - idx++; - if (idx == bufferSize) { + buffer[pos] = (byte) b; + pos++; + if (pos == bufferSize) { flush(); } } @Override public void write(byte[] b, int offset, int length) throws IOException { - if (idx != 0) { + if (pos != 0) { flush(); } out.writeBytes(b, offset, length); @@ -58,8 +58,8 @@ public void write(byte[] b, int offset, int length) throws IOException { @Override public void flush() throws IOException { - out.writeBytes(buffer, 0, idx); - idx = 0; + out.writeBytes(buffer, 0, pos); + pos = 0; } @Override From 660cfc9302df7d90c2a2aabdb50ec4261d64c313 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Mon, 25 Aug 2025 14:56:12 -0400 Subject: [PATCH 05/21] Versioning in Codecs, Formats, Readers, and Writers --- ...Codec.java => CuVS2510GPUSearchCodec.java} | 20 +++++++++------- ...mat.java => CuVS2510GPUVectorsFormat.java} | 18 +++++++-------- ...der.java => CuVS2510GPUVectorsReader.java} | 18 +++++++-------- ...ter.java => CuVS2510GPUVectorsWriter.java} | 23 ++++++++++--------- ...ava => Lucene101AcceleratedHNSWCodec.java} | 16 +++++++------ ...Lucene99AcceleratedHNSWVectorsFormat.java} | 14 ++++++----- ...Lucene99AcceleratedHNSWVectorsWriter.java} | 22 ++++++++++-------- .../java/com/nvidia/cuvs/lucene/Utils.java | 9 ++++---- .../services/org.apache.lucene.codecs.Codec | 4 ++-- .../org.apache.lucene.codecs.KnnVectorsFormat | 4 ++-- ...TestCagraToHnswSerializationAndSearch.java | 4 ++-- .../cuvs/lucene/TestCuVSDeletedDocuments.java | 4 ++-- .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 8 +++---- .../TestCuVSRandomizedVectorSearch.java | 6 ++--- .../cuvs/lucene/TestCuVSVectorsFormat.java | 4 ++-- .../com/nvidia/cuvs/lucene/TestMerge.java | 22 +++++++++--------- 16 files changed, 104 insertions(+), 92 deletions(-) rename src/main/java/com/nvidia/cuvs/lucene/{GPUSearchCodec.java => CuVS2510GPUSearchCodec.java} (75%) rename src/main/java/com/nvidia/cuvs/lucene/{GPUVectorsFormat.java => CuVS2510GPUVectorsFormat.java} (87%) rename src/main/java/com/nvidia/cuvs/lucene/{GPUVectorsReader.java => CuVS2510GPUVectorsReader.java} (96%) rename src/main/java/com/nvidia/cuvs/lucene/{GPUVectorsWriter.java => CuVS2510GPUVectorsWriter.java} (95%) rename src/main/java/com/nvidia/cuvs/lucene/{HNSWSearchCodec.java => Lucene101AcceleratedHNSWCodec.java} (80%) rename src/main/java/com/nvidia/cuvs/lucene/{HNSWVectorsFormat.java => Lucene99AcceleratedHNSWVectorsFormat.java} (91%) rename src/main/java/com/nvidia/cuvs/lucene/{HNSWVectorsWriter.java => Lucene99AcceleratedHNSWVectorsWriter.java} (96%) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java similarity index 75% rename from src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 08e6a1ea..4027d3ab 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -16,18 +16,22 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.GPUVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.FilterCodec; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene101.Lucene101Codec; -/** CuVS based codec for GPU based vector search */ -public class GPUSearchCodec extends FilterCodec { +/** CuVS based codec for GPU based vector search + * + * @apiNote cuVS serialization formats are in experimental phase and hence backward compatibility cannot be guaranteed. + * + * */ +public class CuVS2510GPUSearchCodec extends FilterCodec { - private static final Logger log = Logger.getLogger(GPUSearchCodec.class.getName()); - private static final String NAME = "GPUSearchCodec"; + private static final Logger log = Logger.getLogger(CuVS2510GPUSearchCodec.class.getName()); + private static final String NAME = "CuVS2510GPUSearchCodec"; private static final int DEFAULT_CUVS_WRITER_THREADS = 1; private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; @@ -37,15 +41,15 @@ public class GPUSearchCodec extends FilterCodec { private KnnVectorsFormat format; - public GPUSearchCodec() { + public CuVS2510GPUSearchCodec() { this(NAME, new Lucene101Codec()); } - public GPUSearchCodec(String name, Codec delegate) { + public CuVS2510GPUSearchCodec(String name, Codec delegate) { super(name, delegate); try { format = - new GPUVectorsFormat( + new CuVS2510GPUVectorsFormat( DEFAULT_CUVS_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java similarity index 87% rename from src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 2fb058d7..98dbe2a9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -17,7 +17,7 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.GPUVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.logging.Logger; import org.apache.lucene.codecs.KnnVectorsFormat; @@ -29,9 +29,9 @@ import org.apache.lucene.index.SegmentWriteState; /** CuVS based KnnVectorsFormat for GPU acceleration */ -public class GPUVectorsFormat extends KnnVectorsFormat { +public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { - static final Logger log = Logger.getLogger(GPUVectorsFormat.class.getName()); + static final Logger log = Logger.getLogger(CuVS2510GPUVectorsFormat.class.getName()); // TODO: fix Lucene version in name, to the final targeted release, if any static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; @@ -59,14 +59,14 @@ public class GPUVectorsFormat extends KnnVectorsFormat { final int intGraphDegree; final int graphDegree; final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion - final GPUVectorsWriter.IndexType indexType; // the index type to build, when writing + final CuVS2510GPUVectorsWriter.IndexType indexType; // the index type to build, when writing /** * Creates a CuVSVectorsFormat, with default values. * * @throws LibraryException if the native library fails to load */ - public GPUVectorsFormat() { + public CuVS2510GPUVectorsFormat() { this( DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, @@ -80,7 +80,7 @@ public GPUVectorsFormat() { * * @throws LibraryException if the native library fails to load */ - public GPUVectorsFormat( + public CuVS2510GPUVectorsFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, @@ -106,10 +106,10 @@ private static void checkSupported() { } @Override - public GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + public CuVS2510GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { checkSupported(); var flatWriter = flatVectorsFormat.fieldsWriter(state); - return new GPUVectorsWriter( + return new CuVS2510GPUVectorsWriter( state, cuvsWriterThreads, intGraphDegree, @@ -123,7 +123,7 @@ public GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { checkSupported(); - return new GPUVectorsReader(state, resources, flatVectorsFormat.fieldsReader(state)); + return new CuVS2510GPUVectorsReader(state, resources, flatVectorsFormat.fieldsReader(state)); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java similarity index 96% rename from src/main/java/com/nvidia/cuvs/lucene/GPUVectorsReader.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 541044cd..a9fe61a5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -15,12 +15,12 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_EXT; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.VERSION_CURRENT; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.VERSION_START; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_START; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.BruteForceQuery; @@ -62,10 +62,10 @@ import org.apache.lucene.util.hnsw.IntToIntFunction; /** KnnVectorsReader instance associated with CuVS format */ -public class GPUVectorsReader extends KnnVectorsReader { +public class CuVS2510GPUVectorsReader extends KnnVectorsReader { @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(GPUVectorsReader.class.getName()); + private static final Logger log = Logger.getLogger(CuVS2510GPUVectorsReader.class.getName()); private final CuVSResources resources; private final FlatVectorsReader flatVectorsReader; // for reading the raw vectors @@ -74,7 +74,7 @@ public class GPUVectorsReader extends KnnVectorsReader { private final IntObjectHashMap cuvsIndices; private final IndexInput cuvsIndexInput; - public GPUVectorsReader( + public CuVS2510GPUVectorsReader( SegmentReadState state, CuVSResources resources, FlatVectorsReader flatReader) throws IOException { this.resources = resources; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java similarity index 95% rename from src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 547a875d..69d66db0 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -15,11 +15,11 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_INDEX_EXT; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.CUVS_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.GPUVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -67,12 +67,13 @@ * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into * GPU */ -public class GPUVectorsWriter extends KnnVectorsWriter { +public class CuVS2510GPUVectorsWriter extends KnnVectorsWriter { - private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(GPUVectorsWriter.class); + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(CuVS2510GPUVectorsWriter.class); @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(GPUVectorsWriter.class.getName()); + private static final Logger log = Logger.getLogger(CuVS2510GPUVectorsWriter.class.getName()); /** The name of the CUVS component for the info-stream * */ private static final String CUVS_COMPONENT = "CUVS"; @@ -130,7 +131,7 @@ public boolean hnsw() { } } - public GPUVectorsWriter( + public CuVS2510GPUVectorsWriter( SegmentWriteState state, int cuvsWriterThreads, int intGraphDegree, @@ -425,7 +426,7 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw // Access the CAGRA index for this field from the reader if (knnReader != null) { - if (knnReader instanceof GPUVectorsReader cvr) { + if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { if (cvr != null) { totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); @@ -489,7 +490,7 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws /** * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. */ - private CagraIndex getCagraIndexFromReader(GPUVectorsReader reader, String fieldName) { + private CagraIndex getCagraIndexFromReader(CuVS2510GPUVectorsReader reader, String fieldName) { try { IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); FieldInfos fieldInfos = reader.getFieldInfos(); diff --git a/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java similarity index 80% rename from src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index 32627141..e54271b8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/HNSWSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -23,28 +23,28 @@ import org.apache.lucene.codecs.lucene101.Lucene101Codec; /** CuVS based codec for GPU based vector search */ -public class HNSWSearchCodec extends FilterCodec { +public class Lucene101AcceleratedHNSWCodec extends FilterCodec { - private static final Logger log = Logger.getLogger(HNSWSearchCodec.class.getName()); + private static final Logger log = Logger.getLogger(Lucene101AcceleratedHNSWCodec.class.getName()); private static final int DEFAULT_CUVS_WRITER_THREADS = 1; private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; private static final int DEFAULT_GRAPH_DEGREE = 64; private static final int DEFAULT_HNSW_LAYERS = 1; - private static final String NAME = "HNSWSearchCodec"; + private static final String NAME = "Lucene101AcceleratedHNSWCodec"; private KnnVectorsFormat format; - public HNSWSearchCodec() { + public Lucene101AcceleratedHNSWCodec() { this(NAME, new Lucene101Codec()); } - public HNSWSearchCodec(String name, Codec delegate) { + public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { super(name, delegate); initializeFormatDefaultValues(); } - public HNSWSearchCodec( + public Lucene101AcceleratedHNSWCodec( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { this(NAME, new Lucene101Codec()); initializeFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); @@ -61,7 +61,9 @@ private void initializeFormatDefaultValues() { private void initializeFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { try { - format = new HNSWVectorsFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); + format = + new Lucene99AcceleratedHNSWVectorsFormat( + cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); setKnnFormat(format); } catch (LibraryException ex) { log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); diff --git a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java similarity index 91% rename from src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java rename to src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 07b8a412..c43b2f12 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -29,9 +29,10 @@ import org.apache.lucene.index.SegmentWriteState; /** CuVS based KnnVectorsFormat for GPU acceleration */ -public class HNSWVectorsFormat extends KnnVectorsFormat { +public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { - private static final Logger log = Logger.getLogger(HNSWVectorsFormat.class.getName()); + private static final Logger log = + Logger.getLogger(Lucene99AcceleratedHNSWVectorsFormat.class.getName()); static final int DEFAULT_WRITER_THREADS = 32; static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; @@ -60,7 +61,7 @@ public class HNSWVectorsFormat extends KnnVectorsFormat { * * @throws LibraryException if the native library fails to load */ - public HNSWVectorsFormat() { + public Lucene99AcceleratedHNSWVectorsFormat() { this( DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, @@ -73,7 +74,7 @@ public HNSWVectorsFormat() { * * @throws LibraryException if the native library fails to load */ - public HNSWVectorsFormat( + public Lucene99AcceleratedHNSWVectorsFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { super("CuVSVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; @@ -109,10 +110,11 @@ private static void checkSupported() { } @Override - public HNSWVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + public Lucene99AcceleratedHNSWVectorsWriter fieldsWriter(SegmentWriteState state) + throws IOException { checkSupported(); var flatWriter = flatVectorsFormat.fieldsWriter(state); - return new HNSWVectorsWriter( + return new Lucene99AcceleratedHNSWVectorsWriter( state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java similarity index 96% rename from src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 456b81ee..fd2af2bd 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/HNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -15,10 +15,10 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_INDEX_EXT; -import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.HNSWVectorsFormat.HNSW_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -72,12 +72,14 @@ * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into * GPU */ -public class HNSWVectorsWriter extends KnnVectorsWriter { +public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { - private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(HNSWVectorsWriter.class); + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(Lucene99AcceleratedHNSWVectorsWriter.class); @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(HNSWVectorsWriter.class.getName()); + private static final Logger log = + Logger.getLogger(Lucene99AcceleratedHNSWVectorsWriter.class.getName()); /** The name of the CUVS component for the info-stream * */ private static final String CUVS_COMPONENT = "CUVS"; @@ -96,7 +98,7 @@ public class HNSWVectorsWriter extends KnnVectorsWriter { private String vemFileName; private String vexFileName; - public HNSWVectorsWriter( + public Lucene99AcceleratedHNSWVectorsWriter( SegmentWriteState state, int cuvsWriterThreads, int intGraphDegree, @@ -516,7 +518,7 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw // Access the CAGRA index for this field from the reader if (knnReader != null) { - if (knnReader instanceof GPUVectorsReader cvr) { + if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { if (cvr != null) { totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); @@ -579,7 +581,7 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws /** * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. */ - private CagraIndex getCagraIndexFromReader(GPUVectorsReader reader, String fieldName) { + private CagraIndex getCagraIndexFromReader(CuVS2510GPUVectorsReader reader, String fieldName) { try { IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); FieldInfos fieldInfos = reader.getFieldInfos(); diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index a025fbcc..58a370c8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -57,16 +57,17 @@ static long nanosToMillis(long nanos) { static CuVSResources cuVSResourcesOrNull() { try { - GPUVectorsFormat.resources = CuVSResources.create(); - return GPUVectorsFormat.resources; + CuVS2510GPUVectorsFormat.resources = CuVSResources.create(); + return CuVS2510GPUVectorsFormat.resources; } catch (UnsupportedOperationException uoe) { - GPUVectorsFormat.log.warning( + CuVS2510GPUVectorsFormat.log.warning( "cuvs is not supported on this platform or java version: " + uoe.getMessage()); } catch (Throwable t) { if (t instanceof ExceptionInInitializerError ex) { t = ex.getCause(); } - GPUVectorsFormat.log.warning("Exception occurred during creation of cuvs resources. " + t); + CuVS2510GPUVectorsFormat.log.warning( + "Exception occurred during creation of cuvs resources. " + t); } return null; } diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec index 0f25e335..9865a33d 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -1,2 +1,2 @@ -com.nvidia.cuvs.lucene.HNSWSearchCodec -com.nvidia.cuvs.lucene.GPUSearchCodec +com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec +com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 8c50f8cb..747a1575 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -15,5 +15,5 @@ org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat -com.nvidia.cuvs.lucene.GPUVectorsFormat -com.nvidia.cuvs.lucene.HNSWVectorsFormat +com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat +com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 1fb06296..dcf536df 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -61,7 +61,7 @@ public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); random = new Random(); // Fixed seed so that we can validate against the same result. random.setSeed(222); @@ -71,7 +71,7 @@ public static void beforeClass() throws Exception { @Test public void testCagraToHnswSerializationAndSearch() throws IOException { - Codec codec = new HNSWSearchCodec(); + Codec codec = new Lucene101AcceleratedHNSWCodec(); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); int numDocs = 2000; // random.nextInt(100, 1000); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index 14c1128d..31c6c382 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -56,12 +56,12 @@ public class TestCuVSDeletedDocuments extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); private static Random random; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); random = random(); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index 2054d9e9..c9a05277 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -53,7 +53,7 @@ public class TestCuVSGaps extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); static IndexSearcher searcher; static IndexReader reader; static Directory directory; @@ -70,7 +70,7 @@ public class TestCuVSGaps extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); random = random(); @@ -120,7 +120,7 @@ public static void afterClass() throws Exception { @Test public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; @@ -153,7 +153,7 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { @Test public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 1e83f66f..2988b61d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -56,7 +56,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); static IndexSearcher searcher; static IndexReader reader; static Directory directory; @@ -69,7 +69,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); RandomIndexWriter writer = @@ -184,7 +184,7 @@ private static List> generateExpectedResults( @Test public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuvs not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); Random random = random(); int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 3299ae22..ef7dc505 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -41,12 +41,12 @@ public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { @BeforeClass public static void beforeClass() { - assumeTrue("cuvs is not supported", GPUVectorsFormat.supported()); + assumeTrue("cuvs is not supported", CuVS2510GPUVectorsFormat.supported()); } @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat()); + return TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index a65209eb..0ad4861d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -17,7 +17,7 @@ import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; -import com.nvidia.cuvs.lucene.GPUVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -70,7 +70,7 @@ public class TestMerge extends LuceneTestCase { @BeforeClass public static void beforeClass() { - assumeTrue("cuVS is not supported", GPUVectorsFormat.supported()); + assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported()); } private Directory directory; @@ -128,7 +128,7 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -252,7 +252,7 @@ public void testMergeWithIndexSorting() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setIndexSort(indexSort) // This automatically enables sorting during merges .setMergePolicy(mergePolicy) .setMaxBufferedDocs(maxBufferedDocs) @@ -458,7 +458,7 @@ public void testMergeWithMissingVectors() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -597,7 +597,7 @@ public void testMergeWithDeletions() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -730,8 +730,8 @@ public void testMergeBruteForceIndex() throws IOException { + vectorProbability); // Configure with brute force index type - GPUVectorsFormat bruteForceFormat = - new GPUVectorsFormat( + CuVS2510GPUVectorsFormat bruteForceFormat = + new CuVS2510GPUVectorsFormat( 32, // writer threads 128, // intermediate graph degree 64, // graph degree @@ -882,8 +882,8 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { + vectorProbability); // Configure with CAGRA + brute force combined index type - GPUVectorsFormat combinedFormat = - new GPUVectorsFormat( + CuVS2510GPUVectorsFormat combinedFormat = + new CuVS2510GPUVectorsFormat( 32, // writer threads 128, // intermediate graph degree 64, // graph degree @@ -1058,7 +1058,7 @@ public void testLargeScaleMerge() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new GPUVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); From b9937fa5e4fa78b8ae6aeab294349055d441c78b Mon Sep 17 00:00:00 2001 From: punAhuja Date: Thu, 28 Aug 2025 18:59:33 +0530 Subject: [PATCH 06/21] Rebase with latest cuvs-java --- pom.xml | 2 +- .../cuvs/lucene/FilterCuVSProvider.java | 21 +++++++++++++++++-- .../java/com/nvidia/cuvs/lucene/Utils.java | 9 +++----- ...TestCagraToHnswSerializationAndSearch.java | 8 ++++--- 4 files changed, 28 insertions(+), 12 deletions(-) diff --git a/pom.xml b/pom.xml index cd9ee5bf..818988e0 100644 --- a/pom.xml +++ b/pom.xml @@ -68,7 +68,7 @@ com.nvidia.cuvs cuvs-java - 25.10.0-d0a83-SNAPSHOT + 25.10.0-e3a2d-SNAPSHOT diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index 5c3f7d14..05acbeef 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -68,8 +68,25 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] arg0) throws Throwable { } @Override - public Builder newMatrixBuilder(int size, int dimensions, DataType dataType) { - return delegate.newMatrixBuilder(size, dimensions, dataType); + public com.nvidia.cuvs.GPUInfoProvider gpuInfoProvider() { + return delegate.gpuInfoProvider(); + } + + @Override + public Builder newHostMatrixBuilder(long rows, long cols, DataType dataType) { + return delegate.newHostMatrixBuilder(rows, cols, dataType); + } + + @Override + public Builder newDeviceMatrixBuilder( + CuVSResources resources, long rows, long cols, DataType dataType) { + return delegate.newDeviceMatrixBuilder(resources, rows, cols, dataType); + } + + @Override + public Builder newDeviceMatrixBuilder( + CuVSResources resources, long rows, long cols, int maxRows, int maxCols, DataType dataType) { + return delegate.newDeviceMatrixBuilder(resources, rows, cols, maxRows, maxCols, dataType); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 58a370c8..15ea3b0a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -43,12 +43,9 @@ static void handleThrowable(Throwable t) throws IOException { * @return an instance of {@link CuVSMatrix} */ static CuVSMatrix createFloatMatrix(List data, int dimensions) { - CuVSMatrix.Builder builder = - CuVSMatrix.builder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); - for (float[] vector : data) { - builder.addVector(vector); - } - return builder.build(); + // Convert List to float[][] for the ofArray method + float[][] vectors = data.toArray(new float[0][]); + return CuVSMatrix.ofArray(vectors); } static long nanosToMillis(long nanos) { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index dcf536df..2ce97ac8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -162,9 +162,11 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { @AfterClass public static void afterClass() throws Exception { - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); + if (indexDirPath != null) { + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } } } From 953d0936eb6a9595fe20af648138acbed7e56d80 Mon Sep 17 00:00:00 2001 From: ishan Date: Mon, 1 Sep 2025 22:16:47 +0530 Subject: [PATCH 07/21] Rebase with latest cuvs-java --- pom.xml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pom.xml b/pom.xml index 818988e0..e2af0c65 100644 --- a/pom.xml +++ b/pom.xml @@ -15,7 +15,7 @@ - cuvs-java + searchscale-maven SearchScale Maven https://maven.searchscale.com/snapshots @@ -68,7 +68,7 @@ com.nvidia.cuvs cuvs-java - 25.10.0-e3a2d-SNAPSHOT + 25.10.0-55985-SNAPSHOT From 5553e588c88923a13c9d9e21671edf3aaf2cace5 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Tue, 2 Sep 2025 20:33:46 -0400 Subject: [PATCH 08/21] Multi-layer hnsw bug fix --- .../Lucene99AcceleratedHNSWVectorsWriter.java | 32 +++++++++++++---- ...TestCagraToHnswSerializationAndSearch.java | 35 +++++++++---------- 2 files changed, 42 insertions(+), 25 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index fd2af2bd..e15b743e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -29,6 +29,7 @@ import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; @@ -279,8 +280,7 @@ private GPUBuiltHnswGraph createMultiLayerHnswGraph( while (layerIndex < hnswLayers && currentLayerSize > 1) { // Calculate size for next layer (1/M of current layer) - int nextLayerSize = Math.max(1, currentLayerSize / M); - + int nextLayerSize = Math.max(2, currentLayerSize / M); // Select nodes for this layer SortedSet selectedNodesSet = new TreeSet<>(); @@ -311,7 +311,7 @@ private GPUBuiltHnswGraph createMultiLayerHnswGraph( } // Build CAGRA graph for this layer - layerAdjacencies.add(buildCagraGraphForSubset(selectedVectors)); + layerAdjacencies.add(buildCagraGraphForSubset(selectedVectors, selectedNodes)); // Update for next iteration currentLayerSize = nextLayerSize; @@ -328,8 +328,8 @@ private GPUBuiltHnswGraph createMultiLayerHnswGraph( /** * Builds a CAGRA graph for a subset of vectors */ - private CuVSMatrix buildCagraGraphForSubset(float[][] vectors) throws Throwable { - + private CuVSMatrix buildCagraGraphForSubset(float[][] vectors, int[] selectedNodes) + throws Throwable { // Create CuVSMatrix from the subset vectors CuVSMatrix subsetDataset = CuVSMatrix.ofArray(vectors); @@ -341,8 +341,28 @@ private CuVSMatrix buildCagraGraphForSubset(float[][] vectors) throws Throwable // Get adjacency list from subset CAGRA index CuVSMatrix cagraGraph = subsetIndex.getGraph(); + long numNodes = cagraGraph.size(); + long degree = cagraGraph.columns(); + + // Create a re-mapped adjacency list + int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; + + for (int i = 0; i < numNodes; i++) { + RowView rv = cagraGraph.getRow(i); + for (int j = 0; j < degree && j < rv.size(); j++) { + int subsetIndex1 = rv.getAsInt(j); + // Map subset index to original node ID + if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { + remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; + } else { + // Invalid index, use self-reference + remappedAdjacency[i][j] = selectedNodes[i]; + } + } + } + subsetIndex.close(); - return cagraGraph; + return CuVSMatrix.ofArray(remappedAdjacency); } private void writeMeta( diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 2ce97ac8..c27f7e94 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -21,9 +21,8 @@ import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; -import java.util.ArrayList; import java.util.Arrays; -import java.util.List; +import java.util.HashSet; import java.util.Random; import java.util.UUID; import java.util.logging.Logger; @@ -71,14 +70,15 @@ public static void beforeClass() throws Exception { @Test public void testCagraToHnswSerializationAndSearch() throws IOException { - Codec codec = new Lucene101AcceleratedHNSWCodec(); + Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); - + // TODO: handle random related issues. int numDocs = 2000; // random.nextInt(100, 1000); int dimension = 32; // random.nextInt(8, 1024); - int topK = 100; // random.nextInt(5, 60); - final int COMMIT_FREQ = Math.min(numDocs, random.nextInt(100, 1000)); + int topK = 5; // random.nextInt(5, 60); + final int COMMIT_FREQ = 2000; // Math.min(numDocs, random.nextInt(100, 1000)); int count = COMMIT_FREQ; + final String ID_FIELD = "id"; final String VECTOR_FIELD = "knn1"; float[][] dataset = generateDataset(random, numDocs, dimension); @@ -87,7 +87,7 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { for (int i = 0; i < numDocs; i++) { Document document = new Document(); - document.add(new StringField("id", Integer.toString(i), Field.Store.YES)); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); indexWriter.addDocument(document); count -= 1; @@ -106,7 +106,7 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { int vectorCount = 0; for (LeafReaderContext leafReaderContext : reader.leaves()) { LeafReader leafReader = leafReaderContext.reader(); - FloatVectorValues knnValues = leafReader.getFloatVectorValues("knn1"); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); assertNotNull(knnValues); log.info( VECTOR_FIELD @@ -120,7 +120,7 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { } assertTrue("Dataset size mismatch", vectorCount == numDocs); - log.info("\n2. Testing vector search queries..."); + log.info("\nTesting vector search queries..."); IndexSearcher searcher = new IndexSearcher(reader); float[] queryVector = generateDataset(random, 1, dimension)[0]; @@ -129,30 +129,27 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); TopDocs results = searcher.search(query, topK); - log.info("\nknn1 search results (" + results.totalHits + " total hits):"); - int[] expected = {1803, 1869, 554, 1824, 1982, 1302, 320, 351, 707, 549}; - List res = new ArrayList(); + log.info("\nSearch results (" + results.totalHits + " total hits):"); + Integer[] expected = {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}; + HashSet expectedIds = new HashSet(Arrays.asList(expected)); for (int i = 0; i < results.scoreDocs.length; i++) { ScoreDoc scoreDoc = results.scoreDocs[i]; Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); log.info( " Rank " + (i + 1) + ": doc " + scoreDoc.doc + " (id=" - + doc.get("id") + + id + "), score=" + scoreDoc.score); - res.add(Integer.valueOf(doc.get("id"))); + assertTrue( + "Id " + id + " not found in expectedIds", expectedIds.contains(Integer.valueOf(id))); } - assertTrue("TopK results not returned", results.scoreDocs.length == topK); - // TODO: make this test a bit more meaningful like checking the quality of search results. - for (int i : expected) { - assertTrue("Expected doc id is missing:" + i, res.contains(i)); - } } catch (Exception e) { e.printStackTrace(); From 38e54aceee5593ad65f799ce5b1e698ef944e8af Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 3 Sep 2025 21:07:58 -0400 Subject: [PATCH 09/21] Update format names --- .../com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java | 6 +++--- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 98dbe2a9..c02ae3b3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -62,7 +62,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { final CuVS2510GPUVectorsWriter.IndexType indexType; // the index type to build, when writing /** - * Creates a CuVSVectorsFormat, with default values. + * Creates a CuVS2510GPUVectorsFormat, with default values. * * @throws LibraryException if the native library fails to load */ @@ -76,7 +76,7 @@ public CuVS2510GPUVectorsFormat() { } /** - * Creates a CuVSVectorsFormat, with the given threads, graph degree, etc. + * Creates a CuVS2510GPUVectorsFormat, with the given threads, graph degree, etc. * * @throws LibraryException if the native library fails to load */ @@ -86,7 +86,7 @@ public CuVS2510GPUVectorsFormat( int graphDegree, int hnswLayers, IndexType indexType) { - super("CuVSVectorsFormat"); + super("CuVS2510GPUVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index c43b2f12..66841f6b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -57,7 +57,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion /** - * Creates a CuVSVectorsFormat, with default values. + * Creates a Lucene99AcceleratedHNSWVectorsFormat, with default values. * * @throws LibraryException if the native library fails to load */ @@ -70,13 +70,13 @@ public Lucene99AcceleratedHNSWVectorsFormat() { } /** - * Creates a CuVSVectorsFormat, with the given threads, graph degree, etc. + * Creates a Lucene99AcceleratedHNSWVectorsFormat, with the given threads, graph degree, etc. * * @throws LibraryException if the native library fails to load */ public Lucene99AcceleratedHNSWVectorsFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { - super("CuVSVectorsFormat"); + super("Lucene99AcceleratedHNSWVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; From f26d455379bd08d5549bf3ec83d36ded40cc2e7b Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 4 Sep 2025 22:19:46 -0400 Subject: [PATCH 10/21] Update pom.xml - remove dependencies. --- pom.xml | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/pom.xml b/pom.xml index e2af0c65..81c3d5c8 100644 --- a/pom.xml +++ b/pom.xml @@ -39,21 +39,11 @@ 10.2.0 test - - com.opencsv - opencsv - 5.3 - commons-io commons-io 2.15.1 - - com.github.fommil - jniloader - 1.1 - com.fasterxml.jackson.core jackson-databind From 1a01a6cafe4caf202e35c9f9372119f889da939b Mon Sep 17 00:00:00 2001 From: Ishan Chattopadhyaya Date: Mon, 15 Sep 2025 08:47:01 +0530 Subject: [PATCH 11/21] Remove Jackson dependency and restricting commons-io to test only --- pom.xml | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/pom.xml b/pom.xml index 81c3d5c8..9f5551db 100644 --- a/pom.xml +++ b/pom.xml @@ -42,17 +42,8 @@ commons-io commons-io - 2.15.1 - - - com.fasterxml.jackson.core - jackson-databind - 2.17.0 - - - com.fasterxml.jackson.dataformat - jackson-dataformat-csv - 2.17.0 + 2.18.0 + test com.nvidia.cuvs From 35c749de4bbc4fc849a039ae02405342cd9f7026 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Tue, 16 Sep 2025 13:20:59 -0400 Subject: [PATCH 12/21] Bump up the cuvs-java version --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 9f5551db..2cacd3f7 100644 --- a/pom.xml +++ b/pom.xml @@ -49,7 +49,7 @@ com.nvidia.cuvs cuvs-java - 25.10.0-55985-SNAPSHOT + 25.10.0-9667a-SNAPSHOT From 49d690fa8f20bae5a35409228babc81ae8d447ee Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 17 Sep 2025 23:51:51 -0400 Subject: [PATCH 13/21] Adding fallback support and a test for it --- .../lucene/Lucene101AcceleratedHNSWCodec.java | 26 +++++++-- .../Lucene99AcceleratedHNSWVectorsFormat.java | 56 ++++++++++++++----- ...TestCagraToHnswSerializationAndSearch.java | 56 ++++++++++++------- 3 files changed, 98 insertions(+), 40 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index e54271b8..facd93d8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -15,6 +15,9 @@ */ package com.nvidia.cuvs.lucene; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; + import com.nvidia.cuvs.LibraryException; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -45,9 +48,15 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { } public Lucene101AcceleratedHNSWCodec( - int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + int maxConn, + int beamWidth) { this(NAME, new Lucene101Codec()); - initializeFormat(cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); + initializeFormat( + cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, maxConn, beamWidth); } private void initializeFormatDefaultValues() { @@ -55,15 +64,22 @@ private void initializeFormatDefaultValues() { DEFAULT_CUVS_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, - DEFAULT_HNSW_LAYERS); + DEFAULT_HNSW_LAYERS, + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH); } private void initializeFormat( - int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + int maxConn, + int beamWidth) { try { format = new Lucene99AcceleratedHNSWVectorsFormat( - cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers); + cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, maxConn, beamWidth); setKnnFormat(format); } catch (LibraryException ex) { log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 66841f6b..84345f67 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -15,16 +15,22 @@ */ package com.nvidia.cuvs.lucene; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER; + import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; import java.io.IOException; import java.util.logging.Logger; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; @@ -37,24 +43,29 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { static final int DEFAULT_WRITER_THREADS = 32; static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; static final int DEFAULT_GRAPH_DEGREE = 64; - static final int HNSW_GRAPH_LAYERS = 1; + static final int DEFAULT_HNSW_GRAPH_LAYERS = 1; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; static final String HNSW_INDEX_CODEC_NAME = "Lucene99HnswVectorsFormatIndex"; static final String HNSW_INDEX_EXT = "vex"; - static CuVSResources resources = cuVSResourcesOrNull(); + // Needed to make this public for a test as mocking/env variable manuplation is complicated. + // TODO: Maybe explore a better solution later. + public static CuVSResources resources = cuVSResourcesOrNull(); /** The format for storing, reading, and merging raw vectors on disk. */ private static final FlatVectorsFormat flatVectorsFormat = new Lucene99FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE); - final int maxDimensions = 4096; - final int cuvsWriterThreads; - final int intGraphDegree; - final int graphDegree; - final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion + private final int maxDimensions = 4096; + private final int cuvsWriterThreads; + private final int intGraphDegree; + private final int graphDegree; + private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion + + private final int maxConn; + private final int beamWidth; /** * Creates a Lucene99AcceleratedHNSWVectorsFormat, with default values. @@ -66,7 +77,9 @@ public Lucene99AcceleratedHNSWVectorsFormat() { DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, - HNSW_GRAPH_LAYERS); + DEFAULT_HNSW_GRAPH_LAYERS, + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH); } /** @@ -75,15 +88,22 @@ public Lucene99AcceleratedHNSWVectorsFormat() { * @throws LibraryException if the native library fails to load */ public Lucene99AcceleratedHNSWVectorsFormat( - int cuvsWriterThreads, int intGraphDegree, int graphDegree, int hnswLayers) { + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + int maxConn, + int beamWidth) { super("Lucene99AcceleratedHNSWVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; this.hnswLayers = hnswLayers; + this.maxConn = maxConn; + this.beamWidth = beamWidth; } - private static CuVSResources cuVSResourcesOrNull() { + public static CuVSResources cuVSResourcesOrNull() { try { resources = CuVSResources.create(); return resources; @@ -110,12 +130,18 @@ private static void checkSupported() { } @Override - public Lucene99AcceleratedHNSWVectorsWriter fieldsWriter(SegmentWriteState state) - throws IOException { - checkSupported(); + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { var flatWriter = flatVectorsFormat.fieldsWriter(state); - return new Lucene99AcceleratedHNSWVectorsWriter( - state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); + if (supported()) { + return new Lucene99AcceleratedHNSWVectorsWriter( + state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); + } else { + log.warning( + "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); + // TODO: Make num merge workers configurable. + return new Lucene99HnswVectorsWriter( + state, maxConn, beamWidth, flatWriter, DEFAULT_NUM_MERGE_WORKER, null); + } } @Override diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index c27f7e94..cfe123f4 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -15,19 +15,22 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.cuVSResourcesOrNull; +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.File; import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; +import java.util.List; 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; @@ -56,21 +59,42 @@ public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); private static Random random; - private static Path indexDirPath; + private static List indexDirPath; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); random = new Random(); // Fixed seed so that we can validate against the same result. random.setSeed(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); + indexDirPath = new ArrayList(); + for (int i = 0; i < 2; i++) { + indexDirPath.add(Paths.get(UUID.randomUUID().toString())); + } } @Test public void testCagraToHnswSerializationAndSearch() throws IOException { + log.info("Test Scenario 1 - cuvs supported"); + test( + false, + new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}, + indexDirPath.get(0)); + + log.info("Test Scenario 2 - cuvs NOT supported"); + test(true, new Integer[] {885, 612, 1795, 1806, 1665}, indexDirPath.get(1)); + } + + private void test(boolean disableResources, Integer[] expected, Path indexDirPath) + throws IOException { + + Lucene101AcceleratedHNSWCodec codec = + new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + + if (disableResources) { + Lucene99AcceleratedHNSWVectorsFormat.resources = null; + } - Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); // TODO: handle random related issues. int numDocs = 2000; // random.nextInt(100, 1000); @@ -120,7 +144,7 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { } assertTrue("Dataset size mismatch", vectorCount == numDocs); - log.info("\nTesting vector search queries..."); + log.info("Testing vector search queries..."); IndexSearcher searcher = new IndexSearcher(reader); float[] queryVector = generateDataset(random, 1, dimension)[0]; @@ -129,8 +153,7 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); TopDocs results = searcher.search(query, topK); - log.info("\nSearch results (" + results.totalHits + " total hits):"); - Integer[] expected = {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}; + log.info("Search results (" + results.totalHits + " total hits):"); HashSet expectedIds = new HashSet(Arrays.asList(expected)); for (int i = 0; i < results.scoreDocs.length; i++) { @@ -159,21 +182,14 @@ public void testCagraToHnswSerializationAndSearch() throws IOException { @AfterClass public static void afterClass() throws Exception { - if (indexDirPath != null) { - File indexDirPathFile = indexDirPath.toFile(); + // Reset resources. + Lucene99AcceleratedHNSWVectorsFormat.resources = cuVSResourcesOrNull(); + // Cleanup. + for (Path p : indexDirPath) { + File indexDirPathFile = p.toFile(); if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { FileUtils.deleteDirectory(indexDirPathFile); } } } - - private static float[][] generateDataset(Random random, int datasetSize, int dimensions) { - float[][] dataset = new float[datasetSize][dimensions]; - for (int i = 0; i < datasetSize; i++) { - for (int j = 0; j < dimensions; j++) { - dataset[i][j] = random.nextFloat() * 100; - } - } - return dataset; - } } From cc2913e07276529fe8cbdd63a2efab3159b57f08 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 18 Sep 2025 13:39:11 -0400 Subject: [PATCH 14/21] Code refactoring and cleanup --- .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 28 +++++----- .../Lucene99AcceleratedHNSWVectorsFormat.java | 51 ++++++++----------- .../java/com/nvidia/cuvs/lucene/Utils.java | 12 ++--- ...TestCagraToHnswSerializationAndSearch.java | 28 ++++------ .../cuvs/lucene/TestCuVSDeletedDocuments.java | 4 +- .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 8 +-- .../TestCuVSRandomizedVectorSearch.java | 6 +-- .../cuvs/lucene/TestCuVSVectorsFormat.java | 2 +- 8 files changed, 64 insertions(+), 75 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index c02ae3b3..bccff5ad 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -15,6 +15,8 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; + import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; @@ -35,7 +37,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { // TODO: fix Lucene version in name, to the final targeted release, if any static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; - static final String CUVS_META_CODEC_EXT = "vemc"; // ""cagmf"; + static final String CUVS_META_CODEC_EXT = "vemc"; static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex"; static final String CUVS_INDEX_EXT = "vcag"; @@ -48,7 +50,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; static final int HNSW_GRAPH_LAYERS = 1; - static CuVSResources resources = Utils.cuVSResourcesOrNull(); + static CuVSResources resources = cuVSResourcesOrNull(); /** The format for storing, reading, and merging raw vectors on disk. */ private static final FlatVectorsFormat flatVectorsFormat = @@ -94,17 +96,6 @@ public CuVS2510GPUVectorsFormat( this.indexType = indexType; } - /** Tells whether the platform supports cuvs. */ - public static boolean supported() { - return resources != null; - } - - private static void checkSupported() { - if (!supported()) { - throw new UnsupportedOperationException(); - } - } - @Override public CuVS2510GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { checkSupported(); @@ -142,4 +133,15 @@ public String toString() { sb.append(")"); return sb.toString(); } + + /** Tells whether the platform supports cuVS. */ + public static boolean supported() { + return resources != null; + } + + public static void checkSupported() { + if (!supported()) { + throw new UnsupportedOperationException(); + } + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 84345f67..3a18aaef 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -15,6 +15,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER; @@ -50,9 +51,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { static final String HNSW_INDEX_CODEC_NAME = "Lucene99HnswVectorsFormatIndex"; static final String HNSW_INDEX_EXT = "vex"; - // Needed to make this public for a test as mocking/env variable manuplation is complicated. - // TODO: Maybe explore a better solution later. - public static CuVSResources resources = cuVSResourcesOrNull(); + private static CuVSResources resources = cuVSResourcesOrNull(); /** The format for storing, reading, and merging raw vectors on disk. */ private static final FlatVectorsFormat flatVectorsFormat = @@ -103,36 +102,11 @@ public Lucene99AcceleratedHNSWVectorsFormat( this.beamWidth = beamWidth; } - public static CuVSResources cuVSResourcesOrNull() { - try { - resources = CuVSResources.create(); - return resources; - } catch (UnsupportedOperationException uoe) { - log.warning("cuvs is not supported on this platform or java version: " + uoe.getMessage()); - } catch (Throwable t) { - if (t instanceof ExceptionInInitializerError ex) { - t = ex.getCause(); - } - log.warning("Exception occurred during creation of cuvs resources. " + t); - } - return null; - } - - /** Tells whether the platform supports cuvs. */ - public static boolean supported() { - return resources != null; - } - - private static void checkSupported() { - if (!supported()) { - throw new UnsupportedOperationException(); - } - } - @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { var flatWriter = flatVectorsFormat.fieldsWriter(state); if (supported()) { + log.info("cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter( state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); } else { @@ -165,4 +139,23 @@ public String toString() { sb.append(")"); return sb.toString(); } + + public static CuVSResources getResources() { + return resources; + } + + public static void setResources(CuVSResources resources) { + Lucene99AcceleratedHNSWVectorsFormat.resources = resources; + } + + /** Tells whether the platform supports cuVS. */ + public static boolean supported() { + return resources != null; + } + + public static void checkSupported() { + if (!supported()) { + throw new UnsupportedOperationException(); + } + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 15ea3b0a..3da50061 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -20,9 +20,12 @@ import java.io.IOException; import java.time.Duration; import java.util.List; +import java.util.logging.Logger; public class Utils { + static final Logger log = Logger.getLogger(Utils.class.getName()); + static void handleThrowable(Throwable t) throws IOException { switch (t) { case IOException ioe -> throw ioe; @@ -54,17 +57,14 @@ static long nanosToMillis(long nanos) { static CuVSResources cuVSResourcesOrNull() { try { - CuVS2510GPUVectorsFormat.resources = CuVSResources.create(); - return CuVS2510GPUVectorsFormat.resources; + return CuVSResources.create(); } catch (UnsupportedOperationException uoe) { - CuVS2510GPUVectorsFormat.log.warning( - "cuvs is not supported on this platform or java version: " + uoe.getMessage()); + log.warning("cuVS is not supported on this platform or java version: " + uoe.getMessage()); } catch (Throwable t) { if (t instanceof ExceptionInInitializerError ex) { t = ex.getCause(); } - CuVS2510GPUVectorsFormat.log.warning( - "Exception occurred during creation of cuvs resources. " + t); + log.warning("Exception occurred during creation of cuVS resources. " + t); } return null; } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index cfe123f4..2fbfb698 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -15,8 +15,8 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.cuVSResourcesOrNull; import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.File; @@ -75,26 +75,20 @@ public static void beforeClass() throws Exception { @Test public void testCagraToHnswSerializationAndSearch() throws IOException { - log.info("Test Scenario 1 - cuvs supported"); - test( - false, - new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}, - indexDirPath.get(0)); - - log.info("Test Scenario 2 - cuvs NOT supported"); - test(true, new Integer[] {885, 612, 1795, 1806, 1665}, indexDirPath.get(1)); + log.info("Test Scenario 1 - cuVS supported"); + test(new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}, indexDirPath.get(0)); + + log.info("Test Scenario 2 - cuVS NOT supported"); + // Set resources to null to simulate that cuVS is not supported. + Lucene99AcceleratedHNSWVectorsFormat.setResources(null); + test(new Integer[] {885, 612, 1795, 1806, 1665}, indexDirPath.get(1)); } - private void test(boolean disableResources, Integer[] expected, Path indexDirPath) - throws IOException { + private void test(Integer[] expected, Path indexDirPath) throws IOException { Lucene101AcceleratedHNSWCodec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); - if (disableResources) { - Lucene99AcceleratedHNSWVectorsFormat.resources = null; - } - IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); // TODO: handle random related issues. int numDocs = 2000; // random.nextInt(100, 1000); @@ -103,7 +97,7 @@ private void test(boolean disableResources, Integer[] expected, Path indexDirPat final int COMMIT_FREQ = 2000; // Math.min(numDocs, random.nextInt(100, 1000)); int count = COMMIT_FREQ; final String ID_FIELD = "id"; - final String VECTOR_FIELD = "knn1"; + final String VECTOR_FIELD = "vector_field"; float[][] dataset = generateDataset(random, numDocs, dimension); // Indexing @@ -183,7 +177,7 @@ private void test(boolean disableResources, Integer[] expected, Path indexDirPat @AfterClass public static void afterClass() throws Exception { // Reset resources. - Lucene99AcceleratedHNSWVectorsFormat.resources = cuVSResourcesOrNull(); + Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull()); // Cleanup. for (Path p : indexDirPath) { File indexDirPathFile = p.toFile(); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index 31c6c382..1774587f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -51,7 +51,7 @@ import org.junit.BeforeClass; import org.junit.Test; -@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +@SuppressSysoutChecks(bugUrl = "prints info from within cuVS") public class TestCuVSDeletedDocuments extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); @@ -61,7 +61,7 @@ public class TestCuVSDeletedDocuments extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); random = random(); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index c9a05277..bba34bfa 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -48,7 +48,7 @@ import org.junit.BeforeClass; import org.junit.Test; -@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +@SuppressSysoutChecks(bugUrl = "prints info from within cuVS") public class TestCuVSGaps extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); @@ -70,7 +70,7 @@ public class TestCuVSGaps extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); random = random(); @@ -120,7 +120,7 @@ public static void afterClass() throws Exception { @Test public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; @@ -153,7 +153,7 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { @Test public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 2988b61d..add0fe20 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -51,7 +51,7 @@ import org.junit.BeforeClass; import org.junit.Test; -@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +@SuppressSysoutChecks(bugUrl = "prints info from within cuVS") public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); @@ -69,7 +69,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); RandomIndexWriter writer = @@ -184,7 +184,7 @@ private static List> generateExpectedResults( @Test public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); Random random = random(); int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index ef7dc505..cf78b051 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -41,7 +41,7 @@ public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { @BeforeClass public static void beforeClass() { - assumeTrue("cuvs is not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported()); } @Override From 333182fb42417b19c76adedf30dc4c4267b1e84d Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 18 Sep 2025 23:39:29 -0400 Subject: [PATCH 15/21] Add a test and update the other --- ...TestCagraToHnswSerializationAndSearch.java | 58 ++---- ...ializationAndSearchWithFallbackWriter.java | 173 ++++++++++++++++++ 2 files changed, 191 insertions(+), 40 deletions(-) create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 2fbfb698..4f5f7aba 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -16,21 +16,19 @@ package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.File; import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; -import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; -import java.util.List; 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; @@ -56,48 +54,32 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { - protected static Logger log = + private static Logger log = Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); private static Random random; - private static List indexDirPath; + private static Path indexDirPath; @BeforeClass public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - random = new Random(); // Fixed seed so that we can validate against the same result. - random.setSeed(222); - indexDirPath = new ArrayList(); - for (int i = 0; i < 2; i++) { - indexDirPath.add(Paths.get(UUID.randomUUID().toString())); - } + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); } @Test public void testCagraToHnswSerializationAndSearch() throws IOException { - log.info("Test Scenario 1 - cuVS supported"); - test(new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}, indexDirPath.get(0)); - - log.info("Test Scenario 2 - cuVS NOT supported"); - // Set resources to null to simulate that cuVS is not supported. - Lucene99AcceleratedHNSWVectorsFormat.setResources(null); - test(new Integer[] {885, 612, 1795, 1806, 1665}, indexDirPath.get(1)); - } - - private void test(Integer[] expected, Path indexDirPath) throws IOException { - - Lucene101AcceleratedHNSWCodec codec = - new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); - + Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); - // TODO: handle random related issues. - int numDocs = 2000; // random.nextInt(100, 1000); - int dimension = 32; // random.nextInt(8, 1024); - int topK = 5; // random.nextInt(5, 60); - final int COMMIT_FREQ = 2000; // Math.min(numDocs, random.nextInt(100, 1000)); - int count = COMMIT_FREQ; + + final int COMMIT_FREQ = 2000; final String ID_FIELD = "id"; final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; float[][] dataset = generateDataset(random, numDocs, dimension); // Indexing @@ -148,6 +130,7 @@ private void test(Integer[] expected, Path indexDirPath) throws IOException { TopDocs results = searcher.search(query, topK); log.info("Search results (" + results.totalHits + " total hits):"); + Integer[] expected = new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}; HashSet expectedIds = new HashSet(Arrays.asList(expected)); for (int i = 0; i < results.scoreDocs.length; i++) { @@ -164,7 +147,7 @@ private void test(Integer[] expected, Path indexDirPath) throws IOException { + "), score=" + scoreDoc.score); assertTrue( - "Id " + id + " not found in expectedIds", expectedIds.contains(Integer.valueOf(id))); + "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); } assertTrue("TopK results not returned", results.scoreDocs.length == topK); @@ -176,14 +159,9 @@ private void test(Integer[] expected, Path indexDirPath) throws IOException { @AfterClass public static void afterClass() throws Exception { - // Reset resources. - Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull()); - // Cleanup. - for (Path p : indexDirPath) { - File indexDirPathFile = p.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); } } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java new file mode 100644 index 00000000..792aca96 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -0,0 +1,173 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashSet; +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.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +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; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraToHnswSerializationAndSearchWithFallbackWriter extends LuceneTestCase { + + private static Logger log = + Logger.getLogger(TestCagraToHnswSerializationAndSearchWithFallbackWriter.class.getName()); + + private static Random random; + private static Path indexDirPath; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); + // Set resources to null to simulate that cuVS is not supported. + Lucene99AcceleratedHNSWVectorsFormat.setResources(null); + // Fixed seed so that we can validate against the same result. + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @Test + public void testCagraToHnswSerializationAndSearchWithFallbackWriter() throws IOException { + Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, 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)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath)) { + try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.info("Successfully opened index"); + + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(knnValues); + log.info( + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + } + assertTrue("Dataset size mismatch", vectorCount == numDocs); + + log.info("Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.info("Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.info("Search results (" + results.totalHits + " total hits):"); + Integer[] expected = new Integer[] {1869, 1411, 1497, 351, 554}; + HashSet expectedIds = new HashSet(Arrays.asList(expected)); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.info( + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue( + "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + + } catch (Exception e) { + e.printStackTrace(); + } + } + } + + @AfterClass + public static void afterClass() throws Exception { + // Reset resources for other tests to work + Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull()); + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} From 0a4f7e0f9f4a9125fb31c299a6927528ad56f618 Mon Sep 17 00:00:00 2001 From: Ishan Chattopadhyaya Date: Thu, 25 Sep 2025 02:25:08 +0530 Subject: [PATCH 16/21] Latest cuvs-java changes require cudart to be loaded manually --- pom.xml | 2 +- src/main/java/com/nvidia/cuvs/lucene/Utils.java | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 2cacd3f7..69f3d94a 100644 --- a/pom.xml +++ b/pom.xml @@ -49,7 +49,7 @@ com.nvidia.cuvs cuvs-java - 25.10.0-9667a-SNAPSHOT + 25.10.0 diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 3da50061..dfb9aa0c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -56,6 +56,11 @@ static long nanosToMillis(long nanos) { } static CuVSResources cuVSResourcesOrNull() { + try { + System.loadLibrary("cudart"); + } catch (UnsatisfiedLinkError e) { + log.warning("Could not load CUDA runtime library: " + e.getMessage()); + } try { return CuVSResources.create(); } catch (UnsupportedOperationException uoe) { From b51d76e1128b6717cc99dc61fdbf4c6246946e6f Mon Sep 17 00:00:00 2001 From: Ishan Chattopadhyaya Date: Thu, 25 Sep 2025 04:30:30 +0530 Subject: [PATCH 17/21] Using latest cuvs-java --- build.sh | 2 +- pom.xml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) mode change 100644 => 100755 build.sh diff --git a/build.sh b/build.sh old mode 100644 new mode 100755 index ca83606b..1cf920cd --- a/build.sh +++ b/build.sh @@ -23,4 +23,4 @@ fi mvn verify "${MAVEN_VERIFY_ARGS[@]}" \ && mvn install:install-file -Dfile=./target/cuvs-lucene-$VERSION.jar -DgroupId=$GROUP_ID -DartifactId=cuvs-lucene -Dversion=$VERSION -Dpackaging=jar \ && cp pom.xml ./target/ -- \ No newline at end of file + diff --git a/pom.xml b/pom.xml index 69f3d94a..8d4c5558 100644 --- a/pom.xml +++ b/pom.xml @@ -4,7 +4,7 @@ 4.0.0 com.nvidia.cuvs.lucene cuvs-lucene - 0.0.1-SNAPSHOT + 25.10.0 cuvs-lucene jar From fed1efb272ebeb7c110749efcde137b1d1dd72dd Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 25 Sep 2025 12:11:08 -0400 Subject: [PATCH 18/21] Pin the libcuvs version --- dependencies.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dependencies.yaml b/dependencies.yaml index ee061354..9d065574 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -67,7 +67,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev - - libcuvs + - libcuvs==25.10.0 java: common: - output_types: conda From dc7cc60f51de1ce9bd99d68ee945d29427896e9d Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 25 Sep 2025 12:13:50 -0400 Subject: [PATCH 19/21] Update env yaml files --- conda/environments/all_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-129_arch-x86_64.yaml | 2 +- conda/environments/all_cuda-130_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-130_arch-x86_64.yaml | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 0b310207..e6b390ab 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.0 - maven - openjdk=22.* name: all_cuda-129_arch-aarch64 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index a8f0444b..f23df196 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.0 - maven - openjdk=22.* name: all_cuda-129_arch-x86_64 diff --git a/conda/environments/all_cuda-130_arch-aarch64.yaml b/conda/environments/all_cuda-130_arch-aarch64.yaml index 50e93bd6..50bdf095 100644 --- a/conda/environments/all_cuda-130_arch-aarch64.yaml +++ b/conda/environments/all_cuda-130_arch-aarch64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.0 - maven - openjdk=22.* name: all_cuda-130_arch-aarch64 diff --git a/conda/environments/all_cuda-130_arch-x86_64.yaml b/conda/environments/all_cuda-130_arch-x86_64.yaml index c1e7ab3a..f29dbf58 100644 --- a/conda/environments/all_cuda-130_arch-x86_64.yaml +++ b/conda/environments/all_cuda-130_arch-x86_64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.0 - maven - openjdk=22.* name: all_cuda-130_arch-x86_64 From a78dd2942f849ee2c38c76e3e273fc3ff496a981 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 25 Sep 2025 12:55:49 -0400 Subject: [PATCH 20/21] Pin libcuvs version with variable version --- conda/environments/all_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-129_arch-x86_64.yaml | 2 +- conda/environments/all_cuda-130_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-130_arch-x86_64.yaml | 2 +- dependencies.yaml | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index e6b390ab..05ec85b6 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.10.0 +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-129_arch-aarch64 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index f23df196..f0a25380 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.10.0 +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-129_arch-x86_64 diff --git a/conda/environments/all_cuda-130_arch-aarch64.yaml b/conda/environments/all_cuda-130_arch-aarch64.yaml index 50bdf095..6638808c 100644 --- a/conda/environments/all_cuda-130_arch-aarch64.yaml +++ b/conda/environments/all_cuda-130_arch-aarch64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.10.0 +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-130_arch-aarch64 diff --git a/conda/environments/all_cuda-130_arch-x86_64.yaml b/conda/environments/all_cuda-130_arch-x86_64.yaml index f29dbf58..bb8a746f 100644 --- a/conda/environments/all_cuda-130_arch-x86_64.yaml +++ b/conda/environments/all_cuda-130_arch-x86_64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.10.0 +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-130_arch-x86_64 diff --git a/dependencies.yaml b/dependencies.yaml index 9d065574..b9eb3fcf 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -67,7 +67,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev - - libcuvs==25.10.0 + - libcuvs==25.10.* java: common: - output_types: conda From 8ff4c40e8ef98a421983b9b9c25fe8e49d871824 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Thu, 25 Sep 2025 16:51:18 -0400 Subject: [PATCH 21/21] Update update-version script to update libcuvs version in yaml files and add version update marker for cuvs-java dependency --- ci/release/update-version.sh | 4 ++++ pom.xml | 4 ++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index c7ea7a1b..4a6bf417 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -36,3 +36,7 @@ sed_runner "s/VERSION=\".*\"/VERSION=\"${NEXT_FULL_JAVA_TAG}\"/g" build.sh sed_runner "/.*/s//${NEXT_FULL_JAVA_TAG}<\/version>/g" pom.xml sed_runner "s| CuVS [[:digit:]]\{2\}\.[[:digit:]]\{2\} | CuVS ${NEXT_SHORT_TAG} |g" README.md + +for FILE in dependencies.yaml conda/environments/*.yaml; do + sed_runner "s/libcuvs==.*/libcuvs==${NEXT_SHORT_TAG}.*/g" "${FILE}" +done diff --git a/pom.xml b/pom.xml index aff25cce..dfd94975 100644 --- a/pom.xml +++ b/pom.xml @@ -49,8 +49,8 @@ com.nvidia.cuvs cuvs-java - - 25.10.0 + + 25.10.0