john-mlika commented on code in PR #16030:
URL: https://github.com/apache/lucene/pull/16030#discussion_r4056984060
##########
lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java:
##########
@@ -583,4 +588,844 @@ private void updateVectorMetadataFile(Directory dir,
String fileName) throws Exc
CodecUtil.writeFooter(out);
}
}
+
+ /** Returns a codec that always uses the data-blind variant of this format.
*/
+ private Codec dataBlindCodec() {
+ return TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(
+ encoding,
Lucene104ScalarQuantizedVectorsFormat.Mode.DATA_BLIND_WITHOUT_FLOATS));
+ }
+
+ /**
+ * Returns a codec that always uses the data-blind variant of this format
that still stores
+ * full-precision float vectors.
+ */
+ private Codec dataBlindWithFloatsCodec() {
+ return TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(
+ encoding,
Lucene104ScalarQuantizedVectorsFormat.Mode.DATA_BLIND_WITH_FLOATS));
+ }
+
+ public void testDataBlindSearchCorrectness() throws Exception {
+ String fieldName = "field";
+ int numVectors = random().nextInt(99, 500);
+ int dims = random().nextInt(4, 65);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ KnnFloatVectorField knnField =
+ new KnnFloatVectorField(fieldName, randomVector(dims),
similarityFunction);
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ for (int i = 0; i < numVectors; i++) {
+ Document doc = new Document();
+ knnField.setVectorValue(randomVector(dims));
+ doc.add(knnField);
+ w.addDocument(doc);
+ }
+ w.commit();
+
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ IndexSearcher searcher = new IndexSearcher(reader);
+ final int k = random().nextInt(5, 50);
+ float[] queryVector = randomVector(dims);
+ TopDocs collectedDocs =
+ searcher.search(new KnnFloatVectorQuery(fieldName, queryVector,
k), k);
+ assertEquals(k, collectedDocs.totalHits.value());
+ assertEquals(TotalHits.Relation.EQUAL_TO,
collectedDocs.totalHits.relation());
+ }
+ }
+ }
+ }
+
+ /**
+ * Data-blind segments never write full-precision float vectors, so {@link
FloatVectorValues} must
+ * be a bare dequantizing view rather than one backed by raw vectors.
+ */
+ public void testDataBlindNoRawFloatVectors() throws Exception {
+ String fieldName = "field";
+ int numVectors = random().nextInt(4, 50);
+ int dims = random().nextInt(4, 65);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ for (int i = 0; i < numVectors; i++) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField(fieldName, randomVector(dims),
similarityFunction));
+ w.addDocument(doc);
+ }
+ }
+ try (IndexReader reader = DirectoryReader.open(dir)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ FloatVectorValues vectorValues = r.getFloatVectorValues(fieldName);
+ assertEquals(numVectors, vectorValues.size());
+ assertFalse(
+ vectorValues
+ instanceof
Lucene104ScalarQuantizedVectorsReader.ScalarQuantizedVectorValues);
+ }
+ }
+ }
+
+ public void testDataBlindMultiSegmentMerge() throws Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = random().nextInt(4, 50);
+ int dims = random().nextInt(4, 65);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ for (int s = 0; s < 2; s++) {
+ for (int i = 0; i < numVectorsPerSegment; i++) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField(fieldName, randomVector(dims),
similarityFunction));
+ w.addDocument(doc);
+ }
+ w.commit();
+ }
+ w.forceMerge(1);
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ assertEquals(1, reader.leaves().size());
+ LeafReader r = getOnlyLeafReader(reader);
+ assertEquals(2 * numVectorsPerSegment,
r.getFloatVectorValues(fieldName).size());
+ int k = random().nextInt(5, 20);
+ IndexSearcher searcher = new IndexSearcher(reader);
+ TopDocs td =
+ searcher.search(new KnnFloatVectorQuery(fieldName,
randomVector(dims), k), k);
+ // The index holds exactly 2 * numVectorsPerSegment vectors; when k
exceeds that, the
+ // query can only return the vectors that exist.
+ assertEquals(Math.min(k, 2 * numVectorsPerSegment),
td.totalHits.value());
+ }
+ }
+ }
+ }
+
+ /**
+ * Merging a data-blind (quantized-only) segment together with a segment
that still stores raw
+ * floats must pass the data-blind segment's already-quantized bytes
straight through, rather than
+ * dequantizing and re-quantizing them (which would only add loss).
+ */
+ public void testDataBlindMixedMergeKeepsQuantizedBytes() throws Exception {
+ assertDataBlindMixedMergeKeepsQuantizedBytes(false);
+ }
+
+ /** fp16 counterpart of {@link #testDataBlindMixedMergeKeepsQuantizedBytes}.
*/
+ public void testDataBlindFloat16MixedMergeKeepsQuantizedBytes() throws
Exception {
+ assertDataBlindMixedMergeKeepsQuantizedBytes(true);
+ }
+
+ private void assertDataBlindMixedMergeKeepsQuantizedBytes(boolean
float16Data) throws Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = random().nextInt(4, 50);
+ int dims = float16Data ? 2 * random().nextInt(2, 33) : random().nextInt(4,
65);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ // One segment with no raw floats (data-blind) and one written by the
centered writer (keeps
+ // raw floats). Segment-name order between the two writes is not
asserted.
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setUseCompoundFile(false)
+ .setCodec(dataBlindCodec()))) {
+ addVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment, float16Data);
+ w.commit();
+ }
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setUseCompoundFile(false)
+ .setCodec(
+ TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(encoding,
Mode.CENTERED))))) {
+ addVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment, float16Data);
+ w.commit();
+ }
+ // Capture the data-blind segment's stored quantized bytes and
corrective terms.
+ List<byte[]> sourceQuantized = new ArrayList<>();
+ List<OptimizedScalarQuantizer.QuantizationResult> sourceCorrections =
new ArrayList<>();
+ try (DirectoryReader reader = DirectoryReader.open(dir)) {
+ // Exactly one leaf must be data-blind (no raw floats); the other
keeps raw floats.
+ boolean foundDataBlindLeaf = false;
+ for (LeafReaderContext leaf : reader.leaves()) {
+ Lucene104ScalarQuantizedVectorsReader vectorsReader =
+ (Lucene104ScalarQuantizedVectorsReader)
+ ((CodecReader)
leaf.reader()).getVectorReader().unwrapReaderForField(fieldName);
+ boolean hasRaw =
+ float16Data
+ ? vectorsReader.hasRawFloat16Vectors(fieldName)
+ : vectorsReader.hasRawFloatVectors(fieldName);
+ if (hasRaw == false) {
+ foundDataBlindLeaf = true;
+ captureQuantized(
+ vectorsReader.getQuantizedVectorValues(fieldName),
+ sourceQuantized,
+ sourceCorrections);
+ }
+ }
+ assertTrue("expected one data-blind segment", foundDataBlindLeaf);
+ }
+ assertEquals(numVectorsPerSegment, sourceQuantized.size());
+ // Merge both segments through a data-blind writer.
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergeScheduler(new SerialMergeScheduler())
+ .setCodec(dataBlindCodec()))) {
+ w.forceMerge(1);
+ try (DirectoryReader reader = DirectoryReader.open(w)) {
+ assertEquals(1, reader.leaves().size());
+ Lucene104ScalarQuantizedVectorsReader vectorsReader =
+ (Lucene104ScalarQuantizedVectorsReader)
+ ((CodecReader) getOnlyLeafReader(reader))
+ .getVectorReader()
+ .unwrapReaderForField(fieldName);
+ QuantizedByteVectorValues qvv =
vectorsReader.getQuantizedVectorValues(fieldName);
+ assertEquals(2 * numVectorsPerSegment, qvv.size());
+ List<byte[]> mergedVectors = new ArrayList<>();
+ List<OptimizedScalarQuantizer.QuantizationResult> mergedCorrections
= new ArrayList<>();
+ captureQuantized(qvv, mergedVectors, mergedCorrections);
+ // The data-blind segment's vectors must appear contiguously and
untouched (the pass
+ // through preserves byte and correction terms exactly;
re-quantizing them would not).
+ // The merged block position is not assertable, so accept either a
leading or trailing
+ // block.
+ boolean leading =
+ quantizedBlockEquals(
+ sourceQuantized, sourceCorrections, mergedVectors,
mergedCorrections, 0);
+ boolean trailing =
+ quantizedBlockEquals(
+ sourceQuantized,
+ sourceCorrections,
+ mergedVectors,
+ mergedCorrections,
+ numVectorsPerSegment);
+ assertTrue(
+ "data-blind quantized bytes must be passed through untouched",
leading || trailing);
+ }
+ }
+ }
+ }
+
+ private void addVectorDocs(
+ IndexWriter w,
+ String field,
+ int dims,
+ VectorSimilarityFunction similarity,
+ int count,
+ boolean float16Data)
+ throws IOException {
+ if (float16Data) {
+ addFloat16VectorDocs(w, field, dims, similarity, count);
+ } else {
+ addFloatVectorDocs(w, field, dims, similarity, count);
+ }
+ }
+
+ private static void captureQuantized(
+ QuantizedByteVectorValues qvv,
+ List<byte[]> vectors,
+ List<OptimizedScalarQuantizer.QuantizationResult> corrections)
+ throws IOException {
+ KnnVectorValues.DocIndexIterator it = qvv.iterator();
+ for (int doc = it.nextDoc(); doc != NO_MORE_DOCS; doc = it.nextDoc()) {
+ vectors.add(qvv.vectorValue(it.index()).clone());
+ corrections.add(qvv.getCorrectiveTerms(it.index()));
+ }
+ }
+
+ private static boolean quantizedBlockEquals(
+ List<byte[]> expectedVectors,
+ List<OptimizedScalarQuantizer.QuantizationResult> expectedCorrections,
+ List<byte[]> actualVectors,
+ List<OptimizedScalarQuantizer.QuantizationResult> actualCorrections,
+ int offset) {
+ if (offset + expectedVectors.size() > actualVectors.size()) {
+ return false;
+ }
+ for (int i = 0; i < expectedVectors.size(); i++) {
+ if (Arrays.equals(expectedVectors.get(i), actualVectors.get(offset + i))
== false) {
+ return false;
+ }
+ var expected = expectedCorrections.get(i);
+ var actual = actualCorrections.get(offset + i);
+ if (Float.floatToIntBits(expected.lowerInterval())
+ != Float.floatToIntBits(actual.lowerInterval())) {
+ return false;
+ }
+ if (Float.floatToIntBits(expected.upperInterval())
+ != Float.floatToIntBits(actual.upperInterval())) {
+ return false;
+ }
+ if (Float.floatToIntBits(expected.additionalCorrection())
+ != Float.floatToIntBits(actual.additionalCorrection())) {
+ return false;
+ }
+ if (expected.quantizedComponentSum() != actual.quantizedComponentSum()) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Merging two data-blind segments with different encodings is impossible
because re-quantization
+ * would require raw float vectors that were never written.
+ */
+ public void testDataBlindIncompatibleEncodingMerge() throws Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = 1 + random().nextInt(20);
+ int dims = random().nextInt(4, 33);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ // First segment: data-blind with PACKED_NIBBLE.
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setCodec(
+ TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(
+ ScalarEncoding.PACKED_NIBBLE,
Mode.DATA_BLIND_WITHOUT_FLOATS))))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ }
+ // Second segment: data-blind with a different encoding (UNSIGNED_BYTE).
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setCodec(
+ TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(
+ ScalarEncoding.UNSIGNED_BYTE,
Mode.DATA_BLIND_WITHOUT_FLOATS))))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ }
+ // Merging the two data-blind segments must fail: re-quantization needs
raw floats that were
+ // never written.
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergeScheduler(new SerialMergeScheduler())
+ .setCodec(
+ TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(
+ ScalarEncoding.UNSIGNED_BYTE,
Mode.DATA_BLIND_WITHOUT_FLOATS))))) {
+ Throwable t =
+ expectThrows(
+ Exception.class,
+ () -> {
+ w.forceMerge(1);
+ });
+ assertTrue(
+ "expected encoding-mismatch message, got: " + t,
+ exceptionChainContains(t, "re-quantization requires raw float
vectors"));
+ }
+ }
+ }
+
+ /**
+ * Verifies the on-disk metadata version distinguishes centered (0) from
data-blind (1) writes.
+ */
+ public void testVersionHeaders() throws Exception {
+ String fieldName = "field";
+ int numVectors = random().nextInt(4, 50);
+ int dims = random().nextInt(4, 33);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setUseCompoundFile(false)
+ .setCodec(
+ TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(encoding,
Mode.CENTERED))))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction, numVectors);
+ }
+ assertEquals(Lucene104ScalarQuantizedVectorsFormat.VERSION_START,
metaVersion(dir));
+ assertEquals(
+ Lucene104ScalarQuantizedVectorsFormat.Mode.CENTERED,
+ metaMode(dir)); // v0 metadata implies centered
+ }
+
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setUseCompoundFile(false)
+ .setCodec(dataBlindCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction, numVectors);
+ }
+ assertEquals(Lucene104ScalarQuantizedVectorsFormat.VERSION_DATA_BLIND,
metaVersion(dir));
+ assertEquals(
+
Lucene104ScalarQuantizedVectorsFormat.Mode.DATA_BLIND_WITHOUT_FLOATS,
metaMode(dir));
+ }
+
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setUseCompoundFile(false)
+ .setCodec(dataBlindWithFloatsCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction, numVectors);
+ }
+ assertEquals(Lucene104ScalarQuantizedVectorsFormat.VERSION_DATA_BLIND,
metaVersion(dir));
+ assertEquals(
+ Lucene104ScalarQuantizedVectorsFormat.Mode.DATA_BLIND_WITH_FLOATS,
metaMode(dir));
+ }
+ }
+
+ /**
+ * Merging data-blind segments (which stored no full-precision floats)
through a writer that
+ * stores float vectors must fail: the output format expects floats the
inputs cannot provide.
+ */
+ public void testDataBlindSegmentsCannotMergeIntoCenteredWriter() throws
Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = random().nextInt(1, 20);
+ int dims = random().nextInt(4, 33);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ w.commit();
+ }
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ w.commit();
+ }
+ // Re-encode through the centered writer.
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergeScheduler(new SerialMergeScheduler())
+ .setCodec(
+ TestUtil.alwaysKnnVectorsFormat(
+ new Lucene104ScalarQuantizedVectorsFormat(encoding,
Mode.CENTERED))))) {
+ Throwable t =
+ expectThrows(
+ Exception.class,
+ () -> {
+ w.forceMerge(1);
+ });
+ assertTrue(
+ "expected missing-float-vectors message, got: " + t,
+ exceptionChainContains(t, "without full-precision float vectors"));
+ }
+ }
+ }
+
+ /**
+ * Data-blind mode with floats quantizes against a zero centroid but still
writes full-precision
+ * float vectors, which must be readable and report a zero centroid.
+ */
+ public void testDataBlindWithFloatsKeepsRawFloatVectors() throws Exception {
+ String fieldName = "field";
+ int numVectors = random().nextInt(4, 50);
+ int dims = random().nextInt(4, 65);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setUseCompoundFile(false)
+ .setCodec(dataBlindWithFloatsCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction, numVectors);
+ }
+ try (IndexReader reader = DirectoryReader.open(dir)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ assertEquals(numVectors, r.getFloatVectorValues(fieldName).size());
+ Lucene104ScalarQuantizedVectorsReader vectorsReader =
+ (Lucene104ScalarQuantizedVectorsReader)
+ ((CodecReader)
r).getVectorReader().unwrapReaderForField(fieldName);
+ assertTrue(vectorsReader.hasRawFloatVectors(fieldName));
+ assertZeroCentroid(vectorsReader.getCentroid(fieldName), dims);
+ assertEquals(
+ Lucene104ScalarQuantizedVectorsFormat.Mode.DATA_BLIND_WITH_FLOATS,
metaMode(dir));
+ }
+ }
+ }
+
+ /** fp16 counterpart of {@link
#testDataBlindWithFloatsKeepsRawFloatVectors}. */
+ public void testDataBlindWithFloatsKeepsRawFloat16Vectors() throws Exception
{
+ String fieldName = "field";
+ int numVectors = random().nextInt(4, 50);
+ int dims = 2 * random().nextInt(2, 33);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setCodec(dataBlindWithFloatsCodec()))) {
+ addFloat16VectorDocs(w, fieldName, dims, similarityFunction,
numVectors);
+ }
+ try (IndexReader reader = DirectoryReader.open(dir)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ assertEquals(numVectors, r.getFloat16VectorValues(fieldName).size());
+ Lucene104ScalarQuantizedVectorsReader vectorsReader =
+ (Lucene104ScalarQuantizedVectorsReader)
+ ((CodecReader)
r).getVectorReader().unwrapReaderForField(fieldName);
+ assertTrue(vectorsReader.hasRawFloat16Vectors(fieldName));
+ assertZeroCentroid(vectorsReader.getCentroid(fieldName), dims);
+ }
+ }
+ }
+
+ /** Data-blind-with-floats segments merge with the float vectors intact. */
+ public void testDataBlindWithFloatsMultiSegmentMerge() throws Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = random().nextInt(4, 30);
+ int dims = random().nextInt(4, 65);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindWithFloatsCodec()))) {
+ for (int s = 0; s < 2; s++) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ w.commit();
+ }
+ w.forceMerge(1);
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ assertEquals(1, reader.leaves().size());
+ LeafReader r = getOnlyLeafReader(reader);
+ assertEquals(2 * numVectorsPerSegment,
r.getFloatVectorValues(fieldName).size());
+ Lucene104ScalarQuantizedVectorsReader vectorsReader =
+ (Lucene104ScalarQuantizedVectorsReader)
+ ((CodecReader)
r).getVectorReader().unwrapReaderForField(fieldName);
+ assertTrue(vectorsReader.hasRawFloatVectors(fieldName));
+ assertZeroCentroid(vectorsReader.getCentroid(fieldName), dims);
+ int k = random().nextInt(5, 20);
+ TopDocs td =
+ new IndexSearcher(reader)
+ .search(new KnnFloatVectorQuery(fieldName,
randomVector(dims), k), k);
+ // The index holds exactly 2 * numVectorsPerSegment vectors; when k
exceeds that, the
+ // query can only return the vectors that exist.
+ assertEquals(Math.min(k, 2 * numVectorsPerSegment),
td.totalHits.value());
+ }
+ }
+ }
+ }
+
+ /**
+ * Merging a data-blind (no-floats) segment through a data-blind-with-floats
writer must fail: the
+ * output format expects floats the inputs cannot provide.
+ */
+ public void testDataBlindNoFloatsCannotMergeIntoDataBlindWithFloatsWriter()
throws Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = 1 + random().nextInt(20);
+ int dims = random().nextInt(4, 33);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ w.commit();
+ }
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ addFloatVectorDocs(w, fieldName, dims, similarityFunction,
numVectorsPerSegment);
+ w.commit();
+ }
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMergeScheduler(new SerialMergeScheduler())
+ .setCodec(dataBlindWithFloatsCodec()))) {
+ Throwable t =
+ expectThrows(
+ Exception.class,
+ () -> {
+ w.forceMerge(1);
+ });
+ assertTrue(
+ "expected missing-float-vectors message, got: " + t,
+ exceptionChainContains(t, "without full-precision float vectors"));
+ }
+ }
+ }
+
+ private static void assertZeroCentroid(float[] centroid, int dims) {
+ assertNotNull(centroid);
+ assertEquals(dims, centroid.length);
+ for (float v : centroid) {
+ assertEquals(0f, v, 0f);
+ }
+ }
+
+ /**
+ * Reads the {@link Lucene104ScalarQuantizedVectorsFormat.Mode} wire value
from the metadata of
+ * the first field with vectors, or {@link
Lucene104ScalarQuantizedVectorsFormat.Mode#CENTERED}
+ * for version 0 metadata which implies it.
+ */
+ private Lucene104ScalarQuantizedVectorsFormat.Mode metaMode(Directory dir)
throws IOException {
+ for (String file : dir.listAll()) {
+ if (file.endsWith("." +
Lucene104ScalarQuantizedVectorsFormat.META_EXTENSION)) {
+ try (IndexInput in = dir.openInput(file, IOContext.DEFAULT)) {
+ int version =
+ CodecUtil.checkHeader(
+ in,
+ Lucene104ScalarQuantizedVectorsFormat.META_CODEC_NAME,
+ Lucene104ScalarQuantizedVectorsFormat.VERSION_START,
+ Lucene104ScalarQuantizedVectorsFormat.VERSION_CURRENT);
+ // Consume the segment ID and suffix that checkHeader (unlike
checkIndexHeader) leaves.
+ in.readBytes(new byte[16], 0, 16);
+ in.readString();
+ for (int fieldNumber = in.readInt(); fieldNumber != -1; fieldNumber
= in.readInt()) {
+ in.readInt(); // vector encoding ordinal
+ in.readInt(); // similarity ordinal
+ int dimension = in.readVInt();
+ in.readVLong(); // vector data offset
+ in.readVLong(); // vector data length
+ int count = in.readVInt();
+ if (count > 0) {
+ in.readVInt(); // ScalarEncoding wire number
+ if (version ==
Lucene104ScalarQuantizedVectorsFormat.VERSION_START) {
+ in.readFloats(new float[dimension], 0, dimension);
+ in.readInt(); // centroidDP
+ return Lucene104ScalarQuantizedVectorsFormat.Mode.CENTERED;
+ }
+ return
Lucene104ScalarQuantizedVectorsFormat.Mode.fromWireNumber(in.readByte());
+ }
+ }
+ }
+ }
+ }
+ throw new AssertionError("no metadata file found");
+ }
+
+ public void testDataBlindFloat16Search() throws Exception {
+ String fieldName = "field";
+ int numVectors = random().nextInt(99, 500);
+ int dims = 2 * random().nextInt(2, 33);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ KnnFloat16VectorField knnField =
+ new KnnFloat16VectorField(
+ fieldName, randomNormalizedFloat16Vector(dims),
similarityFunction);
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ for (int i = 0; i < numVectors; i++) {
+ Document doc = new Document();
+ knnField.setVectorValue(randomNormalizedFloat16Vector(dims));
+ doc.add(knnField);
+ w.addDocument(doc);
+ }
+ w.commit();
+
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ IndexSearcher searcher = new IndexSearcher(reader);
+ final int k = random().nextInt(5, 50);
+ short[] queryVector = randomNormalizedFloat16Vector(dims);
+ TopDocs collectedDocs =
+ searcher.search(new KnnFloat16VectorQuery(fieldName,
queryVector, k), k);
+ assertEquals(k, collectedDocs.totalHits.value());
+ assertEquals(TotalHits.Relation.EQUAL_TO,
collectedDocs.totalHits.relation());
+ }
+ }
+ }
+ }
+
+ /**
+ * Data-blind fp16 fields reconstruct as fp16-rounded dequantized values via
{@link
+ * OffHeapScalarQuantizedFloat16VectorValues}, so reads come back as {@link
+ * Lucene104ScalarQuantizedVectorsReader.ScalarQuantizedFloat16VectorValues}.
+ */
+ public void testDataBlindFloat16ReconstructsFp16() throws Exception {
+ String vectorFieldName = "vec1";
+ int numVectors = 1 + random().nextInt(50);
+ int dim = 2 * random().nextInt(1, 33);
+ float eps = (1f / (float) (1 << getQuantizationBits())) + 1e-3f;
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ List<short[]> vectors = new ArrayList<>(numVectors);
+ for (int i = 0; i < numVectors; i++) {
+ vectors.add(randomNormalizedFloat16Vector(dim));
+ }
+
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(
+ dir,
+ newIndexWriterConfig()
+ .setMaxBufferedDocs(numVectors + 1)
+ .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH)
+ .setMergePolicy(NoMergePolicy.INSTANCE)
+ .setCodec(dataBlindCodec()))) {
+ for (int i = 0; i < numVectors; i++) {
+ Document doc = new Document();
+ doc.add(new KnnFloat16VectorField(vectorFieldName, vectors.get(i),
similarityFunction));
+ w.addDocument(doc);
+ }
+ }
+
+ try (IndexReader reader = DirectoryReader.open(dir)) {
+ LeafReader r = getOnlyLeafReader(reader);
+ Float16VectorValues float16VectorValues =
r.getFloat16VectorValues(vectorFieldName);
+ assertTrue(
+ float16VectorValues
+ instanceof
+
Lucene104ScalarQuantizedVectorsReader.ScalarQuantizedFloat16VectorValues);
+ KnnVectorValues.DocIndexIterator iter = float16VectorValues.iterator();
+ for (int docId = iter.nextDoc(); docId != NO_MORE_DOCS; docId =
iter.nextDoc()) {
+ short[] reconstructed =
float16VectorValues.vectorValue(iter.index());
+ short[] original = vectors.get(docId);
+ float mae = 0;
+ for (int i = 0; i < dim; i++) {
+ mae +=
+ Math.abs(
+ Float.float16ToFloat(reconstructed[i]) -
Float.float16ToFloat(original[i]));
+ }
+ mae /= dim;
+ assertTrue(
+ "bits: " + getQuantizationBits() + " mae: " + mae + " > eps: " +
eps, mae <= eps);
+ }
+ }
+ }
+ }
+
+ public void testDataBlindFloat16MultiSegmentMerge() throws Exception {
+ String fieldName = "field";
+ int numVectorsPerSegment = random().nextInt(4, 30);
+ int dims = 2 * random().nextInt(2, 17);
+ VectorSimilarityFunction similarityFunction = randomSimilarity();
+ try (Directory dir = newDirectory()) {
+ try (IndexWriter w =
+ new IndexWriter(dir,
newIndexWriterConfig().setCodec(dataBlindCodec()))) {
+ for (int s = 0; s < 2; s++) {
+ for (int i = 0; i < numVectorsPerSegment; i++) {
+ Document doc = new Document();
+ doc.add(
+ new KnnFloat16VectorField(
+ fieldName, randomNormalizedFloat16Vector(dims),
similarityFunction));
+ w.addDocument(doc);
+ }
+ w.commit();
+ }
+ w.forceMerge(1);
+ try (IndexReader reader = DirectoryReader.open(w)) {
+ assertEquals(1, reader.leaves().size());
+ LeafReader r = getOnlyLeafReader(reader);
+ assertEquals(2 * numVectorsPerSegment,
r.getFloat16VectorValues(fieldName).size());
+ int k = random().nextInt(5, Math.min(15, 2 * numVectorsPerSegment +
1));
+ TopDocs td =
+ new IndexSearcher(reader)
+ .search(
+ new KnnFloat16VectorQuery(fieldName,
randomNormalizedFloat16Vector(dims), k),
+ k);
+ assertEquals(k, td.totalHits.value());
+ }
+ }
+ }
+ }
+
+ public void testDataBlindFloat16IncompatibleEncodingMerge() throws Exception
{
Review Comment:
@mccullocht just a heads up, it seems that this test is flaky. it failed
when i merged it into my PR
you can reproduce on main with:
```gradlew test --tests
TestLucene104ScalarQuantizedVectorsFormat.testDataBlindFloat16IncompatibleEncodingMerge
-Dtests.seed=4A219E3AE1C75137 -Dtests.asserts=false```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
