You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ma...@apache.org on 2022/02/17 10:58:55 UTC
[lucene] branch branch_9x updated: LUCENE-10408 Better encoding of doc Ids in vectors (#649)
This is an automated email from the ASF dual-hosted git repository.
mayya pushed a commit to branch branch_9x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_9x by this push:
new 3355273 LUCENE-10408 Better encoding of doc Ids in vectors (#649)
3355273 is described below
commit 3355273630b8396cfd51a770caf6213a9e2fba3f
Author: Mayya Sharipova <ma...@elastic.co>
AuthorDate: Thu Feb 17 11:34:42 2022 +0100
LUCENE-10408 Better encoding of doc Ids in vectors (#649)
Better encoding of doc Ids in Lucene91HnswVectorsFormat
for a dense case where all docs have vectors.
Currently we write doc Ids of all documents that have vectors
not very efficiently.
This improve their encoding by for a case when all documents
have vectors, we don't write document IDs, but just write a
single short value – a dense marker.
---
lucene/CHANGES.txt | 3 +
.../codecs/lucene91/Lucene91HnswVectorsFormat.java | 6 +-
.../codecs/lucene91/Lucene91HnswVectorsReader.java | 97 +++++++++++++++-------
.../codecs/lucene91/Lucene91HnswVectorsWriter.java | 52 ++++++------
.../org/apache/lucene/index/DocsWithFieldSet.java | 14 +++-
.../tests/index/BaseKnnVectorsFormatTestCase.java | 55 ++++++++++++
6 files changed, 166 insertions(+), 61 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 377b2e2..0d02322 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -182,6 +182,9 @@ Optimizations
* LUCENE-10412: More `Query#rewrite` optimizations for MatchNoDocsQuery.
(Adrien Grand)
+* LUCENE-10408 Better encoding of doc Ids in vectors. (Mayya Sharipova, Julie Tibshirani, Adrien Grand)
+
+
Changes in runtime behavior
---------------------
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsFormat.java
index 2e3fc6c..a4f21b4 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsFormat.java
@@ -26,7 +26,7 @@ import org.apache.lucene.index.SegmentWriteState;
import org.apache.lucene.util.hnsw.HnswGraph;
/**
- * Lucene 9.0 vector format, which encodes numeric vector values and an optional associated graph
+ * Lucene 9.1 vector format, which encodes numeric vector values and an optional associated graph
* connecting the documents having values. The graph is used to power HNSW search. The format
* consists of three files:
*
@@ -67,7 +67,9 @@ import org.apache.lucene.util.hnsw.HnswGraph;
* <li><b>[vlong]</b> length of this field's index data, in bytes
* <li><b>[int]</b> dimension of this field's vectors
* <li><b>[int]</b> the number of documents having values for this field
- * <li><b>array[vint]</b> the docids of documents having vectors, in order
+ * <li><b>[int8]</b> if equals to -1, dense – all documents have values for a field. If equals to
+ * 0, sparse – some documents missing values.
+ * <li><b>array[vint]</b> for sparse case, the docids of documents having vectors, in order
* <li><b>[int]</b> the maximum number of connections (neigbours) that each node can have
* <li><b>[int]</b> number of levels in the graph
* <li>Graph nodes by level. For each level
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java
index 85247df..08016da 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java
@@ -24,6 +24,7 @@ import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
+import java.util.function.IntUnaryOperator;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.index.CorruptIndexException;
@@ -244,7 +245,7 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
int node = results.topNode();
float score = fieldEntry.similarityFunction.convertToScore(results.topScore());
results.pop();
- scoreDocs[scoreDocs.length - ++i] = new ScoreDoc(fieldEntry.ordToDoc[node], score);
+ scoreDocs[scoreDocs.length - ++i] = new ScoreDoc(fieldEntry.ordToDoc(node), score);
}
// always return >= the case where we can assert == is only when there are fewer than topK
// vectors in the index
@@ -256,22 +257,26 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
private OffHeapVectorValues getOffHeapVectorValues(FieldEntry fieldEntry) throws IOException {
IndexInput bytesSlice =
vectorData.slice("vector-data", fieldEntry.vectorDataOffset, fieldEntry.vectorDataLength);
- return new OffHeapVectorValues(fieldEntry.dimension, fieldEntry.ordToDoc, bytesSlice);
+ return new OffHeapVectorValues(
+ fieldEntry.dimension, fieldEntry.size(), fieldEntry.ordToDoc, bytesSlice);
}
private Bits getAcceptOrds(Bits acceptDocs, FieldEntry fieldEntry) {
+ if (fieldEntry.ordToDoc == null) {
+ return acceptDocs;
+ }
if (acceptDocs == null) {
return null;
}
return new Bits() {
@Override
public boolean get(int index) {
- return acceptDocs.get(fieldEntry.ordToDoc[index]);
+ return acceptDocs.get(fieldEntry.ordToDoc(index));
}
@Override
public int length() {
- return fieldEntry.ordToDoc.length;
+ return fieldEntry.size;
}
};
}
@@ -311,7 +316,9 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
final int maxConn;
final int numLevels;
final int dimension;
+ private final int size;
final int[] ordToDoc;
+ private final IntUnaryOperator ordToDocOperator;
final int[][] nodesByLevel;
// for each level the start offsets in vectorIndex file from where to read neighbours
final long[] graphOffsetsByLevel;
@@ -323,12 +330,25 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
vectorIndexOffset = input.readVLong();
vectorIndexLength = input.readVLong();
dimension = input.readInt();
- int size = input.readInt();
- ordToDoc = new int[size];
- for (int i = 0; i < size; i++) {
- int doc = input.readVInt();
- ordToDoc[i] = doc;
+ size = input.readInt();
+
+ int denseSparseMarker = input.readByte();
+ if (denseSparseMarker == -1) {
+ ordToDoc = null; // each document has a vector value
+ } else {
+ assert denseSparseMarker == 0;
+ // TODO: Can we read docIDs from disk directly instead of loading giant arrays in memory?
+ // Or possibly switch to something like DirectMonotonicReader if it doesn't slow down
+ // searches.
+
+ // as not all docs have vector values, fill a mapping from dense vector ordinals to docIds
+ ordToDoc = new int[size];
+ for (int i = 0; i < size; i++) {
+ int doc = input.readVInt();
+ ordToDoc[i] = doc;
+ }
}
+ ordToDocOperator = ordToDoc == null ? IntUnaryOperator.identity() : (ord) -> ordToDoc[ord];
// read nodes by level
maxConn = input.readInt();
@@ -363,7 +383,11 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
}
int size() {
- return ordToDoc.length;
+ return size;
+ }
+
+ int ordToDoc(int ord) {
+ return ordToDocOperator.applyAsInt(ord);
}
}
@@ -371,21 +395,24 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
static class OffHeapVectorValues extends VectorValues
implements RandomAccessVectorValues, RandomAccessVectorValuesProducer {
- final int dimension;
- final int[] ordToDoc;
- final IndexInput dataIn;
-
- final BytesRef binaryValue;
- final ByteBuffer byteBuffer;
- final int byteSize;
- final float[] value;
+ private final int dimension;
+ private final int size;
+ private final int[] ordToDoc;
+ private final IntUnaryOperator ordToDocOperator;
+ private final IndexInput dataIn;
+ private final BytesRef binaryValue;
+ private final ByteBuffer byteBuffer;
+ private final int byteSize;
+ private final float[] value;
- int ord = -1;
- int doc = -1;
+ private int ord = -1;
+ private int doc = -1;
- OffHeapVectorValues(int dimension, int[] ordToDoc, IndexInput dataIn) {
+ OffHeapVectorValues(int dimension, int size, int[] ordToDoc, IndexInput dataIn) {
this.dimension = dimension;
+ this.size = size;
this.ordToDoc = ordToDoc;
+ ordToDocOperator = ordToDoc == null ? IntUnaryOperator.identity() : (ord) -> ordToDoc[ord];
this.dataIn = dataIn;
byteSize = Float.BYTES * dimension;
byteBuffer = ByteBuffer.allocate(byteSize);
@@ -400,7 +427,7 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
@Override
public int size() {
- return ordToDoc.length;
+ return size;
}
@Override
@@ -424,10 +451,10 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
@Override
public int nextDoc() {
- if (++ord >= size()) {
+ if (++ord >= size) {
doc = NO_MORE_DOCS;
} else {
- doc = ordToDoc[ord];
+ doc = ordToDocOperator.applyAsInt(ord);
}
return doc;
}
@@ -435,27 +462,33 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
@Override
public int advance(int target) {
assert docID() < target;
- ord = Arrays.binarySearch(ordToDoc, ord + 1, ordToDoc.length, target);
- if (ord < 0) {
- ord = -(ord + 1);
+
+ if (ordToDoc == null) {
+ ord = target;
+ } else {
+ ord = Arrays.binarySearch(ordToDoc, ord + 1, ordToDoc.length, target);
+ if (ord < 0) {
+ ord = -(ord + 1);
+ }
}
- assert ord <= ordToDoc.length;
- if (ord == ordToDoc.length) {
+
+ assert ord <= size;
+ if (ord == size) {
doc = NO_MORE_DOCS;
} else {
- doc = ordToDoc[ord];
+ doc = ordToDocOperator.applyAsInt(ord);
}
return doc;
}
@Override
public long cost() {
- return ordToDoc.length;
+ return size;
}
@Override
public RandomAccessVectorValues randomAccess() {
- return new OffHeapVectorValues(dimension, ordToDoc, dataIn.clone());
+ return new OffHeapVectorValues(dimension, size, ordToDoc, dataIn.clone());
}
@Override
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsWriter.java
index 2b3eeb4..c6ca1d5 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsWriter.java
@@ -24,15 +24,16 @@ import java.util.Arrays;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.KnnVectorsWriter;
+import org.apache.lucene.index.DocsWithFieldSet;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.RandomAccessVectorValuesProducer;
import org.apache.lucene.index.SegmentWriteState;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.index.VectorValues;
+import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.store.IndexOutput;
-import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator;
@@ -49,6 +50,7 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
private final SegmentWriteState segmentWriteState;
private final IndexOutput meta, vectorData, vectorIndex;
+ private final int maxDoc;
private final int maxConn;
private final int beamWidth;
@@ -102,6 +104,7 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
Lucene91HnswVectorsFormat.VERSION_CURRENT,
state.segmentInfo.getId(),
state.segmentSuffix);
+ maxDoc = state.segmentInfo.maxDoc();
success = true;
} finally {
if (success == false) {
@@ -123,8 +126,7 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
boolean success = false;
try {
// write the vector data to a temporary file
- // TODO - use a better data structure; a bitset? DocsWithFieldSet is p.p. in o.a.l.index
- int[] docIds = writeVectorData(tempVectorData, vectors);
+ DocsWithFieldSet docsWithField = writeVectorData(tempVectorData, vectors);
CodecUtil.writeFooter(tempVectorData);
IOUtils.close(tempVectorData);
@@ -138,9 +140,11 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
long vectorIndexOffset = vectorIndex.getFilePointer();
// build the graph using the temporary vector data
+ // we pass null for ordToDoc mapping, for the graph construction doesn't need to know docIds
+ // TODO: separate random access vector values from DocIdSetIterator?
Lucene91HnswVectorsReader.OffHeapVectorValues offHeapVectors =
new Lucene91HnswVectorsReader.OffHeapVectorValues(
- vectors.dimension(), docIds, vectorDataInput);
+ vectors.dimension(), docsWithField.cardinality(), null, vectorDataInput);
OnHeapHnswGraph graph =
offHeapVectors.size() == 0
? null
@@ -152,7 +156,7 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
vectorDataLength,
vectorIndexOffset,
vectorIndexLength,
- docIds,
+ docsWithField,
graph);
success = true;
} finally {
@@ -168,26 +172,19 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
}
/**
- * Writes the vector values to the output and returns a mapping from dense ordinals to document
- * IDs. The length of the returned array matches the total number of documents with a vector
- * (which excludes deleted documents), so it may be less than {@link VectorValues#size()}.
+ * Writes the vector values to the output and returns a set of documents that contains vectors.
*/
- private static int[] writeVectorData(IndexOutput output, VectorValues vectors)
+ private static DocsWithFieldSet writeVectorData(IndexOutput output, VectorValues vectors)
throws IOException {
- int[] docIds = new int[vectors.size()];
- int count = 0;
- for (int docV = vectors.nextDoc(); docV != NO_MORE_DOCS; docV = vectors.nextDoc(), count++) {
+ DocsWithFieldSet docsWithField = new DocsWithFieldSet();
+ for (int docV = vectors.nextDoc(); docV != NO_MORE_DOCS; docV = vectors.nextDoc()) {
// write vector
BytesRef binaryValue = vectors.binaryValue();
assert binaryValue.length == vectors.dimension() * Float.BYTES;
output.writeBytes(binaryValue.bytes, binaryValue.offset, binaryValue.length);
- docIds[count] = docV;
+ docsWithField.add(docV);
}
-
- if (docIds.length > count) {
- return ArrayUtil.copyOfSubArray(docIds, 0, count);
- }
- return docIds;
+ return docsWithField;
}
private void writeMeta(
@@ -196,7 +193,7 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
long vectorDataLength,
long vectorIndexOffset,
long vectorIndexLength,
- int[] docIds,
+ DocsWithFieldSet docsWithField,
OnHeapHnswGraph graph)
throws IOException {
meta.writeInt(field.number);
@@ -206,14 +203,21 @@ public final class Lucene91HnswVectorsWriter extends KnnVectorsWriter {
meta.writeVLong(vectorIndexOffset);
meta.writeVLong(vectorIndexLength);
meta.writeInt(field.getVectorDimension());
- meta.writeInt(docIds.length);
- for (int docId : docIds) {
- // TODO: delta-encode, or write as bitset
- meta.writeVInt(docId);
+
+ // write docIDs
+ int count = docsWithField.cardinality();
+ meta.writeInt(count);
+ if (count == maxDoc) {
+ meta.writeByte((byte) -1); // dense marker, each document has a vector value
+ } else {
+ meta.writeByte((byte) 0); // sparse marker, some documents don't have vector values
+ DocIdSetIterator iter = docsWithField.iterator();
+ for (int doc = iter.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iter.nextDoc()) {
+ meta.writeVInt(doc);
+ }
}
meta.writeInt(maxConn);
-
// write graph nodes on each level
if (graph == null) {
meta.writeInt(0);
diff --git a/lucene/core/src/java/org/apache/lucene/index/DocsWithFieldSet.java b/lucene/core/src/java/org/apache/lucene/index/DocsWithFieldSet.java
index b4f5ac5..1d018ac 100644
--- a/lucene/core/src/java/org/apache/lucene/index/DocsWithFieldSet.java
+++ b/lucene/core/src/java/org/apache/lucene/index/DocsWithFieldSet.java
@@ -26,7 +26,7 @@ import org.apache.lucene.util.RamUsageEstimator;
* Accumulator for documents that have a value for a field. This is optimized for the case that all
* documents have a value.
*/
-final class DocsWithFieldSet extends DocIdSet {
+public final class DocsWithFieldSet extends DocIdSet {
private static long BASE_RAM_BYTES_USED =
RamUsageEstimator.shallowSizeOfInstance(DocsWithFieldSet.class);
@@ -35,7 +35,15 @@ final class DocsWithFieldSet extends DocIdSet {
private int cardinality = 0;
private int lastDocId = -1;
- void add(int docID) {
+ /** Creates an empty DocsWithFieldSet. */
+ public DocsWithFieldSet() {}
+
+ /**
+ * Add a document to the set
+ *
+ * @param docID – document ID to be added
+ */
+ public void add(int docID) {
if (docID <= lastDocId) {
throw new IllegalArgumentException(
"Out of order doc ids: last=" + lastDocId + ", next=" + docID);
@@ -64,7 +72,7 @@ final class DocsWithFieldSet extends DocIdSet {
}
/** Return the number of documents of this set. */
- int cardinality() {
+ public int cardinality() {
return cardinality;
}
}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java
index 59be16d..72e6ebb 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java
@@ -21,6 +21,7 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
+import java.util.Set;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.document.Document;
@@ -28,6 +29,7 @@ import org.apache.lucene.document.Field;
import org.apache.lucene.document.FieldType;
import org.apache.lucene.document.KnnVectorField;
import org.apache.lucene.document.NumericDocValuesField;
+import org.apache.lucene.document.StoredField;
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.CheckIndex;
import org.apache.lucene.index.CodecReader;
@@ -1018,4 +1020,57 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
}
}
}
+
+ public void testVectorValuesReportCorrectDocs() throws Exception {
+ final int numDocs = atLeast(1000);
+ final int dim = random().nextInt(20) + 1;
+ final VectorSimilarityFunction similarityFunction =
+ VectorSimilarityFunction.values()[
+ random().nextInt(VectorSimilarityFunction.values().length)];
+
+ float fieldValuesCheckSum = 0f;
+ int fieldDocCount = 0;
+ long fieldSumDocIDs = 0;
+
+ try (Directory dir = newDirectory();
+ RandomIndexWriter w = new RandomIndexWriter(random(), dir, newIndexWriterConfig())) {
+ for (int i = 0; i < numDocs; i++) {
+ Document doc = new Document();
+ int docID = random().nextInt(numDocs);
+ doc.add(new StoredField("id", docID));
+ if (random().nextInt(4) == 3) {
+ float[] vector = randomVector(dim);
+ doc.add(new KnnVectorField("knn_vector", vector, similarityFunction));
+ fieldValuesCheckSum += vector[0];
+ fieldDocCount++;
+ fieldSumDocIDs += docID;
+ }
+ w.addDocument(doc);
+ }
+
+ if (random().nextBoolean()) {
+ w.forceMerge(1);
+ }
+
+ try (IndexReader r = w.getReader()) {
+ float checksum = 0;
+ int docCount = 0;
+ long sumDocIds = 0;
+ for (LeafReaderContext ctx : r.leaves()) {
+ VectorValues vectors = ctx.reader().getVectorValues("knn_vector");
+ if (vectors != null) {
+ docCount += vectors.size();
+ while (vectors.nextDoc() != NO_MORE_DOCS) {
+ checksum += vectors.vectorValue()[0];
+ Document doc = ctx.reader().document(vectors.docID(), Set.of("id"));
+ sumDocIds += Integer.parseInt(doc.get("id"));
+ }
+ }
+ }
+ assertEquals(fieldValuesCheckSum, checksum, 1e-3);
+ assertEquals(fieldDocCount, docCount);
+ assertEquals(fieldSumDocIDs, sumDocIds);
+ }
+ }
+ }
}