[WIP] [PARKED] Refactor tests

src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java → src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java

@@ -4,15 +4,19 @@
  */
 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;
 
 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.List;
 import java.util.Random;
 import java.util.UUID;
 import java.util.logging.Level;
@@ -35,100 +39,96 @@
 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 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<Integer> expectedIds = new HashSet<Integer>(Arrays.asList(expected));
+      List<List<Integer>> 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,55 +139,45 @@ 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);
     }
   }
 
   @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();
       }
 
       // Verify the index can be opened and searched
       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));
       }
     }
   }