You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by be...@apache.org on 2024/08/07 12:38:46 UTC
(lucene) branch main updated: Add float|byte vector support to memory index (#13633)
This is an automated email from the ASF dual-hosted git repository.
benwtrent pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new 9e831ee8097 Add float|byte vector support to memory index (#13633)
9e831ee8097 is described below
commit 9e831ee809729054fb01d80e4f5dc736e91e732a
Author: Benjamin Trent <be...@gmail.com>
AuthorDate: Wed Aug 7 08:38:41 2024 -0400
Add float|byte vector support to memory index (#13633)
* Add float|byte vector support to memory index
* adding changes
---
lucene/CHANGES.txt | 2 +
.../apache/lucene/index/memory/MemoryIndex.java | 238 ++++++++++++++++++++-
.../lucene/index/memory/TestMemoryIndex.java | 170 +++++++++++++++
3 files changed, 408 insertions(+), 2 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 265c339145e..f9b989d723f 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -290,6 +290,8 @@ Improvements
* GITHUB#13285: Early terminate graph searches of AbstractVectorSimilarityQuery to follow timeout set from
IndexSearcher#setTimeout(QueryTimeout). (Kaival Parikh)
+
+* GITHUB#13633: Add ability to read/write knn vector values to a MemoryIndex. (Ben Trent)
Optimizations
---------------------
diff --git a/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java b/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java
index 08fb1cf6b5b..9dad312d431 100644
--- a/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java
+++ b/lucene/memory/src/java/org/apache/lucene/index/memory/MemoryIndex.java
@@ -36,15 +36,19 @@ import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.analysis.tokenattributes.TermToBytesRefAttribute;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.FieldType;
+import org.apache.lucene.document.KnnByteVectorField;
+import org.apache.lucene.document.KnnFloatVectorField;
import org.apache.lucene.index.*;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
+import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorable;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.SimpleCollector;
+import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.search.similarities.Similarity;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.ArrayUtil;
@@ -636,6 +640,10 @@ public class MemoryIndex {
if (field.fieldType().stored()) {
storeValues(info, field);
}
+
+ if (field.fieldType().vectorDimension() > 0) {
+ storeVectorValues(info, field);
+ }
}
/**
@@ -749,6 +757,56 @@ public class MemoryIndex {
info.pointValues[info.pointValuesCount++] = BytesRef.deepCopyOf(pointValue);
}
+ private void storeVectorValues(Info info, IndexableField vectorField) {
+ assert vectorField instanceof KnnFloatVectorField || vectorField instanceof KnnByteVectorField;
+ switch (info.fieldInfo.getVectorEncoding()) {
+ case BYTE -> {
+ if (vectorField instanceof KnnByteVectorField byteVectorField) {
+ if (info.byteVectorCount == 1) {
+ throw new IllegalArgumentException(
+ "Only one value per field allowed for byte vector field ["
+ + vectorField.name()
+ + "]");
+ }
+ info.byteVectorCount++;
+ if (info.byteVectorValues == null) {
+ info.byteVectorValues = new byte[1][];
+ }
+ info.byteVectorValues[0] =
+ ArrayUtil.copyOfSubArray(
+ byteVectorField.vectorValue(), 0, info.fieldInfo.getVectorDimension());
+ return;
+ }
+ throw new IllegalArgumentException(
+ "Field ["
+ + vectorField.name()
+ + "] is not a byte vector field, but the field info is configured for byte vectors");
+ }
+ case FLOAT32 -> {
+ if (vectorField instanceof KnnFloatVectorField floatVectorField) {
+ if (info.floatVectorCount == 1) {
+ throw new IllegalArgumentException(
+ "Only one value per field allowed for float vector field ["
+ + vectorField.name()
+ + "]");
+ }
+ info.floatVectorCount++;
+ if (info.floatVectorValues == null) {
+ info.floatVectorValues = new float[1][];
+ }
+ info.floatVectorValues[0] =
+ ArrayUtil.copyOfSubArray(
+ floatVectorField.vectorValue(), 0, info.fieldInfo.getVectorDimension());
+ return;
+ }
+ throw new IllegalArgumentException(
+ "Field ["
+ + vectorField.name()
+ + "] is not a float vector field, but the field info is configured for float vectors");
+ }
+ }
+ }
+
private void storeValues(Info info, IndexableField field) {
if (info.storedValues == null) {
info.storedValues = new ArrayList<>();
@@ -1148,6 +1206,18 @@ public class MemoryIndex {
private BytesRef[] pointValues;
+ /** Number of float vectors added for this field */
+ private int floatVectorCount;
+
+ /** the float vectors added for this field */
+ private float[][] floatVectorValues;
+
+ /** Number of byte vectors added for this field */
+ private int byteVectorCount;
+
+ /** the byte vectors added for this field */
+ private byte[][] byteVectorValues;
+
private byte[] minPackedValue;
private byte[] maxPackedValue;
@@ -1641,12 +1711,20 @@ public class MemoryIndex {
@Override
public FloatVectorValues getFloatVectorValues(String fieldName) {
- return null;
+ Info info = fields.get(fieldName);
+ if (info == null || info.floatVectorValues == null) {
+ return null;
+ }
+ return new MemoryFloatVectorValues(info);
}
@Override
public ByteVectorValues getByteVectorValues(String fieldName) {
- return null;
+ Info info = fields.get(fieldName);
+ if (info == null || info.byteVectorValues == null) {
+ return null;
+ }
+ return new MemoryByteVectorValues(info);
}
@Override
@@ -2204,4 +2282,160 @@ public class MemoryIndex {
return super.clear();
}
}
+
+ private static final class MemoryFloatVectorValues extends FloatVectorValues {
+ private final Info info;
+ private int currentDoc = -1;
+
+ MemoryFloatVectorValues(Info info) {
+ this.info = info;
+ }
+
+ @Override
+ public int dimension() {
+ return info.fieldInfo.getVectorDimension();
+ }
+
+ @Override
+ public int size() {
+ return info.floatVectorCount;
+ }
+
+ @Override
+ public float[] vectorValue() {
+ if (currentDoc == 0) {
+ return info.floatVectorValues[0];
+ } else {
+ return null;
+ }
+ }
+
+ @Override
+ public VectorScorer scorer(float[] query) {
+ if (query.length != info.fieldInfo.getVectorDimension()) {
+ throw new IllegalArgumentException(
+ "query vector dimension "
+ + query.length
+ + " does not match field dimension "
+ + info.fieldInfo.getVectorDimension());
+ }
+ MemoryFloatVectorValues vectorValues = new MemoryFloatVectorValues(info);
+ return new VectorScorer() {
+ @Override
+ public float score() throws IOException {
+ return info.fieldInfo
+ .getVectorSimilarityFunction()
+ .compare(vectorValues.vectorValue(), query);
+ }
+
+ @Override
+ public DocIdSetIterator iterator() {
+ return vectorValues;
+ }
+ };
+ }
+
+ @Override
+ public int docID() {
+ return currentDoc;
+ }
+
+ @Override
+ public int nextDoc() {
+ int doc = ++currentDoc;
+ if (doc == 0) {
+ return doc;
+ } else {
+ return NO_MORE_DOCS;
+ }
+ }
+
+ @Override
+ public int advance(int target) {
+ if (target == 0) {
+ currentDoc = target;
+ return target;
+ } else {
+ return NO_MORE_DOCS;
+ }
+ }
+ }
+
+ private static final class MemoryByteVectorValues extends ByteVectorValues {
+ private final Info info;
+ private int currentDoc = -1;
+
+ MemoryByteVectorValues(Info info) {
+ this.info = info;
+ }
+
+ @Override
+ public int dimension() {
+ return info.fieldInfo.getVectorDimension();
+ }
+
+ @Override
+ public int size() {
+ return info.byteVectorCount;
+ }
+
+ @Override
+ public byte[] vectorValue() {
+ if (currentDoc == 0) {
+ return info.byteVectorValues[0];
+ } else {
+ return null;
+ }
+ }
+
+ @Override
+ public VectorScorer scorer(byte[] query) {
+ if (query.length != info.fieldInfo.getVectorDimension()) {
+ throw new IllegalArgumentException(
+ "query vector dimension "
+ + query.length
+ + " does not match field dimension "
+ + info.fieldInfo.getVectorDimension());
+ }
+ MemoryByteVectorValues vectorValues = new MemoryByteVectorValues(info);
+ return new VectorScorer() {
+ @Override
+ public float score() {
+ return info.fieldInfo
+ .getVectorSimilarityFunction()
+ .compare(vectorValues.vectorValue(), query);
+ }
+
+ @Override
+ public DocIdSetIterator iterator() {
+ return vectorValues;
+ }
+ };
+ }
+
+ @Override
+ public int docID() {
+ return currentDoc;
+ }
+
+ @Override
+ public int nextDoc() {
+ int doc = ++currentDoc;
+ if (doc == 0) {
+ return doc;
+ } else {
+ return NO_MORE_DOCS;
+ }
+ }
+
+ @Override
+ public int advance(int target) {
+ if (target == 0) {
+ currentDoc = target;
+ return target;
+ } else {
+ return NO_MORE_DOCS;
+ }
+ }
+ }
}
diff --git a/lucene/memory/src/test/org/apache/lucene/index/memory/TestMemoryIndex.java b/lucene/memory/src/test/org/apache/lucene/index/memory/TestMemoryIndex.java
index b40409f1f3b..cf5628b9873 100644
--- a/lucene/memory/src/test/org/apache/lucene/index/memory/TestMemoryIndex.java
+++ b/lucene/memory/src/test/org/apache/lucene/index/memory/TestMemoryIndex.java
@@ -43,6 +43,8 @@ import org.apache.lucene.document.FloatPoint;
import org.apache.lucene.document.IntField;
import org.apache.lucene.document.IntPoint;
import org.apache.lucene.document.InvertableType;
+import org.apache.lucene.document.KnnByteVectorField;
+import org.apache.lucene.document.KnnFloatVectorField;
import org.apache.lucene.document.LongPoint;
import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.document.SortedDocValuesField;
@@ -53,8 +55,10 @@ import org.apache.lucene.document.StoredValue;
import org.apache.lucene.document.StringField;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.BinaryDocValues;
+import org.apache.lucene.index.ByteVectorValues;
import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.FieldInvertState;
+import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexableField;
@@ -68,6 +72,7 @@ import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
+import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.search.CollectionStatistics;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
@@ -76,6 +81,7 @@ import org.apache.lucene.search.PhraseQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.TermStatistics;
+import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.search.similarities.BM25Similarity;
import org.apache.lucene.search.similarities.ClassicSimilarity;
import org.apache.lucene.search.similarities.Similarity;
@@ -741,6 +747,170 @@ public class TestMemoryIndex extends LuceneTestCase {
d, "multibinary", new BytesRef[] {new BytesRef("bbar"), new BytesRef("bbaz")});
}
+ public void testKnnFloatVectorOnlyOneVectorAllowed() throws IOException {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField("knnFloatA", new float[] {1.0f, 2.0f}));
+ doc.add(new KnnFloatVectorField("knnFloatA", new float[] {3.0f, 4.0f}));
+ expectThrows(
+ IllegalArgumentException.class,
+ () -> MemoryIndex.fromDocument(doc, new StandardAnalyzer()));
+ }
+
+ public void testKnnFloatVectors() throws IOException {
+ List<IndexableField> fields = new ArrayList<>();
+ fields.add(new KnnFloatVectorField("knnFloatA", new float[] {1.0f, 2.0f}));
+ fields.add(new KnnFloatVectorField("knnFloatB", new float[] {3.0f, 4.0f, 5.0f, 6.0f}));
+ fields.add(
+ new KnnFloatVectorField(
+ "knnFloatC", new float[] {7.0f, 8.0f, 9.0f}, VectorSimilarityFunction.DOT_PRODUCT));
+ Collections.shuffle(fields, random());
+ Document doc = new Document();
+ for (IndexableField f : fields) {
+ doc.add(f);
+ }
+
+ MemoryIndex mi = MemoryIndex.fromDocument(doc, new StandardAnalyzer());
+ assertFloatVectorValue(mi, "knnFloatA", new float[] {1.0f, 2.0f});
+ assertFloatVectorValue(mi, "knnFloatB", new float[] {3.0f, 4.0f, 5.0f, 6.0f});
+ assertFloatVectorValue(mi, "knnFloatC", new float[] {7.0f, 8.0f, 9.0f});
+
+ assertFloatVectorScore(mi, "knnFloatA", new float[] {1.0f, 1.0f}, 0.5f);
+ assertFloatVectorScore(mi, "knnFloatB", new float[] {3.0f, 3.0f, 3.0f, 3.0f}, 0.06666667f);
+ assertFloatVectorScore(mi, "knnFloatC", new float[] {7.0f, 7.0f, 7.0f}, 84.5f);
+
+ assertNull(
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getFloatVectorValues("knnFloatMissing"));
+ assertNull(
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getByteVectorValues("knnByteVectorValue"));
+ }
+
+ public void testKnnByteVectorOnlyOneVectorAllowed() throws IOException {
+ Document doc = new Document();
+ doc.add(new KnnByteVectorField("knnByteA", new byte[] {1, 2}));
+ doc.add(new KnnByteVectorField("knnByteA", new byte[] {3, 4}));
+ expectThrows(
+ IllegalArgumentException.class,
+ () -> MemoryIndex.fromDocument(doc, new StandardAnalyzer()));
+ }
+
+ public void testKnnByteVectors() throws IOException {
+ List<IndexableField> fields = new ArrayList<>();
+ fields.add(new KnnByteVectorField("knnByteA", new byte[] {1, 2}));
+ fields.add(new KnnByteVectorField("knnByteB", new byte[] {3, 4, 5, 6}));
+ fields.add(
+ new KnnByteVectorField(
+ "knnByteC", new byte[] {7, 8, 9}, VectorSimilarityFunction.DOT_PRODUCT));
+ Collections.shuffle(fields, random());
+ Document doc = new Document();
+ for (IndexableField f : fields) {
+ doc.add(f);
+ }
+
+ MemoryIndex mi = MemoryIndex.fromDocument(doc, new StandardAnalyzer());
+ assertByteVectorValue(mi, "knnByteA", new byte[] {1, 2});
+ assertByteVectorValue(mi, "knnByteB", new byte[] {3, 4, 5, 6});
+ assertByteVectorValue(mi, "knnByteC", new byte[] {7, 8, 9});
+
+ assertByteVectorScore(mi, "knnByteA", new byte[] {1, 1}, 0.5f);
+ assertByteVectorScore(mi, "knnByteB", new byte[] {3, 3, 3, 3}, 0.06666667f);
+ assertByteVectorScore(mi, "knnByteC", new byte[] {7, 7, 7}, 0.501709f);
+
+ assertNull(
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getByteVectorValues("knnByteMissing"));
+ assertNull(
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getFloatVectorValues("knnFloatVectorValue"));
+ }
+
+ private static void assertFloatVectorValue(MemoryIndex mi, String fieldName, float[] expected)
+ throws IOException {
+ FloatVectorValues fvv =
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getFloatVectorValues(fieldName);
+ assertNotNull(fvv);
+ assertEquals(0, fvv.nextDoc());
+ assertArrayEquals(expected, fvv.vectorValue(), 1e-6f);
+ assertEquals(DocIdSetIterator.NO_MORE_DOCS, fvv.nextDoc());
+ }
+
+ private static void assertFloatVectorScore(
+ MemoryIndex mi, String fieldName, float[] queryVector, float expectedScore)
+ throws IOException {
+ FloatVectorValues fvv =
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getFloatVectorValues(fieldName);
+ assertNotNull(fvv);
+ if (random().nextBoolean()) {
+ fvv.nextDoc();
+ }
+ VectorScorer scorer = fvv.scorer(queryVector);
+ assertEquals(0, scorer.iterator().nextDoc());
+ assertEquals(expectedScore, scorer.score(), 0.0f);
+ assertEquals(DocIdSetIterator.NO_MORE_DOCS, scorer.iterator().nextDoc());
+ }
+
+ private static void assertByteVectorValue(MemoryIndex mi, String fieldName, byte[] expected)
+ throws IOException {
+ ByteVectorValues bvv =
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getByteVectorValues(fieldName);
+ assertNotNull(bvv);
+ assertEquals(0, bvv.nextDoc());
+ assertArrayEquals(expected, bvv.vectorValue());
+ assertEquals(DocIdSetIterator.NO_MORE_DOCS, bvv.nextDoc());
+ }
+
+ private static void assertByteVectorScore(
+ MemoryIndex mi, String fieldName, byte[] queryVector, float expectedScore)
+ throws IOException {
+ ByteVectorValues bvv =
+ mi.createSearcher()
+ .getIndexReader()
+ .leaves()
+ .get(0)
+ .reader()
+ .getByteVectorValues(fieldName);
+ assertNotNull(bvv);
+ if (random().nextBoolean()) {
+ bvv.nextDoc();
+ }
+ VectorScorer scorer = bvv.scorer(queryVector);
+ assertEquals(0, scorer.iterator().nextDoc());
+ assertEquals(expectedScore, scorer.score(), 0.0f);
+ assertEquals(DocIdSetIterator.NO_MORE_DOCS, scorer.iterator().nextDoc());
+ }
+
private static void assertContains(
Document d, String field, Object expected, Function<IndexableField, Object> value) {
assertNotNull(d.getField(field));