diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java similarity index 64% rename from src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java rename to src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java index 96be6c01..e0895aa5 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java @@ -4,7 +4,11 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -12,7 +16,7 @@ import java.nio.file.Path; import java.nio.file.Paths; import java.util.Arrays; -import java.util.HashSet; +import java.util.List; import java.util.Random; import java.util.UUID; import java.util.logging.Level; @@ -35,6 +39,7 @@ import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.junit.After; @@ -42,93 +47,88 @@ import org.junit.Test; @SuppressSysoutChecks(bugUrl = "") -public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { +public class TestAcceleratedHNSW extends LuceneTestCase { - private static Logger log = - Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); + private static final Logger log = Logger.getLogger(TestAcceleratedHNSW.class.getName()); private static Random random; private static Path indexDirPath; + private static String randomID; + private static Codec codec; + private static TestDataProvider dataProvider; @Before public void beforeTest() throws Exception { assumeTrue("cuVS not supported", isSupported()); + random = new Random(); // Fixed seed so that we can validate against the same result. random = new Random(222); indexDirPath = Paths.get(UUID.randomUUID().toString()); + randomID = UUID.randomUUID().toString(); + dataProvider = new TestDataProvider(random); + codec = new Lucene101AcceleratedHNSWCodec(); } @Test - public void testCagraToHnswSerializationAndSearch() throws Exception { - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params); - 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); - + public void testAcceleratedHNSW() throws Exception { // Indexing try (Directory indexDirectory = FSDirectory.open(indexDirPath); - IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { - for (int i = 0; i < numDocs; i++) { + RandomIndexWriter indexWriter = createWriter(random, indexDirectory, codec)) { + for (int i = 0; i < dataProvider.getDatasetSize(); 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)); + document.add( + new KnnFloatVectorField(VECTOR_FIELD1, dataProvider.getDataset1()[i], EUCLIDEAN)); + document.add( + new KnnFloatVectorField(VECTOR_FIELD2, dataProvider.getDataset2()[i], EUCLIDEAN)); indexWriter.addDocument(document); - count -= 1; - if (count == 0) { - indexWriter.commit(); - count = COMMIT_FREQ; - } } + indexWriter.commit(); } // Searching try (Directory indexDirectory = FSDirectory.open(indexDirPath); DirectoryReader reader = DirectoryReader.open(indexDirectory)) { - log.log(Level.FINE, "Successfully opened index"); + + int datasetSize = dataProvider.getDatasetSize(); + int dimensions = dataProvider.getDimensions(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; int vectorCount = 0; for (LeafReaderContext leafReaderContext : reader.leaves()) { LeafReader leafReader = leafReaderContext.reader(); - FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(knnValues); log.log( Level.FINE, - VECTOR_FIELD + VECTOR_FIELD1 + " field: " + knnValues.size() + " vectors, " + knnValues.dimension() + " dimensions"); vectorCount += knnValues.size(); - assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimensions); } - assertTrue("Dataset size mismatch", vectorCount == numDocs); + assertTrue("Dataset size mismatch", vectorCount == datasetSize); log.log(Level.FINE, "Testing vector search queries..."); IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateDataset(random, 1, dimension)[0]; - log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + log.log(Level.FINER, "Query vector: " + Arrays.toString(queryVector)); - KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); TopDocs results = searcher.search(query, topK); log.log(Level.FINE, "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)); + List> expected = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}); 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); + int id = Integer.valueOf(doc.get(ID_FIELD)); log.log( Level.FINE, " Rank " @@ -139,8 +139,7 @@ public void testCagraToHnswSerializationAndSearch() throws Exception { + id + "), score=" + scoreDoc.score); - assertTrue( - "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + assertTrue("Id: " + id + " expected but not found", expected.get(0).contains(id)); } assertTrue("TopK results not returned", results.scoreDocs.length == topK); } @@ -148,27 +147,17 @@ public void testCagraToHnswSerializationAndSearch() throws Exception { @Test public void testSingleVectorIndex() throws Exception { - // Test single vector index support with dummy HNSW graph - // TODO: This test can be removed once https://github.com/rapidsai/cuvs/pull/1256 is merged - // and CAGRA natively supports single vector indexes - Codec codec = new Lucene101AcceleratedHNSWCodec(); - - final String ID_FIELD = "id"; - final String VECTOR_FIELD = "vector_field"; + try (Directory indexDirectory = newDirectory()) { - int dimension = 32; - float[] vector = generateDataset(random, 1, dimension)[0]; + int dimensions = dataProvider.getDimensions(); + float[] queryVector = dataProvider.getQueries(1)[0]; - // Index a single document with a vector - this should now work with dummy HNSW graph - try (Directory indexDirectory = newDirectory()) { IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); try (IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { Document document = new Document(); - document.add(new StringField(ID_FIELD, "0", Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + document.add(new StringField(ID_FIELD, randomID, Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD1, queryVector, EUCLIDEAN)); indexWriter.addDocument(document); - - // This should now succeed by creating a dummy HNSW graph for the single vector indexWriter.commit(); } @@ -176,18 +165,19 @@ public void testSingleVectorIndex() throws Exception { try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { assertEquals(1, reader.numDocs()); LeafReader leafReader = getOnlyLeafReader(reader); - FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(knnValues); assertEquals(1, knnValues.size()); - assertEquals(dimension, knnValues.dimension()); + assertEquals(dimensions, knnValues.dimension()); // Test search functionality IndexSearcher searcher = new IndexSearcher(reader); - KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, vector, 1); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, 1); TopDocs results = searcher.search(query, 1); assertEquals(1, results.totalHits.value()); assertEquals(1, results.scoreDocs.length); - assertEquals(0, results.scoreDocs[0].doc); + Document doc = reader.storedFields().document(results.scoreDocs[0].doc); + assertEquals(randomID, doc.get(ID_FIELD)); } } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index d164af55..b4dfd10c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -4,9 +4,13 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.TestDataProvider.CATEGORY_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.createWriterConfig; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; import java.util.ArrayList; @@ -23,9 +27,7 @@ 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; @@ -33,8 +35,6 @@ 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; @@ -45,58 +45,66 @@ @SuppressSysoutChecks(bugUrl = "") public class TestAcceleratedHNSWDeletedDocuments extends LuceneTestCase { - protected static Logger log = + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWDeletedDocuments.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + private static Codec codec; private static Random random; + private static float deletionProbability; + private static float vectorProbability; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", isSupported()); + codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); random = random(); + dataProvider = new TestDataProvider(random); + deletionProbability = random.nextFloat() * 0.4f + 0.1f; + vectorProbability = random.nextFloat() * 0.5f + 0.3f; + log.log( + Level.FINE, + "deletion probability: " + + deletionProbability + + ", vector probability: " + + vectorProbability); } @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<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); // Create index with all documents having vectors - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], 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))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log(Level.FINE, "Number of documents deleted: " + deletedDocs.size()); + // 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); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getTopK(); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify we got results @@ -104,8 +112,7 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify no deleted documents in results for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); assertFalse( "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); @@ -114,7 +121,7 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify deleted documents are truly deleted for (int deletedId : deletedDocs) { TopDocs result = - searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + searcher.search(new TermQuery(new Term(ID_FIELD, String.valueOf(deletedId))), 1); assertEquals( "Deleted document " + deletedId + " should not be found", 0, @@ -128,29 +135,25 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { 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<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create index with mixed documents - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { 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 StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); // Randomly assign categories String category = random.nextBoolean() ? "A" : "B"; - doc.add(new StringField("category", category, Field.Store.YES)); + doc.add(new StringField(CATEGORY_FIELD, category, Field.Store.YES)); // Randomly decide whether to add vectors if (random.nextFloat() < vectorProbability) { - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); } else { docsWithoutVectors.add(i); } @@ -160,40 +163,45 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Delete documents randomly for (int i = 0; i < datasetSize; i++) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log( + Level.FINE, + " Documents deleted: " + + deletedDocs.size() + + " Documents without vectors: " + + docsWithoutVectors.size()); + // 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); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, 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); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); assertFalse("Deleted document should not appear", deletedDocs.contains(id)); assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); } // Test filtered search with deletions - Query filter = new TermQuery(new Term("category", "A")); - Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); + Query filter = new TermQuery(new Term(CATEGORY_FIELD, "A")); + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; for (ScoreDoc hit : filteredHits) { Document doc = reader.storedFields().document(hit.doc); - String category = doc.get("category"); + String category = doc.get(CATEGORY_FIELD); assertEquals("Should only match category A", "A", category); - int id = Integer.parseInt(doc.get("id")); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse( "Deleted document should not appear in filtered results", deletedDocs.contains(id)); } @@ -205,37 +213,35 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti 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); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create and delete all documents - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { + // Add all documents 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], 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.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); } writer.commit(); - writer.forceMerge(1); // Force merge to apply deletions + writer.forceMerge(1); } // 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); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); TopDocs results = searcher.search(query, topK); assertEquals( @@ -250,22 +256,20 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { 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 + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - 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 + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -274,7 +278,7 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio List candidatesForDeletion = new ArrayList<>(activeDocIds); for (int docId : candidatesForDeletion) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(docId))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(docId))); activeDocIds.remove(Integer.valueOf(docId)); } } @@ -282,9 +286,8 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -294,15 +297,13 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // 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); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, 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); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); resultIds.add(id); assertTrue("Result should be from active documents", activeDocIds.contains(id)); } @@ -317,19 +318,4 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio } } } - - 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/TestAcceleratedHNSWFallback.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java new file mode 100644 index 00000000..50158d54 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java @@ -0,0 +1,155 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.setCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +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.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.index.RandomIndexWriter; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWFallback extends LuceneTestCase { + + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWFallback.class.getName()); + private static Random random; + private static Path indexDirPath; + private static Codec codec; + private static TestDataProvider dataProvider; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported so skipping these tests", isSupported()); + // Set resources to null to simulate that cuVS is not supported. + setCuVSResourcesInstance(null); + random = new Random(); + dataProvider = new TestDataProvider(random); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + codec = new Lucene101AcceleratedHNSWCodec(); + } + + @Test + public void testAcceleratedHNSWFallback() throws Exception { + + int datasetSize = dataProvider.getDatasetSize(); + int dimensions = dataProvider.getDimensions(); + float[][] dataset = dataProvider.getDataset1(); + float[][] dataset2 = dataProvider.getDataset2(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + RandomIndexWriter indexWriter = createWriter(random, indexDirectory, codec)) { + for (int i = 0; i < datasetSize; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + document.add(new KnnFloatVectorField(VECTOR_FIELD2, dataset2[i], EUCLIDEAN)); + indexWriter.addDocument(document); + } + indexWriter.commit(); + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); + assertNotNull(knnValues); + log.log( + Level.FINE, + VECTOR_FIELD1 + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimensions); + } + assertTrue("Dataset size mismatch", vectorCount == datasetSize); + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + log.log(Level.FINER, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + List> expected = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + int id = Integer.valueOf(doc.get(ID_FIELD)); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue("Id: " + id + " expected but not found", expected.get(0).contains(id)); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + } + } + + @After + public void afterTest() throws Exception { + // Reset resources for other tests to work + setCuVSResourcesInstance(cuVSResourcesOrNull()); + + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java new file mode 100644 index 00000000..04b7b4f9 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java @@ -0,0 +1,150 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.IOException; +import java.util.List; +import java.util.Random; +import java.util.logging.Level; +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.StoredFields; +import org.apache.lucene.index.Term; +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.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 = "") +public class TestAcceleratedHNSWGaps extends LuceneTestCase { + + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWGaps.class.getName()); + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported so skipping these tests", isSupported()); + codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + directory = newDirectory(); + random = random(); + dataProvider = new TestDataProvider(random); + RandomIndexWriter writer = createWriter(random, directory, codec); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + + // 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_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); + + // Only add vectors to even-numbered documents + if (i % 2 == 0) { + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @Test + public void testVectorSearchWithAlternatingDocuments() throws IOException { + + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + + // Perform search + 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) + StoredFields storedFields = reader.storedFields(); + for (ScoreDoc hit : hits) { + int id = Integer.parseInt(storedFields.document(hit.doc).get(ID_FIELD)); + assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); + log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); + } + + // Verify the results match expected top-k + List expectedIds = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}).get(0); + for (ScoreDoc hit : hits) { + int id = Integer.parseInt(storedFields.document(hit.doc).get(ID_FIELD)); + assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); + } + + log.log(Level.FINE, "Alternating documents test passed with " + hits.length + " results"); + } + + @Test + public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { + + int datasetSize = dataProvider.getDatasetSize(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + String randomEvenInRange = String.valueOf(random.nextInt(datasetSize / 2 + 1) * 2); + log.log(Level.FINE, "Randomly chosen even value is: " + randomEvenInRange); + Query filter = new TermQuery(new Term(ID_FIELD, randomEvenInRange)); + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Should only get document (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_FIELD); + assertEquals("Should only return document " + randomEvenInRange, randomEvenInRange, docId); + + log.log( + Level.FINE, + "Filtered alternating documents test passed with " + filteredHits.length + " results"); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java new file mode 100644 index 00000000..1dddd3d3 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java @@ -0,0 +1,161 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.logging.Level; +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.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.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 = "") +public class TestAcceleratedHNSWRandomizedSearch extends LuceneTestCase { + + private static final Logger log = + Logger.getLogger(TestAcceleratedHNSWRandomizedSearch.class.getName()); + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported so skipping these tests", isSupported()); + directory = newDirectory(); + random = random(); + dataProvider = new TestDataProvider(random); + codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + RandomIndexWriter writer = createWriter(random, directory, codec); + + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + float[][] dataset2 = dataProvider.getDataset2(); + + // Add documents + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); + boolean skipVector = random.nextInt(10) < 4; + if (!skipVector || datasetSize < 100) { + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, dataset2[i], EUCLIDEAN)); + } + writer.addDocument(doc); + } + writer.commit(); + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @Test + public void testVectorSearch() throws IOException { + + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + int numQueries = dataProvider.getNumQueries(); + float[][] queries = dataProvider.getQueries(numQueries); + + // Generate queries and expected results for each + List> expected = generateExpectedTopK(topK, dataset, queries); + + for (int i = 0; i < numQueries; i++) { + log.log(Level.FINE, "Running query: " + (i + 1) + " of " + numQueries); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queries[i], topK); + + // Perform search + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); + log.log(Level.FINE, "EXPECTED: " + expected.get(i)); + + // Iterate through the results and assert + for (ScoreDoc hit : hits) { + Document doc = reader.storedFields().document(hit.doc); + int docId = Integer.parseInt(doc.get(ID_FIELD)); + log.log(Level.FINE, "\t" + doc.get(ID_FIELD) + ": " + hit.score); + assertTrue("Result returned was not in topk*2: " + doc, expected.get(i).contains(docId)); + } + } + } + + @Test + public void testVectorSearchWithFilter() throws IOException { + // Find a document that has a vector by doing a search first + + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + Query unfiltered = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, 1); + ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; + + assertTrue( + "Need at least one document with vector for filtering test", unfilteredHits.length > 0); + + Document doc = reader.storedFields().document(unfilteredHits[0].doc); + String targetDocId = doc.get(ID_FIELD); + + // Create a filter that matches only the document we know has a vector + Query filter = new TermQuery(new Term(ID_FIELD, targetDocId)); + + // Test the new constructor with filter + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); + + 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_FIELD); + assertEquals("All results should match the filter", targetDocId, docId); + } + + log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java similarity index 58% rename from src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java rename to src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java index 2428fc50..e1d49d87 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java @@ -4,10 +4,14 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.util.List; +import java.util.Random; import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; @@ -27,11 +31,16 @@ import org.junit.Ignore; @SuppressSysoutChecks(bugUrl = "") -public class TestLucene99AcceleratedHNSWVectorsFormat extends BaseKnnVectorsFormatTestCase { +public class TestAcceleratedHNSWVectorsFormat extends BaseKnnVectorsFormatTestCase { + + private static TestDataProvider dataProvider; + private static Random random; @BeforeClass public static void beforeClass() { assumeTrue("cuVS is not supported", isSupported()); + random = random(); + dataProvider = new TestDataProvider(random); } @Override @@ -40,27 +49,27 @@ protected Codec getCodec() { } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { - float[][] f = new float[][] {randomVector(384), randomVector(384)}; + final int numDocs = 2; + float[][] vectors = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); - w.addDocument(doc1); - w.commit(); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); - w.addDocument(doc2); - w.flush(); - w.commit(); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors[i], EUCLIDEAN)); + w.addDocument(doc); + w.commit(); + } // sanity - verify one doc per leaf try (DirectoryReader reader = DirectoryReader.open(w)) { List subReaders = reader.leaves(); - assertEquals(2, subReaders.size()); - assertEquals(1, subReaders.get(0).reader().getFloatVectorValues("f").size()); - assertEquals(1, subReaders.get(1).reader().getFloatVectorValues("f").size()); + assertEquals(numDocs, subReaders.size()); + for (int i = 0; i < numDocs; i++) { + assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(VECTOR_FIELD1).size()); + } } // now merge to a single segment @@ -69,49 +78,56 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertEquals(numDocs, values.size()); + for (int i = 0; i < numDocs; i++) { + assertArrayEquals(vectors[i], values.vectorValue(i), 0.0f); + } } } } // Basic test for multiple vectors fields per document public void testTwoVectorFieldsPerDoc() throws Exception { - float[][] f1 = new float[][] {randomVector(384), randomVector(384)}; - float[][] f2 = new float[][] {randomVector(384), randomVector(384)}; + + final int numDocs = 2; + float[][] vectors1 = dataProvider.getVectors(numDocs); + float[][] vectors2 = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); - doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); - w.addDocument(doc1); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); - doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); - w.addDocument(doc2); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors1[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, vectors2[i], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f1"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); - - values = r.getFloatVectorValues("f2"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); - // opportunistically check boundary condition - search with a 0 topK - var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors1[i], values.vectorValue(i), 0.0f); + } + + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD2); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors2[i], values.vectorValue(i), 0.0f); + } + + // Check boundary condition - search with a 0 topK + float[] target = dataProvider.getVectors(1)[0]; + var topDocs = r.searchNearestVectors(VECTOR_FIELD1, target, 0, null, 10); assertEquals(0, topDocs.scoreDocs.length); assertEquals(0, topDocs.totalHits.value()); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e660..b8751553 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; @@ -13,7 +13,7 @@ import org.junit.Test; /** - * Tests the backward compatibility mechanism. + * Tests the backward compatibility mechanism for {@link LuceneProvider} * * @since 25.12 */ diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java deleted file mode 100644 index 04ad20a0..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ /dev/null @@ -1,173 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.setCuVSResourcesInstance; -import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import java.io.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Arrays; -import java.util.HashSet; -import java.util.Random; -import java.util.UUID; -import java.util.logging.Level; -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", isSupported()); - // Set resources to null to simulate that cuVS is not supported. - setCuVSResourcesInstance(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 Exception { - AcceleratedHNSWParams params = - new AcceleratedHNSWParams.Builder() - .withHNSWLayer(3) - .withMaxConn(16) - .withBeamWidth(100) - .build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params); - - 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.log(Level.FINE, "Successfully opened index"); - - int vectorCount = 0; - for (LeafReaderContext leafReaderContext : reader.leaves()) { - LeafReader leafReader = leafReaderContext.reader(); - FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); - assertNotNull(knnValues); - log.log( - Level.FINE, - 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.log(Level.FINE, "Testing vector search queries..."); - IndexSearcher searcher = new IndexSearcher(reader); - - float[] queryVector = generateDataset(random, 1, dimension)[0]; - log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); - - KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); - TopDocs results = searcher.search(query, topK); - - log.log(Level.FINE, "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.log( - Level.FINE, - " 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 - setCuVSResourcesInstance(cuVSResourcesOrNull()); - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } -} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java deleted file mode 100644 index d3dc866e..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java +++ /dev/null @@ -1,335 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; - -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.Level; -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 = "") -public class TestCuVSAcceleratedHNSWDeletedDocuments extends LuceneTestCase { - - protected static Logger log = - Logger.getLogger(TestCuVSAcceleratedHNSWDeletedDocuments.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - private static Random random; - - @BeforeClass - public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - 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.log(Level.FINE, "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.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); - } - - // Test filtered search with deletions - Query filter = new TermQuery(new Term("category", "A")); - Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); - 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.log( - Level.FINE, - "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/TestCuVSAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java deleted file mode 100644 index b0dae38e..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java +++ /dev/null @@ -1,191 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; - -import java.io.IOException; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.TreeMap; -import java.util.logging.Level; -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 = "") -public class TestCuVSAcceleratedHNSWGaps extends LuceneTestCase { - - protected static Logger log = Logger.getLogger(TestCuVSAcceleratedHNSWGaps.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - 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", isSupported()); - 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.log(Level.FINE, "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.log(Level.FINE, "Test finished"); - } - - @Test - public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", isSupported()); - - // 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.log(Level.FINE, "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.log(Level.FINE, "Alternating document test passed with " + hits.length + " results"); - } - - @Test - public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", isSupported()); - - // 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 KnnFloatVectorQuery("vector", queryVector, topK, filter); - 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.log( - Level.FINE, - "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/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index de868c7c..c8f0b6b8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -4,9 +4,13 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.TestDataProvider.CATEGORY_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.createWriterConfig; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; import java.util.ArrayList; @@ -23,18 +27,15 @@ 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.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.store.BaseDirectoryWrapper; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.tests.util.TestUtil; @@ -44,57 +45,69 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSDeletedDocuments extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); - - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + private static final Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); + private static Codec codec; private static Random random; + private static TestDataProvider dataProvider; + private static float deletionProbability; + private static float vectorProbability; @BeforeClass public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", isSupported()); random = random(); + dataProvider = new TestDataProvider(random); + codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + deletionProbability = random.nextFloat() * 0.4f + 0.1f; + vectorProbability = random.nextFloat() * 0.5f + 0.3f; + + log.log( + Level.FINE, + ", deletion probability: " + + deletionProbability + + ", vector probability: " + + vectorProbability); } @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 + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - float[][] dataset = generateDataset(random, datasetSize, dimensions); + try (BaseDirectoryWrapper directory = newDirectory()) { Set deletedDocs = new HashSet<>(); // Create index with all documents having vectors - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { + // Add documents 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); } + writer.commit(); // 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))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log(Level.FINE, "Number of documents deleted: " + deletedDocs.size()); + // 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); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify we got results @@ -102,8 +115,8 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify no deleted documents in results for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse( "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); @@ -111,8 +124,8 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify deleted documents are truly deleted for (int deletedId : deletedDocs) { - TopDocs result = - searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + Query termQuery = new TermQuery(new Term(ID_FIELD, String.valueOf(deletedId))); + TopDocs result = searcher.search(termQuery, 1); assertEquals( "Deleted document " + deletedId + " should not be found", 0, @@ -126,29 +139,25 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { 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<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create index with mixed documents - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { 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 StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); // Randomly assign categories String category = random.nextBoolean() ? "A" : "B"; - doc.add(new StringField("category", category, Field.Store.YES)); + doc.add(new StringField(CATEGORY_FIELD, category, Field.Store.YES)); // Randomly decide whether to add vectors if (random.nextFloat() < vectorProbability) { - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); } else { docsWithoutVectors.add(i); } @@ -158,42 +167,42 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Delete documents randomly for (int i = 0; i < datasetSize; i++) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log(Level.FINE, "Number of deleted documents: " + deletedDocs.size()); + // Test vector search behavior try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); 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); + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse("Deleted document should not appear", deletedDocs.contains(id)); assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); } // Test filtered search with deletions - Query filter = new TermQuery(new Term("category", "A")); - GPUKnnFloatVectorQuery filteredQuery = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filter = new TermQuery(new Term(CATEGORY_FIELD, "A")); + Query filteredQuery = + new GPUKnnFloatVectorQuery(VECTOR_FIELD1, 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"); + String category = doc.get(CATEGORY_FIELD); assertEquals("Should only match category A", "A", category); - int id = Integer.parseInt(doc.get("id")); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse( "Deleted document should not appear in filtered results", deletedDocs.contains(id)); } @@ -205,26 +214,24 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti 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); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create and delete all documents - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], 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.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); } writer.commit(); writer.forceMerge(1); // Force merge to apply deletions @@ -233,10 +240,8 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { // Verify search returns no results try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); TopDocs results = searcher.search(query, topK); assertEquals( @@ -251,22 +256,19 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { 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<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Initial indexing - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { - int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -275,7 +277,7 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio List candidatesForDeletion = new ArrayList<>(activeDocIds); for (int docId : candidatesForDeletion) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(docId))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(docId))); activeDocIds.remove(Integer.valueOf(docId)); } } @@ -283,9 +285,8 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -295,16 +296,13 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // Verify search behavior after deletions and additions try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); 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); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); resultIds.add(id); assertTrue("Result should be from active documents", activeDocIds.contains(id)); } @@ -319,19 +317,4 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio } } } - - 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 index a3e2f59a..d03c8d34 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -4,14 +4,17 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.TreeMap; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -21,14 +24,11 @@ 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.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; @@ -41,101 +41,71 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSGaps extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); - - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); - 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; + private static final Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", isSupported()); directory = newDirectory(); random = random(); + dataProvider = new TestDataProvider(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.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - datasetSize = random.nextInt(100, DATASET_SIZE_LIMIT); - dimension = random.nextInt(8, DIMENSIONS_LIMIT); - dataset = generateDataset(random, datasetSize, dimension); + codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + RandomIndexWriter writer = createWriter(random, directory, codec); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); // 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_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)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); } - writer.addDocument(doc); } - + writer.commit(); 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.log(Level.FINE, "Test finished"); - } - @Test public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", isSupported()); - // Use the first vector (from document 0) as query - float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - // Verify we get exactly TOP_K results + // Verify we get exactly topK 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); + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); log.log(Level.FINE, "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); + // Verify the results match expected top-k + List expectedIds = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}).get(0); + for (ScoreDoc hit : hits) { + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); } @@ -144,48 +114,36 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { @Test public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", isSupported()); - // Use the first vector (from document 0) as query - float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + int datasetSize = dataProvider.getDatasetSize(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - // 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 + String randomEvenInRange = String.valueOf(random.nextInt(datasetSize / 2 + 1) * 2); + log.log(Level.FINE, "Randomly chosen even value is: " + randomEvenInRange); + Query filter = new TermQuery(new Term(ID_FIELD, randomEvenInRange)); - GPUKnnFloatVectorQuery filteredQuery = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = + new GPUKnnFloatVectorQuery(VECTOR_FIELD1, 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) + // Should only get document (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); + Document doc = reader.storedFields().document(filteredHits[0].doc); + assertEquals( + "Should only return document " + randomEvenInRange, randomEvenInRange, doc.get(ID_FIELD)); log.log( Level.FINE, "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(); + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java deleted file mode 100644 index c2e2f2f6..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java +++ /dev/null @@ -1,218 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.TreeMap; -import java.util.logging.Level; -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 = "") -public class TestCuVSRandomizedHNSWVectorSearch extends LuceneTestCase { - - protected static Logger log = - Logger.getLogger(TestCuVSRandomizedHNSWVectorSearch.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - static IndexSearcher searcher; - static IndexReader reader; - static Directory directory; - - static int DATASET_SIZE_LIMIT = 1000; - static int DIMENSIONS_LIMIT = 2048; - static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 - static float[][] dataset; - - @BeforeClass - public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - directory = newDirectory(); - - RandomIndexWriter writer = - new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - - log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - Random random = random(); - int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; - int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; - dataset = generateDataset(random, datasetSize, dimensions); - 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)); - boolean skipVector = - 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)); - doc.add(new KnnFloatVectorField("vector2", 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.log(Level.FINE, "Test finished"); - } - - @Test - public void testVectorSearch() throws IOException { - Random random = random(); - int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; - int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - - if (dataset.length < topK) topK = dataset.length; - - float[][] queries = generateQueries(random, dataset[0].length, numQueries); - List> expected = generateExpectedResults(topK, dataset, queries); - - log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); - log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); - log.log(Level.FINE, "TopK: " + topK); - - Query query = new KnnFloatVectorQuery("vector", queries[0], topK); - int correct[] = new int[topK]; - for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); - - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); - log.log(Level.FINE, "EXPECTD: " + expected.get(0)); - - for (ScoreDoc hit : hits) { - log.log( - Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); - } - - for (ScoreDoc hit : hits) { - int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); - assertTrue("Result returned was not in topk*2: " + doc, expected.get(0).contains(doc)); - } - } - - private static List> generateExpectedResults( - int topK, float[][] dataset, float[][] queries) { - List> neighborsResult = new ArrayList<>(); - int dimensions = dataset[0].length; - - for (float[] query : queries) { - Map distances = new TreeMap<>(); - for (int j = 0; j < dataset.length; j++) { - double distance = 0; - for (int k = 0; k < dimensions; k++) { - distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); - } - distances.put(j, (distance)); - } - - Map sorted = new TreeMap(distances); - log.log(Level.FINE, "EXPECTED: " + sorted); - - // Sort by distance and select the topK nearest neighbors - List neighbors = - distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .toList(); - neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); - } - - log.log(Level.FINE, "Expected results generated successfully."); - return neighborsResult; - } - - @Test - public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuVS not supported", isSupported()); - - 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 KnnFloatVectorQuery("vector", queryVector, topK, filter); - - 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.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); - } -} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 3e77e2ed..1750d51a 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -4,17 +4,19 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; -import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.TreeMap; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -24,14 +26,12 @@ 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; @@ -44,50 +44,36 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); + private static final Logger log = + Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); - static IndexSearcher searcher; - static IndexReader reader; - static Directory directory; - - static int DATASET_SIZE_LIMIT = 1000; - static int DIMENSIONS_LIMIT = 2048; - static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 - static float[][] dataset; + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", isSupported()); + assumeTrue("cuVS not supported so skipping these tests", isSupported()); + codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); directory = newDirectory(); + random = random(); + dataProvider = new TestDataProvider(random); + + RandomIndexWriter writer = createWriter(random, directory, codec); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); - RandomIndexWriter writer = - new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - - log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - Random random = random(); - int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; - int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; - dataset = generateDataset(random, datasetSize, dimensions); 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)); - boolean skipVector = - 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)); - doc.add(new KnnFloatVectorField("vector2", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); + boolean skipVector = random.nextInt(10) < 4; + if (!skipVector || datasetSize < 100) { + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, dataset[i], EUCLIDEAN)); } writer.addDocument(doc); @@ -98,109 +84,55 @@ public static void beforeClass() throws Exception { 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.log(Level.FINE, "Test finished"); - } - @Test public void testVectorSearch() throws IOException { - Random random = random(); - int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; - int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - - if (dataset.length < topK) topK = dataset.length; - - float[][] queries = generateQueries(random, dataset[0].length, numQueries); - List> expected = generateExpectedResults(topK, dataset, queries); - - log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); - log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); - log.log(Level.FINE, "TopK: " + topK); - - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queries[0], topK, null, topK, 1); - int correct[] = new int[topK]; - for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); - - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); - log.log(Level.FINE, "EXPECTD: " + expected.get(0)); - - for (ScoreDoc hit : hits) { - log.log( - Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); - } - - for (ScoreDoc hit : hits) { - int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); - assertTrue("Result returned was not in topk*2: " + doc, expected.get(0).contains(doc)); - } - } - - private static List> generateExpectedResults( - int topK, float[][] dataset, float[][] queries) { - List> neighborsResult = new ArrayList<>(); - int dimensions = dataset[0].length; - - for (float[] query : queries) { - Map distances = new TreeMap<>(); - for (int j = 0; j < dataset.length; j++) { - double distance = 0; - for (int k = 0; k < dimensions; k++) { - distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); - } - distances.put(j, (distance)); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + int numQueries = dataProvider.getNumQueries(); + float[][] queries = dataProvider.getQueries(numQueries); + + // Generate queries and expected results for each + List> expected = generateExpectedTopK(topK, dataset, queries); + + for (int i = 0; i < numQueries; i++) { + log.log(Level.FINE, "Running query: " + (i + 1) + " of " + numQueries); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queries[i], topK); + + // Perform search + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); + log.log(Level.FINE, "EXPECTED: " + expected.get(i)); + + // Iterate through the results and assert + for (ScoreDoc hit : hits) { + Document doc = reader.storedFields().document(hit.doc); + int docId = Integer.parseInt(doc.get(ID_FIELD)); + log.log(Level.FINE, "\t" + doc.get(ID_FIELD) + ": " + hit.score); + assertTrue("Result returned was not in topk*2: " + doc, expected.get(i).contains(docId)); } - - Map sorted = new TreeMap(distances); - log.log(Level.FINE, "EXPECTED: " + sorted); - - // Sort by distance and select the topK nearest neighbors - List neighbors = - distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .toList(); - neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); } - - log.log(Level.FINE, "Expected results generated successfully."); - return neighborsResult; } @Test public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuVS not supported", isSupported()); - - 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 - GPUKnnFloatVectorQuery unfiltered = - new GPUKnnFloatVectorQuery("vector", dataset[0], 1, null, 1, 1); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + Query unfiltered = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, 1); ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; - // Skip test if no vectors found at all - assumeTrue( + assertTrue( "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]; + Document doc = reader.storedFields().document(unfilteredHits[0].doc); + String targetDocId = doc.get(ID_FIELD); // Create a filter that matches only the document we know has a vector - Query filter = new TermQuery(new Term("id", targetDocId)); + Query filter = new TermQuery(new Term(ID_FIELD, targetDocId)); // Test the new constructor with filter - GPUKnnFloatVectorQuery filteredQuery = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; @@ -209,10 +141,20 @@ public void testVectorSearchWithFilter() throws IOException { // Verify that all results match the filter for (ScoreDoc hit : filteredHits) { - String docId = reader.storedFields().document(hit.doc).get("id"); + String docId = reader.storedFields().document(hit.doc).get(ID_FIELD); assertEquals("All results should match the filter", targetDocId, docId); } log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.log(Level.FINE, "Test finished"); + } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 1c75cb23..dc118250 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -4,10 +4,14 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.util.List; +import java.util.Random; import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; @@ -29,9 +33,14 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { + private static TestDataProvider dataProvider; + private static Random random; + @BeforeClass public static void beforeClass() { - assumeTrue("cuVS is not supported", isSupported()); + assumeTrue("cuVS is not supported so skipping these tests", isSupported()); + random = random(); + dataProvider = new TestDataProvider(random); } @Override @@ -40,27 +49,27 @@ protected Codec getCodec() { } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { - float[][] f = new float[][] {randomVector(384), randomVector(384)}; + final int numDocs = 2; + float[][] vectors = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); - w.addDocument(doc1); - w.commit(); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); - w.addDocument(doc2); - w.flush(); - w.commit(); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors[i], EUCLIDEAN)); + w.addDocument(doc); + w.commit(); + } // sanity - verify one doc per leaf try (DirectoryReader reader = DirectoryReader.open(w)) { List subReaders = reader.leaves(); - assertEquals(2, subReaders.size()); - assertEquals(1, subReaders.get(0).reader().getFloatVectorValues("f").size()); - assertEquals(1, subReaders.get(1).reader().getFloatVectorValues("f").size()); + assertEquals(numDocs, subReaders.size()); + for (int i = 0; i < numDocs; i++) { + assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(VECTOR_FIELD1).size()); + } } // now merge to a single segment @@ -69,49 +78,56 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertEquals(numDocs, values.size()); + for (int i = 0; i < numDocs; i++) { + assertArrayEquals(vectors[i], values.vectorValue(i), 0.0f); + } } } } // Basic test for multiple vectors fields per document public void testTwoVectorFieldsPerDoc() throws Exception { - float[][] f1 = new float[][] {randomVector(384), randomVector(384)}; - float[][] f2 = new float[][] {randomVector(384), randomVector(384)}; + + final int numDocs = 2; + float[][] vectors1 = dataProvider.getVectors(numDocs); + float[][] vectors2 = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); - doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); - w.addDocument(doc1); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); - doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); - w.addDocument(doc2); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors1[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, vectors2[i], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f1"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); - - values = r.getFloatVectorValues("f2"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); - // opportunistically check boundary condition - search with a 0 topK - var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors1[i], values.vectorValue(i), 0.0f); + } + + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD2); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors2[i], values.vectorValue(i), 0.0f); + } + + // Check boundary condition - search with a 0 topK + float[] target = dataProvider.getVectors(1)[0]; + var topDocs = r.searchNearestVectors(VECTOR_FIELD1, target, 0, null, 10); assertEquals(0, topDocs.scoreDocs.length); assertEquals(0, topDocs.totalHits.value()); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java b/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java new file mode 100644 index 00000000..954bd52d --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVectors; + +import java.util.Random; +import java.util.logging.Level; +import java.util.logging.Logger; + +public class TestDataProvider { + + private static final Logger log = Logger.getLogger(TestDataProvider.class.getName()); + private static final int DATASET_SIZE_LIMIT = 1000; + private static final int DATASET_SIZE_MIN = 200; + private static final int DIMENSIONS_LIMIT = 256; + private static final int DIMENSIONS_MIN = 8; + private static final int TOP_K_LIMIT = 64; + private static final int TOP_K_MIN = 2; + private static final int QUERIES_LIMIT = 50; + private static final int QUERIES_MIN = 2; + + public static final String ID_FIELD = "id"; + public static final String TEXT_FIELD = "some_text_field"; + public static final String CATEGORY_FIELD = "category_field"; + public static final String VECTOR_FIELD1 = "vector_field1"; + public static final String VECTOR_FIELD2 = "vector_field2"; + + private int datasetSize; + private int dimensions; + private int topK; + private float[][] dataset1; + private float[][] dataset2; + private int numQueries; + private Random random; + + public TestDataProvider(Random random) { + datasetSize = random.nextInt(DATASET_SIZE_MIN, DATASET_SIZE_LIMIT); + dimensions = random.nextInt(DIMENSIONS_MIN, DIMENSIONS_LIMIT); + topK = Math.min(random.nextInt(TOP_K_MIN, TOP_K_LIMIT), datasetSize); + dataset1 = generateRandomVectors(random, datasetSize, dimensions); + dataset2 = generateRandomVectors(random, datasetSize, dimensions); + numQueries = random.nextInt(QUERIES_MIN, QUERIES_LIMIT); + this.random = random; + log.log(Level.FINE, "Dataset size: " + datasetSize + "x" + dimensions + ", topK: " + topK); + } + + public int getDatasetSize() { + return datasetSize; + } + + public int getDimensions() { + return dimensions; + } + + public int getTopK() { + return topK; + } + + public float[][] getDataset1() { + return dataset1; + } + + public float[][] getDataset2() { + return dataset2; + } + + public float[][] getQueries(int numQueries) { + return generateRandomVectors(random, numQueries, dimensions); + } + + public float[][] getVectors(int numVectors) { + return generateRandomVectors(random, numVectors, dimensions); + } + + public int getRandom(int min, int max) { + return random.nextInt(min, max); + } + + public double getRandom(double min, double max) { + return random.nextDouble(min, max); + } + + public int getNumQueries() { + return numQueries; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index bcc29178..c2cfb112 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -4,7 +4,11 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomText; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; import com.carrotsearch.randomizedtesting.annotations.Name; @@ -14,10 +18,13 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashSet; import java.util.List; import java.util.Random; +import java.util.Set; import java.util.logging.Level; 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; @@ -28,10 +35,10 @@ 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.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; @@ -58,10 +65,10 @@ 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; + private static Random random; + private static TestDataProvider dataProvider; + private static Directory directory; + private static Codec codec; private static CagraGraphBuildAlgo cagraGraphBuildAlgo; @@ -78,23 +85,15 @@ public static List parameters() { @BeforeClass public static void beforeClass() { assumeTrue("cuVS is not supported", isSupported()); + random = random(); + codec = alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); } - 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.log(Level.FINE, "Using randomized vector dimension: " + vectorDimension); + dataProvider = new TestDataProvider(random); } @After @@ -116,13 +115,11 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { + cagraGraphBuildAlgo); // 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 maxBufferedDocs = dataProvider.getRandom(5, 16); + int totalBatches = dataProvider.getRandom(8, 16); + int docsPerBatch = dataProvider.getRandom(15, 25); int totalDocuments = totalBatches * docsPerBatch; - - // Randomize vector presence probability (60-85%) - double vectorProbability = 0.6 + (random().nextDouble() * 0.25); + double vectorProbability = dataProvider.getRandom(0.6, 0.8); log.log( Level.FINE, @@ -145,10 +142,9 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { IndexWriterConfig config = new IndexWriterConfig() .setCodec(alwaysKnnVectorsFormat(format)) - .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size + .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - List expectedVectors = new ArrayList<>(); List expectedDocIds = new ArrayList<>(); int documentsWithVectors = 0; @@ -158,21 +154,18 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); // 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); + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); expectedDocIds.add(docId); documentsWithVectors++; } - writer.addDocument(doc); } - writer.commit(); // Create a new segment + writer.commit(); } int documentsWithoutVectors = totalDocuments - documentsWithVectors; @@ -181,87 +174,71 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { log.log(Level.FINE, "Documents with vectors: " + documentsWithVectors); log.log(Level.FINE, "Documents without vectors: " + documentsWithoutVectors); - // Force merge to trigger merge logic writer.forceMerge(1); log.log(Level.FINE, "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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = 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()); + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getTopK(); - // Randomize search parameters - int searchK = - Math.min(5 + random().nextInt(10), Math.min(documentsWithVectors, TOP_K_LIMIT)); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs results = searcher.search(query, topK); - 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 some results", results.scoreDocs.length <= documentsWithVectors); - assertTrue("Should find some results after merge", results.scoreDocs.length > 0); - assertTrue( - "Should find reasonable number of results", - results.scoreDocs.length <= documentsWithVectors); - - log.log( - Level.FINE, - "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.log(Level.FINE, "No documents with vectors - skipping vector search verification"); + log.log( + Level.FINE, + "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) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + int docId = Integer.parseInt(doc.get(ID_FIELD)); + assertTrue("Document ID should be valid", expectedDocIds.contains(docId)); } - - log.log(Level.FINE, "Merge verification completed successfully"); } } /** - * Test merging with index sorting enabled using text-based sorting and SortingMergePolicy + * Test merging with index sorting enabled using SortingMergePolicy **/ @Test - public void testMergeWithIndexSorting() throws IOException { + public void testMergeWithIndexSortingStringField() throws IOException { log.log( Level.FINE, "Starting testMergeWithIndexSorting with text-based sorting with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); // 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"; - + final String SORT_FIELD_NAME = "text_sort_key"; + final String ORIGINAL_ORDER = "original_order"; // Configure index sorting by a randomized field - Sort indexSort = new Sort(new SortField(sortFieldName, sortType)); + Sort indexSort = new Sort(new SortField(SORT_FIELD_NAME, SortField.Type.STRING)); // 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 + mergePolicy.setMaxMergedSegmentMB(dataProvider.getRandom(128, 385)); + mergePolicy.setSegmentsPerTier(dataProvider.getRandom(3, 7)); // 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 + int maxBufferedDocs = dataProvider.getRandom(10, 26); + int totalDocuments = dataProvider.getRandom(80, 161); + int segmentSize = dataProvider.getRandom(15, 26); + double vectorProbability = dataProvider.getRandom(0.65, 0.91); - log.log( - Level.FINE, - "Randomized sorting parameters: sortType=" + sortType + ", sortFieldName=" + sortFieldName); log.log( Level.FINE, "Randomized config: maxBufferedDocs=" @@ -286,40 +263,23 @@ public void testMergeWithIndexSorting() throws IOException { .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - // List documents = new ArrayList<>(); - try (IndexWriter writer = new IndexWriter(directory, config)) { // Create documents with randomized sort keys + int numDocsWithVectors = 0; 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)); - } + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new StringField(ORIGINAL_ORDER, String.valueOf(i), Field.Store.YES)); + + String textSortKey = generateRandomText(random, dataProvider.getRandom(4, 21)); + doc.add(new SortedDocValuesField(SORT_FIELD_NAME, new BytesRef(textSortKey))); + doc.add(new StringField(SORT_FIELD_NAME + "_stored", textSortKey, Field.Store.YES)); - if (vector != null) { - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + numDocsWithVectors++; } writer.addDocument(doc); @@ -337,7 +297,13 @@ public void testMergeWithIndexSorting() throws IOException { } } - log.log(Level.FINE, "Created " + totalDocuments + " documents with text-based index sorting"); + log.log( + Level.FINE, + "Number of documents with vectors is: " + + numDocsWithVectors + + " out of a total of " + + totalDocuments + + " documents"); // Force merge with sorting - this will use the sorting merge policy writer.forceMerge(1); @@ -346,95 +312,209 @@ public void testMergeWithIndexSorting() throws IOException { // 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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); - // Verify documents are sorted correctly by the randomized sort field + String previousSortKey = ""; + SortedDocValues sortedValues = leafReader.getSortedDocValues(SORT_FIELD_NAME); + + 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 " + + SORT_FIELD_NAME + + ": '" + + previousSortKey + + "' should be <= '" + + currentSortKey + + "'", + previousSortKey.compareTo(currentSortKey) <= 0); + previousSortKey = currentSortKey; + } + + // Count total vectors by checking if vector field exists and has values + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); + int documentsWithVectors = vectorValues != null ? vectorValues.size() : 0; + log.log( Level.FINE, - "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(); - } + "Found " + documentsWithVectors + " documents with vectors after sorted merge"); - 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.log( - Level.FINE, - "DocId: " - + docId - + ", OriginalOrder: " - + originalOrder - + ", SortKey: '" - + currentSortKey - + "'"); - } + // Test vector search on sorted index + if (documentsWithVectors > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + + KnnFloatVectorQuery query = + new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, Math.min(10, documentsWithVectors)); + TopDocs results = searcher.search(query, 10); + + assertTrue("Should find results in sorted index", results.scoreDocs.length > 0); + log.log( + Level.FINE, + "Vector search on sorted index returned " + results.scoreDocs.length + " results"); + + // Verify that returned documents maintain sort order if we check their sort keys + log.log(Level.FINE, "Verifying vector search results maintain sorting consistency..."); + for (int i = 0; i < Math.min(5, results.scoreDocs.length); i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String originalOrder = doc.get(ORIGINAL_ORDER); + String sortKey = doc.get(SORT_FIELD_NAME + "_stored"); + log.log( + Level.FINE, + "Result " + + i + + ": DocId=" + + scoreDoc.doc + + ", OriginalOrder=" + + originalOrder + + ", SortKey='" + + sortKey + + "', Score=" + + scoreDoc.score); } - } 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.log( - Level.FINE, - "DocId: " - + docId - + ", OriginalOrder: " - + originalOrder - + ", SortKey: " - + currentSortKey); - } + /** + * Test merging with index sorting enabled using SortingMergePolicy + **/ + @Test + public void testMergeWithIndexSortingLongField() throws IOException { + final String SORT_FIELD_NAME = "numeric_sort_key"; + final String ORIGINAL_ORDER = "original_order"; + Sort indexSort = new Sort(new SortField(SORT_FIELD_NAME, SortField.Type.LONG)); + + // Randomize merge policy parameters + TieredMergePolicy mergePolicy = new TieredMergePolicy(); + mergePolicy.setMaxMergedSegmentMB(dataProvider.getRandom(128, 385)); + mergePolicy.setSegmentsPerTier(dataProvider.getRandom(3, 7)); + + // Randomize writer configuration parameters + int maxBufferedDocs = dataProvider.getRandom(10, 26); + int totalDocuments = dataProvider.getRandom(80, 161); + int segmentSize = dataProvider.getRandom(15, 26); + double vectorProbability = dataProvider.getRandom(0.65, 0.91); + + log.log( + Level.FINE, + "Randomized config: maxBufferedDocs=" + + maxBufferedDocs + + ", totalDocuments=" + + totalDocuments + + ", segmentSize=" + + segmentSize + + ", vectorProbability=" + + vectorProbability); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setIndexSort(indexSort) + .setMergePolicy(mergePolicy) + .setMaxBufferedDocs(maxBufferedDocs) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + int numDocsWithVectors = 0; + for (int i = 0; i < totalDocuments; i++) { + + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new StringField(ORIGINAL_ORDER, String.valueOf(i), Field.Store.YES)); + + long numericSortKey = random.nextLong() % 100000; + doc.add(new NumericDocValuesField(SORT_FIELD_NAME, numericSortKey)); + doc.add( + new StringField( + SORT_FIELD_NAME + "_stored", String.valueOf(numericSortKey), Field.Store.YES)); + + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + numDocsWithVectors++; + } + + writer.addDocument(doc); + + // Commit based on randomized segment size + if ((i + 1) % segmentSize == 0) { + writer.commit(); + log.log( + Level.FINE, + "Committed segment " + + ((i + 1) / segmentSize) + + " with " + + (i + 1) + + " total documents"); + } + } + + log.log( + Level.FINE, + "Number of documents with vectors is: " + + numDocsWithVectors + + " out of a total of " + + totalDocuments + + " documents"); + + // Force merge with sorting - this will use the sorting merge policy + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge with text-based sorting completed"); + } + + // Verify the merged and sorted index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Verify numeric-based sorting + long previousSortKey = Long.MIN_VALUE; + var numericValues = leafReader.getNumericDocValues(SORT_FIELD_NAME); + + 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 " + + SORT_FIELD_NAME + + ": " + + 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.log( + Level.FINE, + "DocId: " + + docId + + ", OriginalOrder: " + + originalOrder + + ", SortKey: " + + currentSortKey); } } // Count total vectors by checking if vector field exists and has values - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int documentsWithVectors = vectorValues != null ? vectorValues.size() : 0; log.log( @@ -444,10 +524,10 @@ public void testMergeWithIndexSorting() throws IOException { // Test vector search on sorted index if (documentsWithVectors > 0) { IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); + float[] queryVector = dataProvider.getQueries(1)[0]; KnnFloatVectorQuery query = - new KnnFloatVectorQuery("vector", queryVector, Math.min(10, documentsWithVectors)); + new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, Math.min(10, documentsWithVectors)); TopDocs results = searcher.search(query, 10); assertTrue("Should find results in sorted index", results.scoreDocs.length > 0); @@ -457,12 +537,11 @@ public void testMergeWithIndexSorting() throws IOException { // Verify that returned documents maintain sort order if we check their sort keys log.log(Level.FINE, "Verifying vector search results maintain sorting consistency..."); - for (int i = 0; i < Math.min(3, results.scoreDocs.length); i++) { + for (int i = 0; i < Math.min(5, 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"); + Document doc = searcher.storedFields().document(scoreDoc.doc); + String originalOrder = doc.get(ORIGINAL_ORDER); + String sortKey = doc.get(SORT_FIELD_NAME + "_stored"); log.log( Level.FINE, "Result " @@ -477,8 +556,6 @@ public void testMergeWithIndexSorting() throws IOException { + scoreDoc.score); } } - - log.log(Level.FINE, "Text-based index sorting verification completed successfully"); } } @@ -493,7 +570,8 @@ public void testMergeWithMissingVectors() throws IOException { // Randomize configuration int maxBufferedDocs = 10 + random().nextInt(11); // 10-20 docs per buffer - int numSegments = 3 + random().nextInt(3); // 3-5 segments + int numSegments = dataProvider.getRandom(3, 13); + log.log(Level.FINE, "Randomized parameters: numSegments=" + numSegments); GPUSearchParams params = new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); @@ -518,27 +596,20 @@ public void testMergeWithMissingVectors() throws IOException { 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 docsInSegment = dataProvider.getRandom(15, 100); + double vectorProbability = dataProvider.getRandom(0.1, 0.6); 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); // 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)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); segmentVectorCount++; } - writer.addDocument(doc); } @@ -555,7 +626,7 @@ public void testMergeWithMissingVectors() throws IOException { + " documents, " + segmentVectorCount + " with vectors (probability: " - + String.format("%.2f", vectorProbability) + + vectorProbability + ")"); } @@ -566,13 +637,13 @@ public void testMergeWithMissingVectors() throws IOException { // 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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); // Count actual vectors in merged index - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; log.log( @@ -589,31 +660,22 @@ public void testMergeWithMissingVectors() throws IOException { // Test vector search if we have vectors if (actualVectorCount > 0) { IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = Math.min(dataProvider.getTopK(), actualVectorCount); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); - // 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); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); log.log( Level.FINE, - "Found " - + vectorResults.scoreDocs.length - + " vector results out of " - + actualVectorCount - + " available"); + "Found " + numResults + " vector results out of " + actualVectorCount + " available"); + assertEquals("Search should return exactly topK results", topK, numResults); } else { log.log(Level.FINE, "No vectors in merged index - skipping vector search"); } - - log.log(Level.FINE, "Missing vectors test completed successfully"); } } @@ -622,22 +684,15 @@ public void testMergeWithMissingVectors() throws IOException { **/ @Test public void testMergeWithDeletions() throws IOException { - log.log( - Level.FINE, - "Starting testMergeWithDeletions with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); - - // 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 + int numSegments = dataProvider.getRandom(3, 7); + int docsPerSegment = dataProvider.getRandom(20, 41); + int maxBufferedDocs = dataProvider.getRandom(8, 17); + double vectorProbability = dataProvider.getRandom(0.7, 0.95); + double deletionProbability = dataProvider.getRandom(0.2, 0.5); log.log( Level.FINE, - "Randomized parameters: maxBufferedDocs=" - + maxBufferedDocs - + ", numSegments=" + "Randomized parameters: numSegments=" + numSegments + ", docsPerSegment=" + docsPerSegment @@ -660,6 +715,7 @@ public void testMergeWithDeletions() throws IOException { List expectedRemainingDocs = new ArrayList<>(); List deletedDocs = new ArrayList<>(); int totalDocuments = numSegments * docsPerSegment; + int numDocsWithVectors = 0; try (IndexWriter writer = new IndexWriter(directory, config)) { // Create multiple segments with documents @@ -667,17 +723,14 @@ public void testMergeWithDeletions() throws IOException { 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); // Randomly add vectors - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + numDocsWithVectors++; } - writer.addDocument(doc); } writer.commit(); @@ -692,12 +745,13 @@ public void testMergeWithDeletions() throws IOException { + " documents each (" + totalDocuments + " total)"); + log.log(Level.FINE, "Number of docs with vectors: " + numDocsWithVectors); // 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))); + if (random.nextDouble() < deletionProbability) { + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(docId))); deletedDocs.add(docId); deletedCount++; } else { @@ -710,7 +764,7 @@ public void testMergeWithDeletions() throws IOException { "Deleted " + deletedCount + " documents (" - + String.format("%.1f", (100.0 * deletedCount / totalDocuments)) + + (100.0 * deletedCount / totalDocuments) + "%), remaining: " + expectedRemainingDocs.size()); @@ -723,9 +777,9 @@ public void testMergeWithDeletions() throws IOException { // 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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); int expectedRemaining = expectedRemainingDocs.size(); assertEquals( "Should have correct number of documents after deletions", @@ -738,31 +792,30 @@ public void testMergeWithDeletions() throws IOException { // 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); + TopDocs result = + searcher.search(new TermQuery(new Term(ID_FIELD, 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 + int deletedDocsToCheck = Math.min(10, deletedDocs.size()); 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); + TopDocs result = + searcher.search(new TermQuery(new Term(ID_FIELD, 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.log( - Level.FINE, - "Found " + vectorResults.scoreDocs.length + " vector results after deletions"); - log.log(Level.FINE, "Deletion merge verification completed successfully"); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = Math.min(1, numDocsWithVectors); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results after deletions", numResults > 0); + assertEquals("Search should return exactly topK documents", numResults, topK); + log.log(Level.FINE, "Found " + numResults + " vector results after deletions"); } } @@ -776,10 +829,10 @@ public void testMergeBruteForceIndex() throws IOException { "Starting testMergeBruteForceIndex with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); // 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 + int numSegments = dataProvider.getRandom(3, 10); + int docsPerSegment = dataProvider.getRandom(20, 100); + double vectorProbability = dataProvider.getRandom(0.2, 0.7); + int maxBufferedDocs = dataProvider.getRandom(8, 17); log.log( Level.FINE, @@ -810,27 +863,23 @@ public void testMergeBruteForceIndex() throws IOException { int totalDocuments = numSegments * docsPerSegment; int totalExpectedVectors = 0; + Set docIDsHavingVectors = new HashSet(); 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); // Randomly add vectors based on probability - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getVectors(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + docIDsHavingVectors.add(docId); segmentVectorCount++; } - writer.addDocument(doc); } @@ -865,13 +914,13 @@ public void testMergeBruteForceIndex() throws IOException { // 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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); // Count actual vectors in merged index - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; log.log( @@ -886,61 +935,53 @@ public void testMergeBruteForceIndex() throws IOException { 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)); + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getRandom(1, actualVectorCount); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results in brute force index", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); + log.log( + Level.FINE, + "Brute force search found " + + numResults + + " results out of " + + actualVectorCount + + " available vectors"); + + assertEquals("Search should return exactly topK documents", topK, numResults); + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + String docId = doc.get(ID_FIELD); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); 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.log( - Level.FINE, - "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.log(Level.FINE, "No vectors in brute force merged index - skipping vector search"); + "Document does not have a vector", + docIDsHavingVectors.contains(Integer.parseInt(docId))); } - - log.log(Level.FINE, "Brute force merge verification completed successfully"); } } /** - * Test merging segments for {@link IndexType#CAGRA_AND_BRUTE_FORCE} + * Test merging segments for {@link IndexType#CAGRA} * */ @Test - public void testMergeCagraAndBruteForceIndex() throws IOException { + public void testMergeCagra() throws IOException { log.log( Level.FINE, "Starting testMergeCagraAndBruteForceIndex with CagraGraphBuildAlgo: " + cagraGraphBuildAlgo); // 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 + int numSegments = dataProvider.getRandom(3, 10); + int docsPerSegment = dataProvider.getRandom(20, 100); + double vectorProbability = dataProvider.getRandom(0.2, 0.7); + int maxBufferedDocs = dataProvider.getRandom(8, 17); log.log( Level.FINE, @@ -957,42 +998,36 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { GPUSearchParams params = new GPUSearchParams.Builder() .withCagraGraphBuildAlgo(cagraGraphBuildAlgo) - .withIndexType(IndexType.CAGRA_AND_BRUTE_FORCE) + .withIndexType(IndexType.CAGRA) .build(); - CuVS2510GPUVectorsFormat combinedFormat = - new CuVS2510GPUVectorsFormat(params); // Use combined CAGRA + brute force + CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(combinedFormat)) + .setCodec(alwaysKnnVectorsFormat(format)) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); int totalDocuments = numSegments * docsPerSegment; int totalExpectedVectors = 0; + Set docIDsHavingVectors = new HashSet(); 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)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); - // Add vectors based on probability - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + // Randomly add vectors based on probability + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getVectors(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + docIDsHavingVectors.add(docId); segmentVectorCount++; } - writer.addDocument(doc); } @@ -1001,7 +1036,7 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { log.log( Level.FINE, - "Created CAGRA+brute force segment " + "Created CAGRA segment " + seg + ": " + docsPerSegment @@ -1014,31 +1049,31 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { Level.FINE, "Created " + numSegments - + " CAGRA+brute force segments with " + + " CAGRA segments with " + totalDocuments + " total documents and " + totalExpectedVectors + " vectors"); - // Force merge all CAGRA+brute force segments + // Force merge all CAGRA segments writer.forceMerge(1); - log.log(Level.FINE, "Forced merge of CAGRA+brute force segments completed"); + log.log(Level.FINE, "Forced merge of brute force segments completed"); } - // Verify the merged CAGRA+brute force index + // Verify the merged CAGRA 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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); // Count actual vectors in merged index - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; log.log( Level.FINE, - "CAGRA+brute force merge results: Total documents: " + "CAGRA merge results: Total documents: " + totalDocuments + ", Expected vectors: " + totalExpectedVectors @@ -1047,98 +1082,65 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { 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)); + // Test CAGRA vector search (exact search) + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getRandom(1, actualVectorCount); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results in CAGRA + brute force index", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); + log.log( + Level.FINE, + "CAGRA search found " + + numResults + + " results out of " + + actualVectorCount + + " available vectors"); + + assertEquals("Search should return exactly topK documents", topK, numResults); + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + String docId = doc.get(ID_FIELD); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); 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.log( - Level.FINE, - "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.log( - Level.FINE, - "Trial " + trial + " found " + trialResults.scoreDocs.length + " results"); - } - } else { - log.log( - Level.FINE, "No vectors in CAGRA+brute force merged index - skipping vector search"); + "Document does not have a vector", + docIDsHavingVectors.contains(Integer.parseInt(docId))); } - - log.log(Level.FINE, "CAGRA+brute force merge verification completed successfully"); } } /** - * Test large scale merge to stress test the system - **/ + * Test merging segments for {@link IndexType#CAGRA_AND_BRUTE_FORCE} + * */ @Test - public void testLargeScaleMerge() throws IOException { - assumeTrue( - "testLargeScaleMerge requires -DlargeScale=true", - Boolean.parseBoolean(System.getProperty("largeScale", "false"))); - - log.log(Level.FINE, "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; + public void testMergeCagraAndBruteForceIndex() throws IOException { + int numSegments = dataProvider.getRandom(3, 10); + int docsPerSegment = dataProvider.getRandom(20, 100); + double vectorProbability = dataProvider.getRandom(0.2, 0.7); + int maxBufferedDocs = dataProvider.getRandom(8, 17); log.log( Level.FINE, - "Randomized large scale parameters: maxBufferedDocs=" + "Randomized parameters: maxBufferedDocs=" + maxBufferedDocs - + ", segmentCount=" - + segmentCount + + ", numSegments=" + + numSegments + ", docsPerSegment=" + docsPerSegment - + ", totalDocuments=" - + totalDocuments); + + ", vectorProbability=" + + vectorProbability); GPUSearchParams params = - new GPUSearchParams.Builder().withCagraGraphBuildAlgo(cagraGraphBuildAlgo).build(); + new GPUSearchParams.Builder() + .withCagraGraphBuildAlgo(cagraGraphBuildAlgo) + .withIndexType(IndexType.CAGRA_AND_BRUTE_FORCE) + .build(); CuVS2510GPUVectorsFormat format = new CuVS2510GPUVectorsFormat(params); @@ -1148,113 +1150,109 @@ public void testLargeScaleMerge() throws IOException { .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - try (IndexWriter writer = new IndexWriter(directory, config)) { - for (int seg = 0; seg < segmentCount; seg++) { - log.log(Level.FINE, "Creating segment " + (seg + 1) + "/" + segmentCount); - - // Randomize vector probability per segment - double vectorProbability = - 0.5 + (random().nextDouble() * 0.4); // 50-90% vectors per segment + int totalDocuments = numSegments * docsPerSegment; + int totalExpectedVectors = 0; + Set docIDsHavingVectors = new HashSet(); + try (IndexWriter writer = new IndexWriter(directory, config)) { + 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 NumericDocValuesField("segment", seg)); - doc.add(new NumericDocValuesField("position", i)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); - // 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)); + // Randomly add vectors based on probability + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getVectors(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + docIDsHavingVectors.add(docId); + segmentVectorCount++; } - writer.addDocument(doc); } + writer.commit(); + totalExpectedVectors += segmentVectorCount; + + log.log( + Level.FINE, + "Created CAGRA + brute force segment " + + seg + + ": " + + docsPerSegment + + " documents, " + + segmentVectorCount + + " with vectors"); } log.log( Level.FINE, - "Created " + segmentCount + " segments with " + totalDocuments + " total documents"); + "Created " + + numSegments + + " CAGRA + brute force segments with " + + totalDocuments + + " total documents and " + + totalExpectedVectors + + " vectors"); - // Force merge all segments - long startTime = System.currentTimeMillis(); + // Force merge all CAGRA + brute force segments writer.forceMerge(1); - long mergeTime = System.currentTimeMillis() - startTime; - - log.log(Level.FINE, "Large scale merge completed in " + mergeTime + "ms"); + log.log(Level.FINE, "Forced merge of brute force segments completed"); } - // Verify the large merged index + // 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(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); - // Test vector search performance - var vectorValues = leafReader.getFloatVectorValues("vector"); + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); 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 + log.log( + Level.FINE, + "CAGRA + brute force merge results: Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); - long searchStart = System.currentTimeMillis(); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); - long searchTime = System.currentTimeMillis() - searchStart; + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); - assertTrue("Should find vector results in large index", vectorResults.scoreDocs.length > 0); - log.log( - Level.FINE, - "Vector search in large index returned " - + vectorResults.scoreDocs.length - + " results out of " - + actualVectorCount - + " vectors in " - + searchTime - + "ms"); - } else { - log.log(Level.FINE, "No vectors in large merged index - skipping vector search"); - } + // Test CAGRA + brute force vector search (exact search) + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getRandom(1, actualVectorCount); - log.log(Level.FINE, "Large scale merge verification completed successfully"); - } - } + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results in CAGRA + brute force index", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); - /** 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; + log.log( + Level.FINE, + "CAGRA + Brute force search found " + + numResults + + " results out of " + + actualVectorCount + + " available vectors"); + + assertEquals("Search should return exactly topK documents", topK, numResults); + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + String docId = doc.get(ID_FIELD); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); + assertTrue( + "Document does not have a vector", + docIDsHavingVectors.contains(Integer.parseInt(docId))); } } - 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/TestMultithreadedCuVSGPUSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java index 25dc13b8..c5065d5f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java @@ -5,7 +5,6 @@ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -55,25 +54,26 @@ public class TestMultithreadedCuVSGPUSearch extends LuceneTestCase { private static int numQueries; private static int topK; private static int numThreads; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws IOException { assumeTrue("cuVS not supported", isSupported()); random = random(); + dataProvider = new TestDataProvider(random); directory = newDirectory(new ByteBuffersDirectory()); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec); IndexWriter writer = new IndexWriter(directory, config); - int datasetSize = random.nextInt(500, 2000); - int dimensions = random.nextInt(64, 256); + int datasetSize = dataProvider.getDatasetSize(); topK = random.nextInt(2, 30); log.log(Level.FINE, "Using topK as: " + topK); numThreads = random.nextInt(2, 8); log.log(Level.FINE, "Generating a dataset with " + datasetSize + " vectors"); - float[][] dataset = generateDataset(random, datasetSize, dimensions); + float[][] dataset = dataProvider.getVectors(datasetSize); numQueries = random.nextInt(100, 500); log.log(Level.FINE, "Generating a query set with " + numQueries + " queries"); - float[][] queryVectors = generateDataset(random, numQueries, dimensions); + float[][] queryVectors = dataProvider.getQueries(numQueries); queries = new ArrayBlockingQueue<>(numQueries, true, Arrays.asList(queryVectors)); log.log(Level.FINE, "Indexing " + datasetSize + " vectors"); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index 05e37f84..1217e492 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -1,14 +1,32 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; +import static org.apache.lucene.tests.util.LuceneTestCase.newIndexWriterConfig; +import static org.apache.lucene.tests.util.LuceneTestCase.newTieredMergePolicy; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.IndexWriterConfig; +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; public class TestUtils { - public static float[][] generateDataset(Random random, int size, int dimensions) { + private static final Logger log = Logger.getLogger(TestUtils.class.getName()); + + public static float[][] generateRandomVectors(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++) { @@ -18,22 +36,59 @@ public static float[][] generateDataset(Random random, int size, int dimensions) 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; + public static List> generateExpectedTopK( + int topK, float[][] dataset, float[][] queries) { + List> neighborsResult = new ArrayList<>(); + int dimensions = dataset[0].length; + + for (float[] query : queries) { + Map distances = new TreeMap<>(); + for (int j = 0; j < dataset.length; j++) { + double distance = 0; + for (int k = 0; k < dimensions; k++) { + distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); + } + distances.put(j, (distance)); + } + + Map sorted = new TreeMap(distances); + log.log(Level.FINER, "EXPECTED: " + sorted); + + // Sort by distance and select the topK nearest neighbors + List neighbors = + distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .toList(); + neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); } - return vector; + + log.log(Level.FINE, "Expected results generated successfully."); + return neighborsResult; } - 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; - } + public static RandomIndexWriter createWriter(Random random, Directory directory, Codec codec) + throws IOException { + return new RandomIndexWriter( + random, + directory, + newIndexWriterConfig(new MockAnalyzer(random, MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + } + + public static IndexWriterConfig createWriterConfig(Random random, Codec codec) { + return newIndexWriterConfig(new MockAnalyzer(random, MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy()); + } + + public static 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 queries; + return sb.toString(); } }