Skip to content

Match merged vectors by document id, not by ordinal position #2556

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
imotov wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Match merged vectors by document id, not by ordinal position #2556

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 8 additions & 21 deletions java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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);
}
}
}
Expand All @@ -85,30 +81,21 @@ 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);
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);
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);
Expand Down
29 changes: 8 additions & 21 deletions ...lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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);
}
}
}
Expand All @@ -85,30 +81,21 @@ 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);
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);
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);
Expand Down
27 changes: 6 additions & 21 deletions java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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);
}
}
}
Expand All @@ -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);
Expand All @@ -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);
}
}
}
Expand Down
55 changes: 55 additions & 0 deletions java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
* <p>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<Integer> 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++) {
Expand Down
Loading