You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ch...@apache.org on 2025/04/15 11:00:58 UTC
(lucene) branch branch_10x updated: [10.x] Add support for determining off-heap memory requirements for KnnVectorsReader (#14426) (#14497)
This is an automated email from the ASF dual-hosted git repository.
chegar pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new 547e4aff018 [10.x] Add support for determining off-heap memory requirements for KnnVectorsReader (#14426) (#14497)
547e4aff018 is described below
commit 547e4aff018249d21c5570e86eababf151b34dd6
Author: Chris Hegarty <62...@users.noreply.github.com>
AuthorDate: Tue Apr 15 12:00:52 2025 +0100
[10.x] Add support for determining off-heap memory requirements for KnnVectorsReader (#14426) (#14497)
This PR adds support to KnnVectorsReader in order to determine the off-heap memory requirements.
The motivation here is to give better insight into the size of off-heap memory that will be needed, so that deployments can be better scaled so that vector search workloads fit in memory, in order to provide best execution performance.
---
.../lucene90/Lucene90HnswVectorsReader.java | 11 +++
.../lucene91/Lucene91HnswVectorsReader.java | 11 +++
.../lucene92/Lucene92HnswVectorsReader.java | 11 +++
.../lucene94/Lucene94HnswVectorsReader.java | 18 +++-
.../lucene95/Lucene95HnswVectorsReader.java | 20 ++++
...stLucene99HnswScalarQuantizedVectorsFormat.java | 41 +++++++++
.../simpletext/SimpleTextKnnVectorsReader.java | 12 +++
.../org/apache/lucene/codecs/KnnVectorsFormat.java | 7 ++
.../org/apache/lucene/codecs/KnnVectorsReader.java | 52 +++++++++++
.../Lucene102BinaryQuantizedVectorsReader.java | 16 ++++
.../codecs/lucene99/Lucene99FlatVectorsReader.java | 18 +++-
.../codecs/lucene99/Lucene99HnswVectorsReader.java | 20 +++-
.../Lucene99ScalarQuantizedVectorsReader.java | 15 +++
.../codecs/perfield/PerFieldKnnVectorsFormat.java | 6 ++
.../lucene/index/SlowCodecReaderWrapper.java | 23 +++++
.../index/SlowCompositeCodecReaderWrapper.java | 9 ++
.../apache/lucene/index/SortingCodecReader.java | 5 +
...tLucene102HnswBinaryQuantizedVectorsFormat.java | 34 +++++++
.../lucene99/TestLucene99HnswVectorsFormat.java | 37 ++++++++
.../asserting/AssertingKnnVectorsFormat.java | 12 ++-
.../tests/index/BaseKnnVectorsFormatTestCase.java | 101 +++++++++++++++++++++
21 files changed, 470 insertions(+), 9 deletions(-)
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java
index a9021780a60..8a78f93e8d3 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java
@@ -17,9 +17,12 @@
package org.apache.lucene.backward_codecs.lucene90;
+import static org.apache.lucene.backward_codecs.lucene90.Lucene90HnswVectorsFormat.VECTOR_DATA_EXTENSION;
+import static org.apache.lucene.backward_codecs.lucene90.Lucene90HnswVectorsFormat.VECTOR_INDEX_EXTENSION;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
+import java.util.Map;
import java.util.SplittableRandom;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
@@ -306,6 +309,14 @@ public final class Lucene90HnswVectorsReader extends KnnVectorsReader {
return new OffHeapHnswGraph(entry, bytesSlice);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ FieldEntry entry = getFieldEntry(fieldInfo.name);
+ var raw = Map.entry(VECTOR_DATA_EXTENSION, entry.vectorDataLength);
+ var graph = Map.entry(VECTOR_INDEX_EXTENSION, entry.indexDataLength);
+ return Map.ofEntries(raw, graph);
+ }
+
@Override
public void close() throws IOException {
IOUtils.close(vectorData, vectorIndex);
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java
index 18c105a3e01..472c2a95944 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java
@@ -17,10 +17,13 @@
package org.apache.lucene.backward_codecs.lucene91;
+import static org.apache.lucene.backward_codecs.lucene91.Lucene91HnswVectorsFormat.VECTOR_DATA_EXTENSION;
+import static org.apache.lucene.backward_codecs.lucene91.Lucene91HnswVectorsFormat.VECTOR_INDEX_EXTENSION;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
import java.util.Arrays;
+import java.util.Map;
import java.util.function.IntUnaryOperator;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
@@ -297,6 +300,14 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
return new OffHeapHnswGraph(entry, bytesSlice);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ FieldEntry entry = getFieldEntry(fieldInfo.name);
+ var raw = Map.entry(VECTOR_DATA_EXTENSION, entry.vectorDataLength);
+ var graph = Map.entry(VECTOR_INDEX_EXTENSION, entry.vectorIndexLength);
+ return Map.ofEntries(raw, graph);
+ }
+
@Override
public void close() throws IOException {
IOUtils.close(vectorData, vectorIndex);
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java
index 1b9ff59dec5..df46cf2f276 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java
@@ -17,10 +17,13 @@
package org.apache.lucene.backward_codecs.lucene92;
+import static org.apache.lucene.backward_codecs.lucene92.Lucene92HnswVectorsFormat.VECTOR_DATA_EXTENSION;
+import static org.apache.lucene.backward_codecs.lucene92.Lucene92HnswVectorsFormat.VECTOR_INDEX_EXTENSION;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
import java.util.Arrays;
+import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
@@ -263,6 +266,14 @@ public final class Lucene92HnswVectorsReader extends KnnVectorsReader {
return new OffHeapHnswGraph(entry, bytesSlice);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ FieldEntry entry = getFieldEntry(fieldInfo.name);
+ var raw = Map.entry(VECTOR_DATA_EXTENSION, entry.vectorDataLength);
+ var graph = Map.entry(VECTOR_INDEX_EXTENSION, entry.vectorIndexLength);
+ return Map.ofEntries(raw, graph);
+ }
+
@Override
public void close() throws IOException {
IOUtils.close(vectorData, vectorIndex);
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java
index 61ad082da16..6bf791a309c 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java
@@ -17,10 +17,13 @@
package org.apache.lucene.backward_codecs.lucene94;
+import static org.apache.lucene.backward_codecs.lucene94.Lucene94HnswVectorsFormat.VECTOR_DATA_EXTENSION;
+import static org.apache.lucene.backward_codecs.lucene94.Lucene94HnswVectorsFormat.VECTOR_INDEX_EXTENSION;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
import java.util.Arrays;
+import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
@@ -231,12 +234,17 @@ public final class Lucene94HnswVectorsReader extends KnnVectorsReader {
CodecUtil.checksumEntireFile(vectorIndex);
}
- private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
+ private FieldEntry getFieldEntryOrThrow(String field) {
final FieldInfo info = fieldInfos.fieldInfo(field);
final FieldEntry fieldEntry;
if (info == null || (fieldEntry = fields.get(info.number)) == null) {
throw new IllegalArgumentException("field=\"" + field + "\" not found");
}
+ return fieldEntry;
+ }
+
+ private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
+ final FieldEntry fieldEntry = getFieldEntryOrThrow(field);
if (fieldEntry.vectorEncoding != expectedEncoding) {
throw new IllegalArgumentException(
"field=\""
@@ -305,6 +313,14 @@ public final class Lucene94HnswVectorsReader extends KnnVectorsReader {
return new OffHeapHnswGraph(entry, bytesSlice);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ FieldEntry entry = getFieldEntryOrThrow(fieldInfo.name);
+ var raw = Map.entry(VECTOR_DATA_EXTENSION, entry.vectorDataLength);
+ var graph = Map.entry(VECTOR_INDEX_EXTENSION, entry.vectorIndexLength);
+ return Map.ofEntries(raw, graph);
+ }
+
@Override
public void close() throws IOException {
IOUtils.close(vectorData, vectorIndex);
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java
index 3c68a8a6353..d8ce482907a 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java
@@ -17,10 +17,13 @@
package org.apache.lucene.backward_codecs.lucene95;
+import static org.apache.lucene.backward_codecs.lucene95.Lucene95HnswVectorsFormat.VECTOR_DATA_EXTENSION;
+import static org.apache.lucene.backward_codecs.lucene95.Lucene95HnswVectorsFormat.VECTOR_INDEX_EXTENSION;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
import java.util.Arrays;
+import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
@@ -237,6 +240,15 @@ public final class Lucene95HnswVectorsReader extends KnnVectorsReader implements
CodecUtil.checksumEntireFile(vectorIndex);
}
+ private FieldEntry getFieldEntryOrThrow(String field) {
+ final FieldInfo info = fieldInfos.fieldInfo(field);
+ final FieldEntry fieldEntry;
+ if (info == null || (fieldEntry = fields.get(info.number)) == null) {
+ throw new IllegalArgumentException("field=\"" + field + "\" not found");
+ }
+ return fieldEntry;
+ }
+
private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
final FieldInfo info = fieldInfos.fieldInfo(field);
final FieldEntry fieldEntry;
@@ -358,6 +370,14 @@ public final class Lucene95HnswVectorsReader extends KnnVectorsReader implements
return new OffHeapHnswGraph(entry, vectorIndex);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ FieldEntry entry = getFieldEntryOrThrow(fieldInfo.name);
+ var raw = Map.entry(VECTOR_DATA_EXTENSION, entry.vectorDataLength);
+ var graph = Map.entry(VECTOR_INDEX_EXTENSION, entry.vectorIndexLength);
+ return Map.ofEntries(raw, graph);
+ }
+
@Override
public void close() throws IOException {
IOUtils.close(vectorData, vectorIndex);
diff --git a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene99/TestLucene99HnswScalarQuantizedVectorsFormat.java b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene99/TestLucene99HnswScalarQuantizedVectorsFormat.java
index 268d231c4f4..d7af766ee07 100644
--- a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene99/TestLucene99HnswScalarQuantizedVectorsFormat.java
+++ b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene99/TestLucene99HnswScalarQuantizedVectorsFormat.java
@@ -17,7 +17,20 @@
package org.apache.lucene.backward_codecs.lucene99;
+import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
+
+import java.io.IOException;
import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase;
import org.apache.lucene.tests.util.TestUtil;
@@ -26,4 +39,32 @@ public class TestLucene99HnswScalarQuantizedVectorsFormat extends BaseKnnVectors
protected Codec getCodec() {
return TestUtil.alwaysKnnVectorsFormat(new Lucene99RWHnswScalarQuantizationVectorsFormat());
}
+
+ public void testSimpleOffHeapSize() throws IOException {
+ float[] vector = randomVector(random().nextInt(12, 500));
+ try (Directory dir = newDirectory();
+ IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT));
+ w.addDocument(doc);
+ w.commit();
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ if (r instanceof CodecReader codecReader) {
+ KnnVectorsReader knnVectorsReader = codecReader.getVectorReader();
+ if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
+ knnVectorsReader = fieldsReader.getFieldReader("f");
+ }
+ var fieldInfo = r.getFieldInfos().fieldInfo("f");
+ var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
+ assertEquals(vector.length * Float.BYTES, (long) offHeap.get("vec"));
+ assertEquals(1L, (long) offHeap.get("vex"));
+ long corrections = Float.BYTES;
+ long expected = fieldInfo.getVectorDimension() + corrections;
+ assertEquals(expected, (long) offHeap.get("veq"));
+ assertEquals(3, offHeap.size());
+ }
+ }
+ }
+ }
}
diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java
index 6c7c53a38d0..a1bc4fc5257 100644
--- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java
+++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java
@@ -26,6 +26,8 @@ import static org.apache.lucene.codecs.simpletext.SimpleTextKnnVectorsWriter.VEC
import java.io.IOException;
import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.Objects;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.index.ByteVectorValues;
import org.apache.lucene.index.CorruptIndexException;
@@ -275,6 +277,16 @@ public class SimpleTextKnnVectorsReader extends KnnVectorsReader {
}
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ Objects.requireNonNull(fieldInfo);
+ FieldEntry fieldEntry = fieldEntries.get(fieldInfo.number);
+ if (fieldEntry == null) {
+ return null;
+ }
+ return Map.of(); // all in-heap
+ }
+
@Override
public void close() throws IOException {
dataIn.close();
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsFormat.java
index ad6e4aba607..99439f38780 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsFormat.java
@@ -18,8 +18,10 @@
package org.apache.lucene.codecs;
import java.io.IOException;
+import java.util.Map;
import java.util.Set;
import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.SegmentReadState;
import org.apache.lucene.index.SegmentWriteState;
@@ -148,6 +150,11 @@ public abstract class KnnVectorsFormat implements NamedSPILoader.NamedSPI {
throw new UnsupportedOperationException();
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ return Map.of();
+ }
+
@Override
public void close() {}
};
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java
index 54d070fa599..e15fe1b5b3a 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java
@@ -19,6 +19,9 @@ package org.apache.lucene.codecs;
import java.io.Closeable;
import java.io.IOException;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
import org.apache.lucene.index.ByteVectorValues;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.FloatVectorValues;
@@ -26,6 +29,7 @@ import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TotalHits;
+import org.apache.lucene.util.Accountable;
import org.apache.lucene.util.Bits;
/** Reads vectors from an index. */
@@ -130,4 +134,52 @@ public abstract class KnnVectorsReader implements Closeable {
* <p>The default implementation is empty
*/
public void finishMerge() throws IOException {}
+
+ /**
+ * Returns the desired size of off-heap memory the given field. This size can be used to help
+ * determine the memory requirements for optimal search performance, which can be greatly affected
+ * by page faults when not enough memory is available.
+ *
+ * <p>For reporting purposes, the size of the off-heap index structures is broken down by their
+ * file extension, which provides a logical categorization of their purpose, e.g. the {@code
+ * Lucene99HnswVectorsFormat} stores the HNSW graph neighbours lists in a file with the "vex"
+ * extension.
+ *
+ * <p>The long value is the size in bytes of the off-heap space needed if the associated index
+ * structure were to be fully loaded in memory. While somewhat analogous to {@link
+ * Accountable#ramBytesUsed()} (which reports actual on-heap memory usage), the sizes reported by
+ * this method are not actual usage but rather the amount of available memory needed to fully load
+ * the index into memory, rather than an actual RAM usage requirement.
+ *
+ * <p>To determine the total desired off-heap memory size for the given field:
+ *
+ * <pre>{@code
+ * getOffHeapByteSize(field).values().stream().mapToLong(Long::longValue).sum();
+ * }</pre>
+ *
+ * <p>The default implementation returns an empty map.
+ *
+ * @param fieldInfo the fieldInfo
+ * @return a map of the desired off-heap memory requirements by category
+ * @lucene.experimental
+ */
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ return Map.of();
+ }
+
+ /**
+ * Merges the Maps returned by {@link #getOffHeapByteSize(FieldInfo)}.
+ *
+ * <p>This method is a convenience for aggregating the desired off-heap memory requirements for
+ * several fields. The keys in the returned map are a union of the keys in the given maps. Entries
+ * with the same key are summed.
+ *
+ * @lucene.experimental
+ */
+ public static Map<String, Long> mergeOffHeapByteSizeMaps(
+ Map<String, Long> map1, Map<String, Long> map2) {
+ return Stream.of(map1, map2)
+ .flatMap(map -> map.entrySet().stream())
+ .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, Long::sum));
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102BinaryQuantizedVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102BinaryQuantizedVectorsReader.java
index a698ea9a754..87d57b2da84 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102BinaryQuantizedVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102BinaryQuantizedVectorsReader.java
@@ -16,6 +16,7 @@
*/
package org.apache.lucene.codecs.lucene102;
+import static org.apache.lucene.codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat.VECTOR_DATA_EXTENSION;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.discretize;
@@ -23,7 +24,9 @@ import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.discr
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
+import java.util.Objects;
import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
import org.apache.lucene.index.ByteVectorValues;
@@ -257,6 +260,19 @@ class Lucene102BinaryQuantizedVectorsReader extends FlatVectorsReader {
return size;
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ Objects.requireNonNull(fieldInfo);
+ var raw = rawVectorsReader.getOffHeapByteSize(fieldInfo);
+ var fieldEntry = fields.get(fieldInfo.name);
+ if (fieldEntry == null) {
+ assert fieldInfo.getVectorEncoding() == VectorEncoding.BYTE;
+ return raw;
+ }
+ var quant = Map.of(VECTOR_DATA_EXTENSION, fieldEntry.vectorDataLength());
+ return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, quant);
+ }
+
public float[] getCentroid(String field) {
FieldEntry fieldEntry = fields.get(field);
if (fieldEntry != null) {
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsReader.java
index 52839000137..62bf6ef3f16 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsReader.java
@@ -22,6 +22,7 @@ import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVe
import java.io.IOException;
import java.io.UncheckedIOException;
+import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
@@ -166,6 +167,12 @@ public final class Lucene99FlatVectorsReader extends FlatVectorsReader {
return Lucene99FlatVectorsReader.SHALLOW_SIZE + fields.ramBytesUsed();
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ final FieldEntry entry = getFieldEntryOrThrow(fieldInfo.name);
+ return Map.of(Lucene99FlatVectorsFormat.VECTOR_DATA_EXTENSION, entry.vectorDataLength());
+ }
+
@Override
public void checkIntegrity() throws IOException {
CodecUtil.checksumEntireFile(vectorData);
@@ -182,12 +189,17 @@ public final class Lucene99FlatVectorsReader extends FlatVectorsReader {
}
}
- private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
+ private FieldEntry getFieldEntryOrThrow(String field) {
final FieldInfo info = fieldInfos.fieldInfo(field);
- final FieldEntry fieldEntry;
- if (info == null || (fieldEntry = fields.get(info.number)) == null) {
+ final FieldEntry entry;
+ if (info == null || (entry = fields.get(info.number)) == null) {
throw new IllegalArgumentException("field=\"" + field + "\" not found");
}
+ return entry;
+ }
+
+ private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
+ final FieldEntry fieldEntry = getFieldEntryOrThrow(field);
if (fieldEntry.vectorEncoding != expectedEncoding) {
throw new IllegalArgumentException(
"field=\""
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
index 14838c8dd01..21ab010e528 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
@@ -22,6 +22,7 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
+import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
@@ -264,12 +265,17 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
return flatVectorsReader.getByteVectorValues(field);
}
- private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
+ private FieldEntry getFieldEntryOrThrow(String field) {
final FieldInfo info = fieldInfos.fieldInfo(field);
- final FieldEntry fieldEntry;
- if (info == null || (fieldEntry = fields.get(info.number)) == null) {
+ final FieldEntry entry;
+ if (info == null || (entry = fields.get(info.number)) == null) {
throw new IllegalArgumentException("field=\"" + field + "\" not found");
}
+ return entry;
+ }
+
+ private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
+ final FieldEntry fieldEntry = getFieldEntryOrThrow(field);
if (fieldEntry.vectorEncoding != expectedEncoding) {
throw new IllegalArgumentException(
"field=\""
@@ -367,6 +373,14 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
return new OffHeapHnswGraph(entry, vectorIndex);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ FieldEntry entry = getFieldEntryOrThrow(fieldInfo.name);
+ var flat = flatVectorsReader.getOffHeapByteSize(fieldInfo);
+ var graph = Map.of(Lucene99HnswVectorsFormat.VECTOR_INDEX_EXTENSION, entry.vectorIndexLength);
+ return KnnVectorsReader.mergeOffHeapByteSizeMaps(flat, graph);
+ }
+
@Override
public void close() throws IOException {
IOUtils.close(flatVectorsReader, vectorIndex);
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99ScalarQuantizedVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99ScalarQuantizedVectorsReader.java
index 712e9b91f9d..81ad8f46d68 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99ScalarQuantizedVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99ScalarQuantizedVectorsReader.java
@@ -19,9 +19,12 @@ package org.apache.lucene.codecs.lucene99;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
+import static org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat.VECTOR_DATA_EXTENSION;
import java.io.IOException;
+import java.util.Map;
import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
@@ -282,6 +285,18 @@ public final class Lucene99ScalarQuantizedVectorsReader extends FlatVectorsReade
return SHALLOW_SIZE + fields.ramBytesUsed() + rawVectorsReader.ramBytesUsed();
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ var raw = rawVectorsReader.getOffHeapByteSize(fieldInfo);
+ var fieldEntry = fields.get(fieldInfo.number);
+ if (fieldEntry == null) {
+ assert fieldInfo.getVectorEncoding() == VectorEncoding.BYTE;
+ return raw;
+ }
+ var quant = Map.of(VECTOR_DATA_EXTENSION, fieldEntry.vectorDataLength());
+ return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, quant);
+ }
+
private FieldEntry readField(IndexInput input, int versionMeta, FieldInfo info)
throws IOException {
VectorEncoding vectorEncoding = readVectorEncoding(input);
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/perfield/PerFieldKnnVectorsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/perfield/PerFieldKnnVectorsFormat.java
index bc18b231e74..884e9210c64 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/perfield/PerFieldKnnVectorsFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/perfield/PerFieldKnnVectorsFormat.java
@@ -322,6 +322,12 @@ public abstract class PerFieldKnnVectorsFormat extends KnnVectorsFormat {
reader.search(field, target, knnCollector, acceptDocs);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ KnnVectorsReader knnVectorsReader = fields.get(fieldInfo.number);
+ return knnVectorsReader.getOffHeapByteSize(fieldInfo);
+ }
+
@Override
public void close() throws IOException {
List<KnnVectorsReader> readers = new ArrayList<>(fields.size());
diff --git a/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
index 57836ae482f..d7ed3d2f06f 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
@@ -21,6 +21,7 @@ import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
+import java.util.Map;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
import org.apache.lucene.codecs.KnnVectorsReader;
@@ -28,6 +29,7 @@ import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
import org.apache.lucene.codecs.TermVectorsReader;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.util.Bits;
@@ -189,6 +191,27 @@ public final class SlowCodecReaderWrapper {
// We already checkIntegrity the entire reader up front
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ SegmentReader segmentReader = segmentReader(reader);
+ var vectorsReader = segmentReader.getVectorReader();
+ if (vectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
+ vectorsReader = fieldsReader.getFieldReader(fieldInfo.name);
+ }
+ return vectorsReader.getOffHeapByteSize(fieldInfo);
+ }
+
+ static SegmentReader segmentReader(LeafReader reader) {
+ if (reader instanceof SegmentReader) {
+ return (SegmentReader) reader;
+ } else if (reader instanceof final FilterLeafReader fReader) {
+ return segmentReader(FilterLeafReader.unwrap(fReader));
+ } else if (reader instanceof final FilterCodecReader fReader) {
+ return segmentReader(FilterCodecReader.unwrap(fReader));
+ }
+ throw new AssertionError("unexpected reader [" + reader + "]");
+ }
+
@Override
public void close() {}
};
diff --git a/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
index 48f4e19033e..7cb12a271ab 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
@@ -852,6 +852,15 @@ final class SlowCompositeCodecReaderWrapper extends CodecReader {
return new MergedFloatVectorValues(dimension, size, subs);
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ Map<String, Long> map = new HashMap<>();
+ for (var reader : readers) {
+ map = KnnVectorsReader.mergeOffHeapByteSizeMaps(map, reader.getOffHeapByteSize(fieldInfo));
+ }
+ return map;
+ }
+
class MergedFloatVectorValues extends FloatVectorValues {
final int dimension;
final int size;
diff --git a/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java b/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java
index ab9964026ad..5a24f7782ac 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java
@@ -590,6 +590,11 @@ public final class SortingCodecReader extends FilterCodecReader {
throw new UnsupportedOperationException();
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ return delegate.getOffHeapByteSize(fieldInfo);
+ }
+
@Override
public void close() throws IOException {
delegate.close();
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene102/TestLucene102HnswBinaryQuantizedVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene102/TestLucene102HnswBinaryQuantizedVectorsFormat.java
index 4d15d2a6c52..cadfd9439ed 100644
--- a/lucene/core/src/test/org/apache/lucene/codecs/lucene102/TestLucene102HnswBinaryQuantizedVectorsFormat.java
+++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene102/TestLucene102HnswBinaryQuantizedVectorsFormat.java
@@ -17,18 +17,24 @@
package org.apache.lucene.codecs.lucene102;
import static java.lang.String.format;
+import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.discretize;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.oneOf;
+import java.io.IOException;
import java.util.Arrays;
import java.util.Locale;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.FilterCodec;
import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.index.CodecReader;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.IndexReader;
@@ -137,4 +143,32 @@ public class TestLucene102HnswBinaryQuantizedVectorsFormat extends BaseKnnVector
var expectedValues = Arrays.stream(VectorSimilarityFunction.values()).toList();
assertEquals(Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS, expectedValues);
}
+
+ public void testSimpleOffHeapSize() throws IOException {
+ float[] vector = randomVector(random().nextInt(12, 500));
+ try (Directory dir = newDirectory();
+ IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT));
+ w.addDocument(doc);
+ w.commit();
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ if (r instanceof CodecReader codecReader) {
+ KnnVectorsReader knnVectorsReader = codecReader.getVectorReader();
+ if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
+ knnVectorsReader = fieldsReader.getFieldReader("f");
+ }
+ var fieldInfo = r.getFieldInfos().fieldInfo("f");
+ var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
+ assertEquals(vector.length * Float.BYTES, (long) offHeap.get("vec"));
+ assertEquals(1L, (long) offHeap.get("vex"));
+ long corrections = Float.BYTES + Float.BYTES + Float.BYTES + Short.BYTES;
+ long expected = discretize(fieldInfo.getVectorDimension(), 64) / Byte.SIZE + corrections;
+ assertEquals(expected, (long) offHeap.get("veb"));
+ assertEquals(3, offHeap.size());
+ }
+ }
+ }
+ }
}
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene99/TestLucene99HnswVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene99/TestLucene99HnswVectorsFormat.java
index aea32b0a13d..dccc51c35a7 100644
--- a/lucene/core/src/test/org/apache/lucene/codecs/lucene99/TestLucene99HnswVectorsFormat.java
+++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene99/TestLucene99HnswVectorsFormat.java
@@ -17,13 +17,25 @@
package org.apache.lucene.codecs.lucene99;
import static java.lang.String.format;
+import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.oneOf;
+import java.io.IOException;
import java.util.Locale;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.FilterCodec;
import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase;
import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.SameThreadExecutorService;
@@ -61,4 +73,29 @@ public class TestLucene99HnswVectorsFormat extends BaseKnnVectorsFormatTestCase
IllegalArgumentException.class,
() -> new Lucene99HnswVectorsFormat(20, 100, 1, new SameThreadExecutorService()));
}
+
+ public void testSimpleOffHeapSize() throws IOException {
+ float[] vector = randomVector(random().nextInt(12, 500));
+ try (Directory dir = newDirectory();
+ IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT));
+ w.addDocument(doc);
+ w.commit();
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ if (r instanceof CodecReader codecReader) {
+ KnnVectorsReader knnVectorsReader = codecReader.getVectorReader();
+ if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
+ knnVectorsReader = fieldsReader.getFieldReader("f");
+ }
+ var fieldInfo = r.getFieldInfos().fieldInfo("f");
+ var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
+ assertEquals(vector.length * Float.BYTES, (long) offHeap.get("vec"));
+ assertEquals(1L, (long) offHeap.get("vex"));
+ assertEquals(2, offHeap.size());
+ }
+ }
+ }
+ }
}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java b/lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java
index bd2911fc50a..4d471ceb715 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java
@@ -18,6 +18,7 @@
package org.apache.lucene.tests.codecs.asserting;
import java.io.IOException;
+import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.lucene.codecs.KnnFieldVectorsWriter;
import org.apache.lucene.codecs.KnnVectorsFormat;
@@ -103,8 +104,10 @@ public class AssertingKnnVectorsFormat extends KnnVectorsFormat {
}
}
- static class AssertingKnnVectorsReader extends KnnVectorsReader implements HnswGraphProvider {
- final KnnVectorsReader delegate;
+ /** Wraps a AssertingKnnVectorsReader providing additional assertions. */
+ public static class AssertingKnnVectorsReader extends KnnVectorsReader
+ implements HnswGraphProvider {
+ public final KnnVectorsReader delegate;
final FieldInfos fis;
final boolean mergeInstance;
AtomicInteger mergeInstanceCount = new AtomicInteger();
@@ -213,6 +216,11 @@ public class AssertingKnnVectorsFormat extends KnnVectorsFormat {
finishMergeCount.incrementAndGet();
}
+ @Override
+ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+ return delegate.getOffHeapByteSize(fieldInfo);
+ }
+
@Override
public void close() throws IOException {
assert !mergeInstance;
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 2bbeca0eee9..617f2d604e3 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
@@ -25,6 +25,7 @@ import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
+import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;
@@ -34,7 +35,9 @@ import org.apache.lucene.codecs.KnnFieldVectorsWriter;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.KnnVectorsWriter;
+import org.apache.lucene.codecs.hnsw.HnswGraphProvider;
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
+import org.apache.lucene.codecs.simpletext.SimpleTextKnnVectorsReader;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.KnnByteVectorField;
@@ -77,6 +80,7 @@ import org.apache.lucene.search.TotalHits;
import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
+import org.apache.lucene.tests.codecs.asserting.AssertingKnnVectorsFormat;
import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.BytesRef;
@@ -85,6 +89,7 @@ import org.apache.lucene.util.InfoStream;
import org.apache.lucene.util.StringHelper;
import org.apache.lucene.util.VectorUtil;
import org.apache.lucene.util.Version;
+import org.apache.lucene.util.quantization.QuantizedVectorsReader;
import org.junit.Before;
/**
@@ -450,6 +455,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
assertEquals(0, iterator.nextDoc());
assertEquals(0, vectorValues.vectorValue(0)[0], 0);
assertEquals(NO_MORE_DOCS, iterator.nextDoc());
+ assertOffHeapByteSize(r, fieldName);
}
}
}
@@ -1066,6 +1072,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
assertNotNull(values);
assertEquals(0, values.size());
assertNull(values.scorer(new float[] {2, 3, 5, 6}));
+ assertOffHeapByteSize(r, "v");
}
}
}
@@ -1092,6 +1099,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
assertNotNull(values);
assertEquals(0, values.size());
assertNull(values.scorer(new byte[] {2, 3, 5, 6}));
+ assertOffHeapByteSize(r, "v");
}
}
}
@@ -1484,6 +1492,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
fieldName, randomNormalizedVector(dimension), k, liveDocs, visitedLimit);
assertEquals(TotalHits.Relation.EQUAL_TO, results.totalHits.relation());
assertTrue(results.totalHits.value() <= visitedLimit);
+ assertOffHeapByteSize(ctx.reader(), fieldName);
}
}
}
@@ -1567,6 +1576,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
for (int i = 0; i < k - 1; i++) {
assertTrue(results.scoreDocs[i].score >= results.scoreDocs[i + 1].score);
}
+ assertOffHeapByteSize(ctx.reader(), fieldName);
}
}
}
@@ -1760,6 +1770,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
i = iter.docID();
}
}
+ assertOffHeapByteSize(r, fieldName);
}
}
}
@@ -1828,6 +1839,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
storedFields.document(byteVectorValues.ordToDoc(ord), Set.of("id"));
sumOrdToDocIds += Integer.parseInt(doc.get("id"));
}
+ assertOffHeapByteSize(ctx.reader(), "knn_vector");
}
}
}
@@ -1848,6 +1860,7 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
Document doc = storedFields.document(vectorValues.ordToDoc(ord), Set.of("id"));
sumOrdToDocIds += Integer.parseInt(doc.get("id"));
}
+ assertOffHeapByteSize(ctx.reader(), "knn_vector");
}
}
}
@@ -1914,4 +1927,92 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
IOUtils.close(reader, w2, dir1, dir2);
}
+
+ static void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOException {
+ var fieldInfo = r.getFieldInfos().fieldInfo(fieldName);
+
+ if (r instanceof CodecReader codecReader) {
+ KnnVectorsReader knnVectorsReader = codecReader.getVectorReader();
+ if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
+ knnVectorsReader = fieldsReader.getFieldReader(fieldName);
+ }
+ var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
+ long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();
+ if (knnVectorsReader instanceof SimpleTextKnnVectorsReader) {
+ assertEquals(0L, offHeap.size()); // all vectors are in memory
+ assertEquals(0L, totalByteSize);
+ } else {
+ if (getNumVectors(knnVectorsReader, fieldInfo) == 0) {
+ assertEquals(0L, totalByteSize);
+ } else {
+ assertTrue(totalByteSize > 0);
+ assertTrue(offHeap.get("vec") > 0L);
+
+ if (hasHNSW(knnVectorsReader, fieldInfo)) {
+ assertTrue(offHeap.get("vex") > 0L);
+ var quant = offHeap.getOrDefault("veq", offHeap.get("veb"));
+ assertTrue(quant == null || quant > 0L);
+ } else {
+ assertTrue(offHeap.get("vex") == null);
+ }
+
+ if (hasQuantized(knnVectorsReader, fieldInfo)) {
+ assertTrue(offHeap.getOrDefault("veq", offHeap.get("veb")) > 0L);
+ }
+ }
+ }
+ } else {
+ throw new AssertionError("unexpected:" + r.getClass());
+ }
+ }
+
+ static int getNumVectors(KnnVectorsReader reader, FieldInfo fieldInfo) throws IOException {
+ return switch (fieldInfo.getVectorEncoding()) {
+ case BYTE -> reader.getByteVectorValues(fieldInfo.getName()).size();
+ case FLOAT32 -> reader.getFloatVectorValues(fieldInfo.getName()).size();
+ };
+ }
+
+ static boolean hasQuantized(KnnVectorsReader knnVectorsReader, FieldInfo fieldInfo)
+ throws IOException {
+ if (fieldInfo.getVectorEncoding() == VectorEncoding.BYTE) {
+ return false; // byte vectors are never auto-quantized
+ }
+ if (knnVectorsReader
+ instanceof AssertingKnnVectorsFormat.AssertingKnnVectorsReader assertingReader) {
+ knnVectorsReader = assertingReader.delegate;
+ }
+
+ if (knnVectorsReader instanceof QuantizedVectorsReader quantReader) {
+ return quantReader.getQuantizedVectorValues(fieldInfo.name) != null;
+ }
+
+ String name = knnVectorsReader.getClass().getSimpleName().toLowerCase(Locale.ROOT);
+ return name.contains("quantized");
+ }
+
+ static boolean hasHNSW(KnnVectorsReader knnVectorsReader, FieldInfo fieldInfo)
+ throws IOException {
+ if (knnVectorsReader
+ instanceof AssertingKnnVectorsFormat.AssertingKnnVectorsReader assertingReader) {
+ knnVectorsReader = assertingReader.delegate;
+ }
+ if (knnVectorsReader instanceof HnswGraphProvider graphProvider) {
+ return graphProvider.getGraph(fieldInfo.name) != null;
+ }
+ String name = knnVectorsReader.getClass().getSimpleName().toLowerCase(Locale.ROOT);
+ return name.contains("hnsw");
+ }
+
+ public void testMergeOffHeapByteSizeMaps() {
+ Map<String, Long> map1 = Map.of("a", 0L, "b", 1L, "c", 100L);
+ Map<String, Long> map2 = Map.of("b", 2L, "c", 5L, "d", 101L, "e", 6L);
+ var r = KnnVectorsReader.mergeOffHeapByteSizeMaps(map1, map2);
+ assertEquals(r.size(), 5);
+ assertEquals(0L, (long) r.get("a"));
+ assertEquals(3L, (long) r.get("b"));
+ assertEquals(105L, (long) r.get("c"));
+ assertEquals(101L, (long) r.get("d"));
+ assertEquals(6L, (long) r.get("e"));
+ }
}