diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 2b960bf3cd..6db380a3ce 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -14,7 +15,6 @@ 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.LeafReader; import org.apache.lucene.index.LeafReaderContext; @@ -44,12 +44,12 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); w.addDocument(doc2); w.flush(); @@ -69,11 +69,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertVectorsKeepTheirDocuments(r, "f", f); } } } @@ -85,12 +81,12 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); w.addDocument(doc2); @@ -98,17 +94,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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); + assertVectorsKeepTheirDocuments(r, "f1", f1); + assertVectorsKeepTheirDocuments(r, "f2", f2); // opportunistically check boundary condition - search with a 0 topK var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 9bb140228b..707ee35709 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -14,7 +15,6 @@ 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.LeafReader; import org.apache.lucene.index.LeafReaderContext; @@ -44,12 +44,12 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); w.addDocument(doc2); w.flush(); @@ -69,11 +69,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertVectorsKeepTheirDocuments(r, "f", f); } } } @@ -85,12 +81,12 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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 StringField("id", "0", Field.Store.YES)); 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 StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); w.addDocument(doc2); @@ -98,17 +94,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { 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); + assertVectorsKeepTheirDocuments(r, "f1", f1); + assertVectorsKeepTheirDocuments(r, "f2", f2); // opportunistically check boundary condition - search with a 0 topK var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 5140782fc7..3b05beaacf 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.COSINE; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -77,7 +78,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { for (int i = 0; i < R; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(F, f[i], EUCLIDEAN)); w.addDocument(doc); w.commit(); @@ -95,12 +96,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues(F); - assertNotNull(values); - assertEquals(R, values.size()); - for (int i = 0; i < R; i++) { - assertArrayEquals(f[i], values.vectorValue(i), 0.0f); - } + assertVectorsKeepTheirDocuments(r, F, f); } } } @@ -121,7 +117,7 @@ public void testTwoVectorFieldsPerDoc() throws Exception { for (int i = 0; i < R; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(F1, f1[i], EUCLIDEAN)); doc.add(new KnnFloatVectorField(F2, f2[i], EUCLIDEAN)); w.addDocument(doc); @@ -130,19 +126,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues(F1); - assertNotNull(values); - assertEquals(R, values.size()); - for (int i = 0; i < R; i++) { - assertArrayEquals(f1[i], values.vectorValue(i), 0.0f); - } - - values = r.getFloatVectorValues(F2); - assertNotNull(values); - assertEquals(R, values.size()); - for (int i = 0; i < R; i++) { - assertArrayEquals(f2[i], values.vectorValue(i), 0.0f); - } + assertVectorsKeepTheirDocuments(r, F1, f1); + assertVectorsKeepTheirDocuments(r, F2, f2); } } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index 88ca3150fd..22cd509e7b 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -4,10 +4,65 @@ */ package com.nvidia.cuvs.lucene; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +import java.util.HashSet; import java.util.Random; +import java.util.Set; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.StoredFields; public class TestUtils { + /** + * Asserts that every vector in {@code field} is still paired with the document that indexed it, + * matching on the document's stored {@code id} rather than on ordinal position. + * + *

Nothing fixes the order documents land in after a merge, and the randomized test + * framework's {@code MockRandomMergePolicy} disturbs it two independent ways: {@code + * findForcedMerges} shuffles the segments it is about to merge, and {@code + * MockRandomOneMerge.reorder} reverses doc IDs outright. The second applies even when there is + * only one segment, so a test that never commits between documents is no safer than one that + * does. Asserting {@code vectorValue(i) == expectedById[i]} therefore fails on some seeds with + * nothing wrong. What has to hold is that no document loses its own vector, which is what this + * checks. + * + * @param expectedById the vector indexed for each document id, indexed by that id + */ + public static void assertVectorsKeepTheirDocuments( + LeafReader reader, String field, float[][] expectedById) throws IOException { + FloatVectorValues values = reader.getFloatVectorValues(field); + assertNotNull("no vector values for field " + field, values); + assertEquals(expectedById.length, values.size()); + StoredFields storedFields = reader.storedFields(); + Set seen = new HashSet<>(); + int previousDoc = -1; + for (int ord = 0; ord < values.size(); ord++) { + int doc = values.ordToDoc(ord); + // Ordinals are assigned in docID order. Where every document has a vector -- every caller + // today -- ordToDoc is the hardcoded identity and this cannot fire. It earns its keep only + // if a caller passes a field that some documents lack, where the id lookup below would + // otherwise be free to agree with a garbled mapping. + assertTrue("ordToDoc went backwards at ordinal " + ord, doc > previousDoc); + previousDoc = doc; + String storedId = storedFields.document(doc).get("id"); + assertNotNull("document at ordinal " + ord + " has no stored id", storedId); + int id = Integer.parseInt(storedId); + assertTrue("document id " + id + " appeared twice", seen.add(id)); + assertArrayEquals( + "vector for document id " + id + " (doc " + doc + ", ordinal " + ord + ")", + expectedById[id], + values.vectorValue(ord), + 0.0f); + } + assertEquals("not every document was found", expectedById.length, seen.size()); + } + public static float[][] generateDataset(Random random, int size, int dimensions) { float[][] dataset = new float[size][dimensions]; for (int i = 0; i < size; i++) {