Apache
/
lucene
mirror of https://github.com/apache/lucene.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fixed index order needed for TestKnnVectorQuery.testScoreEuclidean (#11732)

This commit is contained in:
Michael Sokolov 2022-09-01 09:52:48 -04:00 committed by Michael Sokolov
parent 1649964f07
commit 0462a0ad73
1 changed files with 34 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
lucene/core/src/test/org/apache/lucene/search/TestKnnVectorQuery.java
View File

@ -245,7 +245,7 @@ public class TestKnnVectorQuery extends LuceneTestCase {
for (int j = 0; j < 5; j++) { for (int j = 0; j < 5; j++) {
vectors[j] = new float[] {j, j}; vectors[j] = new float[] {j, j};
} }
try (Directory d = getIndexStore("field", 1, vectors); try (Directory d = getStableIndexStore("field", vectors);
IndexReader reader = DirectoryReader.open(d)) { IndexReader reader = DirectoryReader.open(d)) {
IndexSearcher searcher = new IndexSearcher(reader); IndexSearcher searcher = new IndexSearcher(reader);
KnnVectorQuery query = new KnnVectorQuery("field", new float[] {2, 3}, 3); KnnVectorQuery query = new KnnVectorQuery("field", new float[] {2, 3}, 3);
@ -756,13 +756,8 @@ public class TestKnnVectorQuery extends LuceneTestCase {
} }
} }
private Directory getIndexStore(String field, float[]... contents) throws IOException {
return getIndexStore(field, -1, contents);
}
/** Creates a new directory and adds documents with the given vectors as kNN vector fields */ /** Creates a new directory and adds documents with the given vectors as kNN vector fields */
private Directory getIndexStore(String field, int forceMerge, float[]... contents) private Directory getIndexStore(String field, float[]... contents) throws IOException {
throws IOException {
Directory indexStore = newDirectory(); Directory indexStore = newDirectory();
RandomIndexWriter writer = new RandomIndexWriter(random(), indexStore); RandomIndexWriter writer = new RandomIndexWriter(random(), indexStore);
VectorEncoding encoding = randomVectorEncoding(); VectorEncoding encoding = randomVectorEncoding();
@ -786,13 +781,42 @@ public class TestKnnVectorQuery extends LuceneTestCase {
doc.add(new StringField("other", "value", Field.Store.NO)); doc.add(new StringField("other", "value", Field.Store.NO));
writer.addDocument(doc); writer.addDocument(doc);
} }
if (forceMerge > 0) {
writer.forceMerge(forceMerge);
}
writer.close(); writer.close();
return indexStore; return indexStore;
} }
/**
* Creates a new directory and adds documents with the given vectors as kNN vector fields,
* preserving the order of the added documents.
*/
private Directory getStableIndexStore(String field, float[]... contents) throws IOException {
Directory indexStore = newDirectory();
try (IndexWriter writer = new IndexWriter(indexStore, newIndexWriterConfig())) {
VectorEncoding encoding = randomVectorEncoding();
for (int i = 0; i < contents.length; ++i) {
Document doc = new Document();
if (encoding == VectorEncoding.BYTE) {
BytesRef v = new BytesRef(new byte[contents[i].length]);
for (int j = 0; j < v.length; j++) {
v.bytes[j] = (byte) contents[i][j];
}
doc.add(new KnnVectorField(field, v, EUCLIDEAN));
} else {
doc.add(new KnnVectorField(field, contents[i]));
}
doc.add(new StringField("id", "id" + i, Field.Store.YES));
writer.addDocument(doc);
}
// Add some documents without a vector
for (int i = 0; i < 5; i++) {
Document doc = new Document();
doc.add(new StringField("other", "value", Field.Store.NO));
writer.addDocument(doc);
}
}
return indexStore;
}
private void assertMatches(IndexSearcher searcher, Query q, int expectedMatches) private void assertMatches(IndexSearcher searcher, Query q, int expectedMatches)
throws IOException { throws IOException {
ScoreDoc[] result = searcher.search(q, 1000).scoreDocs; ScoreDoc[] result = searcher.search(q, 1000).scoreDocs;