diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 92e7fed208e2..c1b23d51e910 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -129,6 +129,8 @@ New Features * GITHUB#16383: Add fp16 vector encoding support. (Pulkit Gupta) +* GITHUB#16473: Add scalar quantization support in Fp16 vector encoding. (Pulkit Gupta) + Improvements --------------------- * GITHUB#15704: Replace LinkedList with more efficient data structure. (Renato Haeberli) diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java index 70dad47d928b..d061fe6db55b 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java @@ -111,6 +111,43 @@ public RandomVectorScorer getRandomVectorScorer( public RandomVectorScorer getRandomVectorScorer( VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, short[] target) throws IOException { + if (vectorValues instanceof QuantizedByteVectorValues qv) { + FlatVectorsScorer.checkDimensions(target.length, qv.dimension()); + OptimizedScalarQuantizer quantizer = qv.getQuantizer(); + ScalarEncoding scalarEncoding = qv.getScalarEncoding(); + byte[] scratch = new byte[scalarEncoding.getDiscreteDimensions(qv.dimension())]; + final byte[] targetQuantized; + if (scalarEncoding.isAsymmetric() == false) { + targetQuantized = scratch; + } else { + // This is asymmetric quantization, we will pack the vector + targetQuantized = new byte[scalarEncoding.getQueryPackedLength(scratch.length)]; + } + // Inflate the fp16 query to fp32 and normalize there; quantization operates on fp32. + float[] copy = new float[target.length]; + for (int i = 0; i < target.length; i++) { + copy[i] = Float.float16ToFloat(target[i]); + } + if (similarityFunction == COSINE) { + VectorUtil.l2normalize(copy); + } + var targetCorrectiveTerms = + quantizer.scalarQuantize(copy, scratch, scalarEncoding.getQueryBits(), qv.getCentroid()); + // for asymmetric encodings with 4-bit query, we need to transpose the nibbles for fast + // scoring comparisons + if (scalarEncoding == ScalarEncoding.SINGLE_BIT_QUERY_NIBBLE + || scalarEncoding == ScalarEncoding.DIBIT_QUERY_NIBBLE) { + OptimizedScalarQuantizer.transposeHalfByte(scratch, targetQuantized); + } + return new RandomVectorScorer.AbstractRandomVectorScorer(qv) { + @Override + public float score(int node) throws IOException { + return quantizedScore( + targetQuantized, targetCorrectiveTerms, qv, node, similarityFunction); + } + }; + } + // It is possible to get to this branch during initial indexing and flush return nonQuantizedDelegate.getRandomVectorScorer(similarityFunction, vectorValues, target); } diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java index ef458b05162c..c4ad2eb2122d 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java @@ -224,7 +224,26 @@ public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) thr @Override public RandomVectorScorer getRandomVectorScorer(String field, short[] target) throws IOException { - return rawVectorsReader.getRandomVectorScorer(field, target); + FieldEntry fi = fields.get(field); + if (fi == null) { + return null; + } + return vectorScorer.getRandomVectorScorer( + fi.similarityFunction, + OffHeapScalarQuantizedVectorValues.load( + fi.ordToDocDISIReaderConfiguration, + fi.dimension, + fi.size, + new OptimizedScalarQuantizer(fi.similarityFunction), + fi.scalarEncoding, + fi.similarityFunction, + vectorScorer, + fi.centroid, + fi.centroidDP, + fi.vectorDataOffset, + fi.vectorDataLength, + quantizedVectorData), + target); } @Override @@ -289,7 +308,57 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException { @Override public Float16VectorValues getFloat16VectorValues(String field) throws IOException { - return rawVectorsReader.getFloat16VectorValues(field); + FieldEntry fi = fields.get(field); + if (fi == null) { + return null; + } + if (fi.vectorEncoding != VectorEncoding.FLOAT16) { + throw new IllegalArgumentException( + "field=\"" + + field + + "\" is encoded as: " + + fi.vectorEncoding + + " expected: " + + VectorEncoding.FLOAT16); + } + + Float16VectorValues rawFloat16VectorValues = rawVectorsReader.getFloat16VectorValues(field); + + OffHeapScalarQuantizedVectorValues sqvv = + OffHeapScalarQuantizedVectorValues.load( + fi.ordToDocDISIReaderConfiguration, + fi.dimension, + fi.size, + new OptimizedScalarQuantizer(fi.similarityFunction), + fi.scalarEncoding, + fi.similarityFunction, + vectorScorer, + fi.centroid, + fi.centroidDP, + fi.vectorDataOffset, + fi.vectorDataLength, + quantizedVectorData); + + if (rawFloat16VectorValues.size() == 0) { + // The raw float16 vectors were dropped, so reads reconstruct values by dequantizing. Pair + // that view with sqvv so scorer() scores in quantized space while vectorValue() and + // rescorer() dequantize. + Float16VectorValues dequantizedRawVectorValues = + OffHeapScalarQuantizedFloat16VectorValues.load( + fi.ordToDocDISIReaderConfiguration, + fi.dimension, + fi.size, + fi.scalarEncoding, + fi.similarityFunction, + vectorScorer, + fi.centroid, + fi.vectorDataOffset, + fi.vectorDataLength, + quantizedVectorData); + return new ScalarQuantizedFloat16VectorValues(dequantizedRawVectorValues, sqvv); + } + + return new ScalarQuantizedFloat16VectorValues(rawFloat16VectorValues, sqvv); } @Override @@ -302,11 +371,25 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Accep public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { if (knnCollector.k() == 0) return; - final RandomVectorScorer scorer = getRandomVectorScorer(field, target); + exhaustiveBulkScore(getRandomVectorScorer(field, target), knnCollector, acceptDocs); + } + + @Override + public void search(String field, short[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) + throws IOException { + if (knnCollector.k() == 0) return; + exhaustiveBulkScore(getRandomVectorScorer(field, target), knnCollector, acceptDocs); + } + + /** + * Scores every accepted vector with the given scorer, collecting into {@code knnCollector}. + * Scoring happens in batches of {@link #EXHAUSTIVE_BULK_SCORE_ORDS} ordinals. + */ + private static void exhaustiveBulkScore( + RandomVectorScorer scorer, KnnCollector knnCollector, AcceptDocs acceptDocs) + throws IOException { if (scorer == null) return; Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); - // if k is larger than the number of vectors we expect to visit in an HNSW search, - // we can just iterate over all vectors and collect them. int[] ords = new int[EXHAUSTIVE_BULK_SCORE_ORDS]; float[] scores = new float[EXHAUSTIVE_BULK_SCORE_ORDS]; int numOrds = 0; @@ -339,12 +422,6 @@ public void search(String field, float[] target, KnnCollector knnCollector, Acce } } - @Override - public void search(String field, short[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) - throws IOException { - rawVectorsReader.search(field, target, knnCollector, acceptDocs); - } - @Override public void close() throws IOException { IOUtils.close(quantizedVectorData, rawVectorsReader); @@ -366,10 +443,7 @@ public Map getOffHeapByteSize(FieldInfo fieldInfo) { var raw = rawVectorsReader.getOffHeapByteSize(fieldInfo); var fieldEntry = fields.get(fieldInfo.name); if (fieldEntry == null) { - // Only FLOAT32 fields are scalar-quantized by this format; BYTE and FLOAT16 fields are - // stored raw by the delegate and therefore have no quantized field entry here. - assert fieldInfo.getVectorEncoding() == VectorEncoding.BYTE - || fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT16; + assert fieldInfo.getVectorEncoding() == VectorEncoding.BYTE; return raw; } var quant = Map.of(VECTOR_DATA_EXTENSION, fieldEntry.vectorDataLength()); @@ -442,14 +516,16 @@ public QuantizedByteVectorValues getQuantizedVectorValues(String field) throws I if (fi == null) { return null; } - if (fi.vectorEncoding != VectorEncoding.FLOAT32) { + if (fi.vectorEncoding.isFloatingPoint() == false) { throw new IllegalArgumentException( "field=\"" + field + "\" is encoded as: " + fi.vectorEncoding + " expected: " - + VectorEncoding.FLOAT32); + + VectorEncoding.FLOAT32 + + " or " + + VectorEncoding.FLOAT16); } return OffHeapScalarQuantizedVectorValues.load( fi.ordToDocDISIReaderConfiguration, @@ -692,4 +768,66 @@ QuantizedByteVectorValues getQuantizedVectorValues() throws IOException { return quantizedVectorValues; } } + + /** Vector values holding raw and quantized vector values */ + protected static final class ScalarQuantizedFloat16VectorValues extends Float16VectorValues { + private final Float16VectorValues rawVectorValues; + private final QuantizedByteVectorValues quantizedVectorValues; + + ScalarQuantizedFloat16VectorValues( + Float16VectorValues rawVectorValues, QuantizedByteVectorValues quantizedVectorValues) { + this.rawVectorValues = rawVectorValues; + this.quantizedVectorValues = quantizedVectorValues; + } + + @Override + public int dimension() { + return rawVectorValues.dimension(); + } + + @Override + public int size() { + return rawVectorValues.size(); + } + + @Override + public short[] vectorValue(int ord) throws IOException { + return rawVectorValues.vectorValue(ord); + } + + @Override + public ScalarQuantizedFloat16VectorValues copy() throws IOException { + return new ScalarQuantizedFloat16VectorValues( + rawVectorValues.copy(), quantizedVectorValues.copy()); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return rawVectorValues.getAcceptOrds(acceptDocs); + } + + @Override + public int ordToDoc(int ord) { + return rawVectorValues.ordToDoc(ord); + } + + @Override + public DocIndexIterator iterator() { + return rawVectorValues.iterator(); + } + + @Override + public VectorScorer scorer(short[] query) throws IOException { + return quantizedVectorValues.scorer(query); + } + + @Override + public VectorScorer rescorer(short[] target) throws IOException { + return rawVectorValues.rescorer(target); + } + + QuantizedByteVectorValues getQuantizedVectorValues() throws IOException { + return quantizedVectorValues; + } + } } diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java index af2374775be8..6d764048d057 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java @@ -36,6 +36,7 @@ import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.Float16VectorValues; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.KnnVectorValues; @@ -43,9 +44,7 @@ import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.Sorter; import org.apache.lucene.index.VectorEncoding; -import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.internal.hppc.FloatArrayList; -import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.VectorScorer; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; @@ -65,7 +64,7 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter { shallowSizeOfInstance(Lucene104ScalarQuantizedVectorsWriter.class); private final SegmentWriteState segmentWriteState; - private final List fields = new ArrayList<>(); + private final List> fields = new ArrayList<>(); private final IndexOutput meta, vectorData; private final ScalarEncoding encoding; private final FlatVectorsWriter rawVectorDelegate; @@ -118,10 +117,8 @@ public Lucene104ScalarQuantizedVectorsWriter( @Override public FlatFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { FlatFieldVectorsWriter rawVectorDelegate = this.rawVectorDelegate.addField(fieldInfo); - if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) { - @SuppressWarnings("unchecked") - FieldWriter fieldWriter = - new FieldWriter(fieldInfo, (FlatFieldVectorsWriter) rawVectorDelegate); + if (fieldInfo.getVectorEncoding().isFloatingPoint()) { + FieldWriter fieldWriter = FieldWriter.create(fieldInfo, rawVectorDelegate); fields.add(fieldWriter); return fieldWriter; } @@ -131,25 +128,11 @@ public FlatFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOExceptio @Override public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { rawVectorDelegate.flush(maxDoc, sortMap); - for (FieldWriter field : fields) { - // after raw vectors are written, normalize vectors for clustering and quantization - if (VectorSimilarityFunction.COSINE == field.fieldInfo.getVectorSimilarityFunction()) { - field.normalizeVectors(); - } - final float[] clusterCenter; - int vectorCount = field.flatFieldVectorsWriter.getVectors().size(); - clusterCenter = new float[field.dimensionSums.length]; - if (vectorCount > 0) { - for (int i = 0; i < field.dimensionSums.length; i++) { - clusterCenter[i] = field.dimensionSums[i] / vectorCount; - } - if (VectorSimilarityFunction.COSINE == field.fieldInfo.getVectorSimilarityFunction()) { - VectorUtil.l2normalize(clusterCenter); - } - } + for (FieldWriter field : fields) { + final float[] clusterCenter = field.centroid(); if (segmentWriteState.infoStream.isEnabled(QUANTIZED_VECTOR_COMPONENT)) { segmentWriteState.infoStream.message( - QUANTIZED_VECTOR_COMPONENT, "Vectors' count:" + vectorCount); + QUANTIZED_VECTOR_COMPONENT, "Vectors' count:" + field.getVectors().size()); } OptimizedScalarQuantizer quantizer = new OptimizedScalarQuantizer(field.fieldInfo.getVectorSimilarityFunction()); @@ -163,7 +146,10 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { } private void writeField( - FieldWriter fieldData, float[] clusterCenter, int maxDoc, OptimizedScalarQuantizer quantizer) + FieldWriter fieldData, + float[] clusterCenter, + int maxDoc, + OptimizedScalarQuantizer quantizer) throws IOException { // write vector values long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES); @@ -183,7 +169,7 @@ private void writeField( } private void writeVectors( - FieldWriter fieldData, float[] clusterCenter, OptimizedScalarQuantizer scalarQuantizer) + FieldWriter fieldData, float[] clusterCenter, OptimizedScalarQuantizer scalarQuantizer) throws IOException { byte[] scratch = new byte[encoding.getDiscreteDimensions(fieldData.fieldInfo.getVectorDimension())]; @@ -194,9 +180,9 @@ private void writeVectors( new byte[encoding.getDocPackedLength(scratch.length)]; }; for (int i = 0; i < fieldData.getVectors().size(); i++) { - float[] v = fieldData.getVectors().get(i); OptimizedScalarQuantizer.QuantizationResult corrections = - scalarQuantizer.scalarQuantize(v, scratch, encoding.getBits(), clusterCenter); + scalarQuantizer.scalarQuantize( + fieldData.floatVectorValue(i), scratch, encoding.getBits(), clusterCenter); switch (encoding) { case PACKED_NIBBLE -> OffHeapScalarQuantizedVectorValues.packNibbles(scratch, vector); case SINGLE_BIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.packAsBinary(scratch, vector); @@ -212,7 +198,7 @@ private void writeVectors( } private void writeSortingField( - FieldWriter fieldData, + FieldWriter fieldData, float[] clusterCenter, int maxDoc, Sorter.DocMap sortMap, @@ -241,7 +227,7 @@ private void writeSortingField( } private void writeSortedVectors( - FieldWriter fieldData, + FieldWriter fieldData, float[] clusterCenter, int[] ordMap, OptimizedScalarQuantizer scalarQuantizer) @@ -255,9 +241,9 @@ private void writeSortedVectors( new byte[encoding.getDocPackedLength(scratch.length)]; }; for (int ordinal : ordMap) { - float[] v = fieldData.getVectors().get(ordinal); OptimizedScalarQuantizer.QuantizationResult corrections = - scalarQuantizer.scalarQuantize(v, scratch, encoding.getBits(), clusterCenter); + scalarQuantizer.scalarQuantize( + fieldData.floatVectorValue(ordinal), scratch, encoding.getBits(), clusterCenter); switch (encoding) { case PACKED_NIBBLE -> OffHeapScalarQuantizedVectorValues.packNibbles(scratch, vector); case SINGLE_BIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.packAsBinary(scratch, vector); @@ -319,12 +305,27 @@ public void finish() throws IOException { } } + private QuantizedByteVectorValues mergedQuantizedVectorValues( + FieldInfo fieldInfo, MergeState mergeState, float[] centroid) throws IOException { + OptimizedScalarQuantizer quantizer = + new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction()); + FloatVectorValues vectorValues = + fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT16 + ? new Float16AsFloatVectorValues( + MergedVectorValues.mergeFloat16VectorValues(fieldInfo, mergeState)) + : MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + if (fieldInfo.getVectorSimilarityFunction() == COSINE) { + vectorValues = new NormalizedFloatVectorValues(vectorValues); + } + return new QuantizedFloatVectorValues(vectorValues, quantizer, encoding, centroid); + } + @Override public void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { // Don't need access to the random vectors, we can just use the merged rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState); - if (!fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) { + if (fieldInfo.getVectorEncoding().isFloatingPoint() == false) { return; } final float[] centroid; @@ -335,17 +336,8 @@ public void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState mergeState) segmentWriteState.infoStream.message( QUANTIZED_VECTOR_COMPONENT, "Vectors' count:" + vectorCount); } - FloatVectorValues floatVectorValues = - MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - if (fieldInfo.getVectorSimilarityFunction() == COSINE) { - floatVectorValues = new NormalizedFloatVectorValues(floatVectorValues); - } - QuantizedFloatVectorValues quantizedVectorValues = - new QuantizedFloatVectorValues( - floatVectorValues, - new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction()), - encoding, - centroid); + QuantizedByteVectorValues quantizedVectorValues = + mergedQuantizedVectorValues(fieldInfo, mergeState, centroid); long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES); DocsWithFieldSet docsWithField = writeVectorData(vectorData, quantizedVectorValues); long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset; @@ -428,21 +420,40 @@ static float[] getCentroid(KnnVectorsReader vectorsReader, String fieldName) { return null; } + /** + * Returns the reader's floating-point vectors viewed as fp32, inflating fp16 on read, or null + * when the field is absent from this reader or is byte-encoded. + */ + private static FloatVectorValues floatingPointVectorValues( + KnnVectorsReader reader, FieldInfo fieldInfo) throws IOException { + return switch (fieldInfo.getVectorEncoding()) { + case FLOAT32 -> reader.getFloatVectorValues(fieldInfo.name); + case FLOAT16 -> { + Float16VectorValues f16 = reader.getFloat16VectorValues(fieldInfo.name); + yield f16 == null ? null : new Float16AsFloatVectorValues(f16); + } + case BYTE -> null; + }; + } + static int mergeAndRecalculateCentroids( MergeState mergeState, FieldInfo fieldInfo, float[] mergedCentroid) throws IOException { boolean recalculate = false; int totalVectorCount = 0; for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i]; - if (knnVectorsReader == null - || knnVectorsReader.getFloatVectorValues(fieldInfo.name) == null) { + if (knnVectorsReader == null) { continue; } - float[] centroid = getCentroid(knnVectorsReader, fieldInfo.name); - int vectorCount = knnVectorsReader.getFloatVectorValues(fieldInfo.name).size(); + KnnVectorValues values = floatingPointVectorValues(knnVectorsReader, fieldInfo); + if (values == null) { + continue; + } + int vectorCount = values.size(); if (vectorCount == 0) { continue; } + float[] centroid = getCentroid(knnVectorsReader, fieldInfo.name); totalVectorCount += vectorCount; // If there aren't centroids, or previously clustered with more than one cluster // or if there are deleted docs, we must recalculate the centroid @@ -471,28 +482,14 @@ static int mergeAndRecalculateCentroids( static int calculateCentroid(MergeState mergeState, FieldInfo fieldInfo, float[] centroid) throws IOException { - assert fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32); + assert fieldInfo.getVectorEncoding().isFloatingPoint(); // clear out the centroid Arrays.fill(centroid, 0); int count = 0; for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i]; if (knnVectorsReader == null) continue; - FloatVectorValues vectorValues = - mergeState.knnVectorsReaders[i].getFloatVectorValues(fieldInfo.name); - if (vectorValues == null) { - continue; - } - KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator(); - for (int doc = iterator.nextDoc(); - doc != DocIdSetIterator.NO_MORE_DOCS; - doc = iterator.nextDoc()) { - ++count; - float[] vector = vectorValues.vectorValue(iterator.index()); - for (int j = 0; j < vector.length; j++) { - centroid[j] += vector[j]; - } - } + count += accumulateCentroid(knnVectorsReader, fieldInfo, centroid); } if (count == 0) { return count; @@ -506,6 +503,24 @@ static int calculateCentroid(MergeState mergeState, FieldInfo fieldInfo, float[] return count; } + private static int accumulateCentroid( + KnnVectorsReader reader, FieldInfo fieldInfo, float[] centroid) throws IOException { + FloatVectorValues vectorValues = floatingPointVectorValues(reader, fieldInfo); + if (vectorValues == null) { + return 0; + } + int count = 0; + KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator(); + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + count++; + float[] vector = vectorValues.vectorValue(iterator.index()); + for (int j = 0; j < vector.length; j++) { + centroid[j] += vector[j]; + } + } + return count; + } + @Override public long ramBytesUsed() { long total = SHALLOW_RAM_BYTES_USED; @@ -514,7 +529,7 @@ public long ramBytesUsed() { // For float32 fields, this covers the flat vector data; our FieldWriter adds the // quantization-specific overhead (magnitudes, dimensionSums) on top. total += rawVectorDelegate.ramBytesUsed(); - for (FieldWriter field : fields) { + for (FieldWriter field : fields) { // quantizationOverheadBytesUsed() intentionally excludes flatFieldVectorsWriter // because rawVectorDelegate.ramBytesUsed() already accounts for all flat vector // data at the writer level. Calling field.ramBytesUsed() here would double-count. @@ -523,33 +538,44 @@ public long ramBytesUsed() { return total; } - static class FieldWriter extends FlatFieldVectorsWriter { + abstract static class FieldWriter extends FlatFieldVectorsWriter { private static final long SHALLOW_SIZE = shallowSizeOfInstance(FieldWriter.class); - private final FieldInfo fieldInfo; + protected final FieldInfo fieldInfo; private boolean finished; - private final FlatFieldVectorsWriter flatFieldVectorsWriter; + protected final FlatFieldVectorsWriter flatFieldVectorsWriter; private final float[] dimensionSums; private final FloatArrayList magnitudes = new FloatArrayList(); + protected final int dim; - FieldWriter(FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + FieldWriter(FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { this.fieldInfo = fieldInfo; this.flatFieldVectorsWriter = flatFieldVectorsWriter; - this.dimensionSums = new float[fieldInfo.getVectorDimension()]; + this.dim = fieldInfo.getVectorDimension(); + this.dimensionSums = new float[dim]; + } + + @SuppressWarnings("unchecked") + static FieldWriter create( + FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + return switch (fieldInfo.getVectorEncoding()) { + case BYTE -> throw new UnsupportedOperationException("Byte Vectors aren't supported"); + case FLOAT32 -> + new Float32FieldWriter( + fieldInfo, (FlatFieldVectorsWriter) flatFieldVectorsWriter); + case FLOAT16 -> + new Float16FieldWriter( + fieldInfo, (FlatFieldVectorsWriter) flatFieldVectorsWriter); + }; } @Override - public List getVectors() { + public List getVectors() { return flatFieldVectorsWriter.getVectors(); } - public void normalizeVectors() { - for (int i = 0; i < flatFieldVectorsWriter.getVectors().size(); i++) { - float[] vector = flatFieldVectorsWriter.getVectors().get(i); - float magnitude = magnitudes.get(i); - for (int j = 0; j < vector.length; j++) { - vector[j] /= magnitude; - } - } + @Override + public T copyValue(T vectorValue) { + throw new UnsupportedOperationException(); } @Override @@ -571,26 +597,55 @@ public boolean isFinished() { return finished && flatFieldVectorsWriter.isFinished(); } - @Override - public void addValue(int docID, float[] vectorValue) throws IOException { - flatFieldVectorsWriter.addValue(docID, vectorValue); + /** The ordinal's stored vector as fp32, ready for quantization (unit-length for COSINE). */ + abstract float[] floatVectorValue(int ord); + + /** + * Adds {@code vector} to the centroid sums, unit-scaled when COSINE, and caches its magnitude + * for {@link #scaleToUnitLength}. + */ + protected final void accumulate(float[] vector) { if (fieldInfo.getVectorSimilarityFunction() == COSINE) { - float dp = VectorUtil.dotProduct(vectorValue, vectorValue); + float dp = VectorUtil.dotProduct(vector, vector); float divisor = (float) Math.sqrt(dp); magnitudes.add(divisor); - for (int i = 0; i < vectorValue.length; i++) { - dimensionSums[i] += (vectorValue[i] / divisor); + for (int i = 0; i < vector.length; i++) { + dimensionSums[i] += (vector[i] / divisor); } } else { - for (int i = 0; i < vectorValue.length; i++) { - dimensionSums[i] += vectorValue[i]; + for (int i = 0; i < vector.length; i++) { + dimensionSums[i] += vector[i]; } } } - @Override - public float[] copyValue(float[] vectorValue) { - throw new UnsupportedOperationException(); + /** + * The mean of the accumulated vectors, unit-length for COSINE, used as the quantization + * centroid. All zeroes when no vectors were added. + */ + protected final float[] centroid() { + float[] centroid = new float[dim]; + int vectorCount = getVectors().size(); + if (vectorCount > 0) { + for (int i = 0; i < dim; i++) { + centroid[i] = dimensionSums[i] / vectorCount; + } + if (fieldInfo.getVectorSimilarityFunction() == COSINE) { + VectorUtil.l2normalize(centroid); + } + } + return centroid; + } + + /** + * Writes {@code src} into {@code dst} scaled to unit length, using the magnitude cached for + * {@code ord}. The two arrays may be the same. + */ + protected final void scaleToUnitLength(float[] src, float[] dst, int ord) { + float magnitude = magnitudes.get(ord); + for (int i = 0; i < src.length; i++) { + dst[i] = src[i] / magnitude; + } } /** @@ -613,6 +668,75 @@ public long ramBytesUsed() { } } + private static class Float32FieldWriter extends FieldWriter { + private final float[] normalized; + + Float32FieldWriter( + FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + super(fieldInfo, flatFieldVectorsWriter); + this.normalized = new float[dim]; + } + + @Override + public void addValue(int docID, float[] vectorValue) throws IOException { + flatFieldVectorsWriter.addValue(docID, vectorValue); + accumulate(vectorValue); + } + + @Override + float[] floatVectorValue(int ord) { + float[] vector = flatFieldVectorsWriter.getVectors().get(ord); + if (fieldInfo.getVectorSimilarityFunction() == COSINE) { + scaleToUnitLength(vector, normalized, ord); + return normalized; + } + return vector; + } + + @Override + long quantizationOverheadBytesUsed() { + return super.quantizationOverheadBytesUsed() + RamUsageEstimator.sizeOf(normalized); + } + } + + private static class Float16FieldWriter extends FieldWriter { + private final float[] inflated; + + Float16FieldWriter( + FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + super(fieldInfo, flatFieldVectorsWriter); + this.inflated = new float[dim]; + } + + @Override + public void addValue(int docID, short[] vectorValue) throws IOException { + flatFieldVectorsWriter.addValue(docID, vectorValue); + inflate(vectorValue); + accumulate(inflated); + } + + @Override + float[] floatVectorValue(int ord) { + inflate(flatFieldVectorsWriter.getVectors().get(ord)); + if (fieldInfo.getVectorSimilarityFunction() == COSINE) { + scaleToUnitLength(inflated, inflated, ord); + } + return inflated; + } + + /** Inflates an fp16 vector into {@link #inflated}. */ + private void inflate(short[] vectorValue) { + for (int i = 0; i < vectorValue.length; i++) { + inflated[i] = Float.float16ToFloat(vectorValue[i]); + } + } + + @Override + long quantizationOverheadBytesUsed() { + return super.quantizationOverheadBytesUsed() + RamUsageEstimator.sizeOf(inflated); + } + } + static class QuantizedFloatVectorValues extends QuantizedByteVectorValues { private OptimizedScalarQuantizer.QuantizationResult corrections; private final byte[] quantized; @@ -700,6 +824,11 @@ public VectorScorer scorer(float[] target) throws IOException { throw new UnsupportedOperationException(); } + @Override + public VectorScorer scorer(short[] target) throws IOException { + throw new UnsupportedOperationException(); + } + @Override public QuantizedByteVectorValues copy() throws IOException { return new QuantizedFloatVectorValues(values.copy(), quantizer, encoding, centroid); @@ -728,6 +857,54 @@ public int ordToDoc(int ord) { } } + /** + * Exposes a {@link Float16VectorValues} as {@link FloatVectorValues}, inflating fp16 to fp32 on + * read, so the merge path can reuse the fp32 quantization classes. + */ + static final class Float16AsFloatVectorValues extends FloatVectorValues { + private final Float16VectorValues values; + private final float[] floatVector; + + Float16AsFloatVectorValues(Float16VectorValues values) { + this.values = values; + this.floatVector = new float[values.dimension()]; + } + + @Override + public int dimension() { + return values.dimension(); + } + + @Override + public int size() { + return values.size(); + } + + @Override + public int ordToDoc(int ord) { + return values.ordToDoc(ord); + } + + @Override + public float[] vectorValue(int ord) throws IOException { + short[] v = values.vectorValue(ord); + for (int i = 0; i < v.length; i++) { + floatVector[i] = Float.float16ToFloat(v[i]); + } + return floatVector; + } + + @Override + public DocIndexIterator iterator() { + return values.iterator(); + } + + @Override + public Float16AsFloatVectorValues copy() throws IOException { + return new Float16AsFloatVectorValues(values.copy()); + } + } + static final class NormalizedFloatVectorValues extends FloatVectorValues { private final FloatVectorValues values; private final float[] normalizedVector; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java new file mode 100644 index 000000000000..10cc1e3795d4 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java @@ -0,0 +1,398 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.codecs.lucene104; + +import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.deQuantize; + +import java.io.IOException; +import java.nio.ByteBuffer; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene90.IndexedDISI; +import org.apache.lucene.codecs.lucene95.HasIndexSlice; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.packed.DirectMonotonicReader; +import org.apache.lucene.util.quantization.OptimizedScalarQuantizer; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding; + +/** + * Reads quantized vector values from the index input and returns float16 vector values after + * dequantizing them. + * + *

Used for read-only indexes whose raw float16 vectors have been dropped to save storage: only + * the scalar-quantized bytes remain, so {@link #vectorValue(int)} reconstructs float16 values by + * dequantizing them, with some precision loss. + * + * @lucene.internal + */ +abstract class OffHeapScalarQuantizedFloat16VectorValues extends Float16VectorValues + implements HasIndexSlice { + + final int dimension; + final int size; + final VectorSimilarityFunction similarityFunction; + final FlatVectorsScorer vectorsScorer; + + final IndexInput slice; + final short[] vectorValue; + final byte[] byteValue; + final ByteBuffer byteBuffer; + final byte[] unpackedByteVectorValue; + final int byteSize; + private int lastOrd = -1; + final float[] correctiveValues; + int quantizedComponentSum; + final ScalarEncoding encoding; + final float[] centroid; + + OffHeapScalarQuantizedFloat16VectorValues( + int dimension, + int size, + float[] centroid, + ScalarEncoding encoding, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + IndexInput slice) { + this.dimension = dimension; + this.size = size; + this.similarityFunction = similarityFunction; + this.vectorsScorer = vectorsScorer; + this.slice = slice; + this.centroid = centroid; + this.correctiveValues = new float[3]; + this.encoding = encoding; + int docPackedLength = encoding.getDocPackedLength(dimension); + this.byteSize = docPackedLength + (Float.BYTES * 3) + Integer.BYTES; + this.byteBuffer = ByteBuffer.allocate(docPackedLength); + this.vectorValue = new short[dimension]; + this.byteValue = byteBuffer.array(); + this.unpackedByteVectorValue = new byte[dimension]; + } + + @Override + public int dimension() { + return dimension; + } + + @Override + public int size() { + return size; + } + + @Override + public short[] vectorValue(int targetOrd) throws IOException { + if (lastOrd == targetOrd) { + return vectorValue; + } + + // read quantized byte vector, correctiveValues and quantizedComponentSum + slice.seek((long) targetOrd * byteSize); + slice.readBytes(byteBuffer.array(), byteBuffer.arrayOffset(), byteValue.length); + slice.readFloats(correctiveValues, 0, 3); + quantizedComponentSum = slice.readInt(); + + // unpack bytes + switch (encoding) { + case PACKED_NIBBLE -> + OffHeapScalarQuantizedVectorValues.unpackNibbles(byteValue, unpackedByteVectorValue); + case SINGLE_BIT_QUERY_NIBBLE -> + OptimizedScalarQuantizer.unpackBinary(byteValue, unpackedByteVectorValue); + case DIBIT_QUERY_NIBBLE -> + OptimizedScalarQuantizer.untransposeDibit(byteValue, unpackedByteVectorValue); + case UNSIGNED_BYTE, SEVEN_BIT -> { + deQuantize( + byteValue, + vectorValue, + encoding.getBits(), + correctiveValues[0], + correctiveValues[1], + centroid); + lastOrd = targetOrd; + return vectorValue; + } + } + + // dequantize + deQuantize( + unpackedByteVectorValue, + vectorValue, + encoding.getBits(), + correctiveValues[0], + correctiveValues[1], + centroid); + + lastOrd = targetOrd; + return vectorValue; + } + + public OptimizedScalarQuantizer.QuantizationResult getCorrectiveTerms(int targetOrd) + throws IOException { + if (lastOrd == targetOrd) { + return new OptimizedScalarQuantizer.QuantizationResult( + correctiveValues[0], correctiveValues[1], correctiveValues[2], quantizedComponentSum); + } + slice.seek(((long) targetOrd * byteSize) + byteValue.length); + slice.readFloats(correctiveValues, 0, 3); + quantizedComponentSum = slice.readInt(); + return new OptimizedScalarQuantizer.QuantizationResult( + correctiveValues[0], correctiveValues[1], correctiveValues[2], quantizedComponentSum); + } + + @Override + public int getVectorByteLength() { + // Length of the packed quantized vector payload, excluding the corrective terms stored after + // it. This differs from the logical dimension for packed encodings such as PACKED_NIBBLE. + return byteValue.length; + } + + @Override + public IndexInput getSlice() { + return slice; + } + + static OffHeapScalarQuantizedFloat16VectorValues load( + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int size, + ScalarEncoding encoding, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + float[] centroid, + long quantizedVectorDataOffset, + long quantizedVectorDataLength, + IndexInput vectorData) + throws IOException { + if (configuration.isEmpty()) { + return new OffHeapScalarQuantizedFloat16VectorValues.EmptyOffHeapVectorValues( + dimension, similarityFunction, vectorsScorer); + } + assert centroid != null; + IndexInput bytesSlice = + vectorData.slice( + "scalar-quantized-float16-vector-data", + quantizedVectorDataOffset, + quantizedVectorDataLength); + if (configuration.isDense()) { + return new OffHeapScalarQuantizedFloat16VectorValues.DenseOffHeapVectorValues( + dimension, size, centroid, encoding, similarityFunction, vectorsScorer, bytesSlice); + } else { + return new OffHeapScalarQuantizedFloat16VectorValues.SparseOffHeapVectorValues( + configuration, + dimension, + size, + centroid, + encoding, + vectorData, + similarityFunction, + vectorsScorer, + bytesSlice); + } + } + + /** Dense off-heap scalar quantized vector values */ + private static class DenseOffHeapVectorValues extends OffHeapScalarQuantizedFloat16VectorValues { + DenseOffHeapVectorValues( + int dimension, + int size, + float[] centroid, + ScalarEncoding encoding, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + IndexInput slice) { + super(dimension, size, centroid, encoding, similarityFunction, vectorsScorer, slice); + } + + @Override + public OffHeapScalarQuantizedFloat16VectorValues.DenseOffHeapVectorValues copy() + throws IOException { + return new OffHeapScalarQuantizedFloat16VectorValues.DenseOffHeapVectorValues( + dimension, size, centroid, encoding, similarityFunction, vectorsScorer, slice.clone()); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return acceptDocs; + } + + @Override + public VectorScorer scorer(short[] target) throws IOException { + OffHeapScalarQuantizedFloat16VectorValues.DenseOffHeapVectorValues copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer scorer = + vectorsScorer.getRandomVectorScorer(similarityFunction, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return scorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + + @Override + public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) { + return Bulk.fromRandomScorerDense(scorer, iterator, matchingDocs); + } + }; + } + + @Override + public DocIndexIterator iterator() { + return createDenseIterator(); + } + } + + /** Sparse off-heap scalar quantized vector values */ + private static class SparseOffHeapVectorValues extends OffHeapScalarQuantizedFloat16VectorValues { + private final DirectMonotonicReader ordToDoc; + private final IndexedDISI disi; + // dataIn was used to init a new IndexedDISI for #randomAccess() + private final IndexInput dataIn; + private final OrdToDocDISIReaderConfiguration configuration; + + SparseOffHeapVectorValues( + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int size, + float[] centroid, + ScalarEncoding encoding, + IndexInput dataIn, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + IndexInput slice) + throws IOException { + super(dimension, size, centroid, encoding, similarityFunction, vectorsScorer, slice); + this.configuration = configuration; + this.dataIn = dataIn; + this.ordToDoc = configuration.getDirectMonotonicReader(dataIn); + this.disi = configuration.getIndexedDISI(dataIn); + } + + @Override + public OffHeapScalarQuantizedFloat16VectorValues.SparseOffHeapVectorValues copy() + throws IOException { + return new OffHeapScalarQuantizedFloat16VectorValues.SparseOffHeapVectorValues( + configuration, + dimension, + size, + centroid, + encoding, + dataIn, + similarityFunction, + vectorsScorer, + slice.clone()); + } + + @Override + public int ordToDoc(int ord) { + return (int) ordToDoc.get(ord); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + if (acceptDocs == null) { + return null; + } + return new Bits() { + @Override + public boolean get(int index) { + return acceptDocs.get(ordToDoc(index)); + } + + @Override + public int length() { + return size; + } + }; + } + + @Override + public DocIndexIterator iterator() { + return IndexedDISI.asDocIndexIterator(disi); + } + + @Override + public VectorScorer scorer(short[] target) throws IOException { + OffHeapScalarQuantizedFloat16VectorValues.SparseOffHeapVectorValues copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer scorer = + vectorsScorer.getRandomVectorScorer(similarityFunction, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return scorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + + @Override + public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) { + return Bulk.fromRandomScorerSparse(scorer, iterator, matchingDocs); + } + }; + } + } + + /** Empty vector values */ + private static class EmptyOffHeapVectorValues extends OffHeapScalarQuantizedFloat16VectorValues { + EmptyOffHeapVectorValues( + int dimension, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer) { + super( + dimension, + 0, + null, + ScalarEncoding.UNSIGNED_BYTE, + similarityFunction, + vectorsScorer, + null); + } + + @Override + public DocIndexIterator iterator() { + return createDenseIterator(); + } + + @Override + public OffHeapScalarQuantizedFloat16VectorValues.DenseOffHeapVectorValues copy() { + throw new UnsupportedOperationException(); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return null; + } + + @Override + public VectorScorer scorer(short[] target) { + return null; + } + } +} diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedVectorValues.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedVectorValues.java index a91ace8ff9f9..f16ec3726ca5 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedVectorValues.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedVectorValues.java @@ -334,6 +334,31 @@ public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) { }; } + @Override + public VectorScorer scorer(short[] target) throws IOException { + assert isQuerySide == false; + OffHeapScalarQuantizedVectorValues.DenseOffHeapVectorValues copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer scorer = + vectorsScorer.getRandomVectorScorer(similarityFunction, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return scorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + + @Override + public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) { + return Bulk.fromRandomScorerDense(scorer, iterator, matchingDocs); + } + }; + } + @Override public DocIndexIterator iterator() { return createDenseIterator(); @@ -447,6 +472,31 @@ public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) { } }; } + + @Override + public VectorScorer scorer(short[] target) throws IOException { + assert isQuerySide == false; + SparseOffHeapVectorValues copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer scorer = + vectorsScorer.getRandomVectorScorer(similarityFunction, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return scorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + + @Override + public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) { + return Bulk.fromRandomScorerSparse(scorer, iterator, matchingDocs); + } + }; + } } private static class EmptyOffHeapVectorValues extends OffHeapScalarQuantizedVectorValues { @@ -486,5 +536,10 @@ public Bits getAcceptOrds(Bits acceptDocs) { public VectorScorer scorer(float[] target) { return null; } + + @Override + public VectorScorer scorer(short[] target) { + return null; + } } } diff --git a/lucene/core/src/java/org/apache/lucene/index/VectorEncoding.java b/lucene/core/src/java/org/apache/lucene/index/VectorEncoding.java index f4582180f193..33b2bb074b87 100644 --- a/lucene/core/src/java/org/apache/lucene/index/VectorEncoding.java +++ b/lucene/core/src/java/org/apache/lucene/index/VectorEncoding.java @@ -46,4 +46,9 @@ public enum VectorEncoding { VectorEncoding(int byteSize) { this.byteSize = byteSize; } + + /** Returns true if this is a floating-point encoding ({@link #FLOAT32} or {@link #FLOAT16}). */ + public boolean isFloatingPoint() { + return this == FLOAT32 || this == FLOAT16; + } } diff --git a/lucene/core/src/java/org/apache/lucene/util/quantization/BaseQuantizedByteVectorValues.java b/lucene/core/src/java/org/apache/lucene/util/quantization/BaseQuantizedByteVectorValues.java index 5e18fbaee1ea..a609143f3ec8 100644 --- a/lucene/core/src/java/org/apache/lucene/util/quantization/BaseQuantizedByteVectorValues.java +++ b/lucene/core/src/java/org/apache/lucene/util/quantization/BaseQuantizedByteVectorValues.java @@ -40,6 +40,16 @@ public VectorScorer scorer(float[] query) throws IOException { throw new UnsupportedOperationException(); } + /** + * Return a {@link VectorScorer} for the given fp16 query vector. + * + * @param query the query vector + * @return a {@link VectorScorer} instance or null + */ + public VectorScorer scorer(short[] query) throws IOException { + throw new UnsupportedOperationException(); + } + @Override public IndexInput getSlice() { return null; diff --git a/lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java b/lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java index 179799ff83af..6038df8d5ebf 100644 --- a/lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java +++ b/lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java @@ -266,6 +266,36 @@ public static float[] deQuantize( return dequantized; } + /** + * Dequantizes a quantized byte vector back to float16 values. + * + *

Behaves as {@link #deQuantize(byte[], float[], byte, float, float, float[])}, narrowing each + * reconstructed value to float16. + * + * @param quantized the quantized byte vector to dequantize + * @param dequantized the output array to store dequantized float16 bit patterns + * @param bits the number of bits used for quantization + * @param lowerInterval lower value of quantization range + * @param upperInterval upper value of quantization range + * @param centroid the centroid vector that was subtracted during quantization + * @return the dequantized float16 array (same as dequantized parameter) + */ + public static short[] deQuantize( + byte[] quantized, + short[] dequantized, + byte bits, + float lowerInterval, + float upperInterval, + float[] centroid) { + int nSteps = (1 << bits) - 1; + double step = (upperInterval - lowerInterval) / nSteps; + for (int h = 0; h < quantized.length; h++) { + double xi = (double) (quantized[h] & 0xFF) * step + lowerInterval; + dequantized[h] = Float.floatToFloat16((float) (xi + centroid[h])); + } + return dequantized; + } + /** * Compute the loss of the vector given the interval. Effectively, we are computing the MSE of a * dequantized vector with the raw vector. diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java index 7825d92706e5..3b0834704e2d 100644 --- a/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java +++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java @@ -23,33 +23,47 @@ import static org.hamcrest.Matchers.oneOf; import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; import java.util.Locale; +import java.util.Map; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.codecs.FilterCodec; import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloat16VectorField; import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.Float16VectorValues; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloat16VectorQuery; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.TotalHits; +import org.apache.lucene.search.VectorScorer; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.store.BaseDirectoryWrapper; import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.VectorUtil; import org.apache.lucene.util.quantization.OptimizedScalarQuantizer; import org.apache.lucene.util.quantization.QuantizedByteVectorValues; import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding; @@ -105,6 +119,42 @@ public void testSearch() throws Exception { } } + public void testFloat16Search() throws Exception { + String fieldName = "field"; + int numVectors = random().nextInt(99, 500); + int dims = random().nextInt(4, 65); + if (dims % 2 == 1) { + dims++; + } + VectorSimilarityFunction similarityFunction = randomSimilarity(); + KnnFloat16VectorField knnField = + new KnnFloat16VectorField( + fieldName, randomNormalizedFloat16Vector(dims), similarityFunction); + IndexWriterConfig iwc = newIndexWriterConfig(); + try (Directory dir = newDirectory()) { + try (IndexWriter w = new IndexWriter(dir, iwc)) { + 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); + // Routes scoring through Lucene104ScalarQuantizedVectorScorer's short[] (fp16) branch. + Query q = new KnnFloat16VectorQuery(fieldName, queryVector, k); + TopDocs collectedDocs = searcher.search(q, k); + assertEquals(k, collectedDocs.totalHits.value()); + assertEquals(TotalHits.Relation.EQUAL_TO, collectedDocs.totalHits.relation()); + } + } + } + } + public void testToString() { FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) { @@ -214,6 +264,256 @@ public void testQuantizedVectorsWriteAndRead() throws IOException { } } + /** + * fp16 counterpart of {@link #testQuantizedVectorsWriteAndRead()}: indexes float16 vectors and + * verifies the persisted quantized bytes + corrective terms match a reference re-quantization. + * The reference mirrors the writer's fp16 path exactly — inflate fp16->fp32 (normalizing + * for COSINE) before quantizing — so the comparison is byte-exact, not MAE-based. + */ + public void testFloat16QuantizedVectorsWriteAndRead() throws IOException { + String fieldName = "field"; + int numVectors = random().nextInt(99, 500); + int dims = random().nextInt(4, 65); + if (dims % 2 == 1) { + dims++; + } + + VectorSimilarityFunction similarityFunction = randomSimilarity(); + KnnFloat16VectorField knnField = + new KnnFloat16VectorField( + fieldName, randomNormalizedFloat16Vector(dims), similarityFunction); + try (Directory dir = newDirectory()) { + try (IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + for (int i = 0; i < numVectors; i++) { + Document doc = new Document(); + knnField.setVectorValue(randomNormalizedFloat16Vector(dims)); + doc.add(knnField); + w.addDocument(doc); + if (i % 101 == 0) { + w.commit(); + } + } + w.commit(); + w.forceMerge(1); + + try (IndexReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + Float16VectorValues vectorValues = r.getFloat16VectorValues(fieldName); + assertEquals(vectorValues.size(), numVectors); + QuantizedByteVectorValues qvectorValues = + ((Lucene104ScalarQuantizedVectorsReader.ScalarQuantizedFloat16VectorValues) + vectorValues) + .getQuantizedVectorValues(); + float[] centroid = qvectorValues.getCentroid(); + assertEquals(centroid.length, dims); + + OptimizedScalarQuantizer quantizer = new OptimizedScalarQuantizer(similarityFunction); + byte[] scratch = new byte[encoding.getDiscreteDimensions(dims)]; + byte[] expectedVector = new byte[encoding.getDocPackedLength(scratch.length)]; + float[] inflated = new float[dims]; + + KnnVectorValues.DocIndexIterator docIndexIterator = vectorValues.iterator(); + while (docIndexIterator.nextDoc() != NO_MORE_DOCS) { + // Reproduce the writer's fp16 reference: inflate to fp32, normalize for COSINE. + short[] raw = vectorValues.vectorValue(docIndexIterator.index()); + for (int i = 0; i < dims; i++) { + inflated[i] = Float.float16ToFloat(raw[i]); + } + if (similarityFunction == VectorSimilarityFunction.COSINE) { + VectorUtil.l2normalize(inflated); + } + OptimizedScalarQuantizer.QuantizationResult corrections = + quantizer.scalarQuantize(inflated, scratch, encoding.getBits(), centroid); + switch (encoding) { + case UNSIGNED_BYTE, SEVEN_BIT -> + System.arraycopy(scratch, 0, expectedVector, 0, dims); + case PACKED_NIBBLE -> + OffHeapScalarQuantizedVectorValues.packNibbles(scratch, expectedVector); + case SINGLE_BIT_QUERY_NIBBLE -> + OptimizedScalarQuantizer.packAsBinary(scratch, expectedVector); + case DIBIT_QUERY_NIBBLE -> + OptimizedScalarQuantizer.transposeDibit(scratch, expectedVector); + } + assertArrayEquals(expectedVector, qvectorValues.vectorValue(docIndexIterator.index())); + var actualCorrections = qvectorValues.getCorrectiveTerms(docIndexIterator.index()); + assertEquals(corrections.lowerInterval(), actualCorrections.lowerInterval(), 0.00001f); + assertEquals(corrections.upperInterval(), actualCorrections.upperInterval(), 0.00001f); + assertEquals( + corrections.additionalCorrection(), + actualCorrections.additionalCorrection(), + 0.00001f); + assertEquals( + corrections.quantizedComponentSum(), actualCorrections.quantizedComponentSum()); + } + } + } + } + } + + /** + * Reads float16 vectors back from an index whose raw {@code .vec} data has been dropped, so + * values are reconstructed through {@link OffHeapScalarQuantizedFloat16VectorValues}, and asserts + * they stay within the quantization error bound. + */ + public void testReadQuantizedFloat16VectorWithEmptyRawVectors() throws Exception { + String vectorFieldName = "vec1"; + int numVectors = 1 + random().nextInt(50); + int dim = random().nextInt(64) + 1; + if (dim % 2 == 1) { + dim++; + } + // Quantization error bound, plus a small slack for the extra fp16 rounding applied on top of + // quantization (both the stored input and the dequantized output are fp16-rounded). + float eps = (1f / (float) (1 << getQuantizationBits())) + 1e-3f; + VectorSimilarityFunction similarityFunction = randomSimilarity(); + + // Build fp16 (short-bit) source vectors; keep them to form the MAE reference on read-back. + List vectors = new ArrayList<>(numVectors); + for (int i = 0; i < numVectors; i++) { + vectors.add(randomNormalizedFloat16Vector(dim)); + } + + try (BaseDirectoryWrapper dir = newDirectory()) { + dir.setCheckIndexOnClose(false); // raw .vec is deliberately emptied below + + try (IndexWriter w = + new IndexWriter( + dir, + new IndexWriterConfig() + .setMaxBufferedDocs(numVectors + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE) + .setUseCompoundFile(false) + .setCodec(getCodecForFloatVectorFallbackTest()))) { + for (int i = 0; i < numVectors; i++) { + Document doc = new Document(); + doc.add(new KnnFloat16VectorField(vectorFieldName, vectors.get(i), similarityFunction)); + w.addDocument(doc); + } + } + + // Drop the raw float16 vectors, leaving only the quantized data. + simulateEmptyRawVectors(dir); + + try (IndexReader reader = DirectoryReader.open(dir)) { + LeafReader r = getOnlyLeafReader(reader); + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + knnVectorsReader = knnVectorsReader.unwrapReaderForField(vectorFieldName); + // With raw vectors dropped this routes through OffHeapScalarQuantizedFloat16VectorValues. + Float16VectorValues float16VectorValues = + knnVectorsReader.getFloat16VectorValues(vectorFieldName); + if (float16VectorValues.size() > 0) { + KnnVectorValues.DocIndexIterator iter = float16VectorValues.iterator(); + for (int docId = iter.nextDoc(); docId != NO_MORE_DOCS; docId = iter.nextDoc()) { + short[] dequantizedVector = float16VectorValues.vectorValue(iter.index()); + short[] originalVector = vectors.get(docId); + float mae = 0; + for (int i = 0; i < dim; i++) { + mae += + Math.abs( + Float.float16ToFloat(dequantizedVector[i]) + - Float.float16ToFloat(originalVector[i])); + } + mae /= dim; + assertTrue( + "bits: " + getQuantizationBits() + " mae: " + mae + " > eps: " + eps, mae <= eps); + } + } else { + fail("float16VectorValues size should be non zero"); + } + } else { + fail("reader is not CodecReader"); + } + } + } + } + + /** + * Tests that dropping the raw float16 vectors does not change scoring. {@link + * Float16VectorValues#scorer(short[])} is documented to score against the quantized vectors when + * the underlying format quantizes, so a quantized index must produce identical scores before and + * after its raw vector file is emptied. + */ + public void testFloat16ScoresUnchangedWithEmptyRawVectors() throws Exception { + String vectorFieldName = "vec1"; + int numVectors = 1 + random().nextInt(50); + int dim = random().nextInt(64) + 1; + if (dim % 2 == 1) { + dim++; + } + VectorSimilarityFunction similarityFunction = randomSimilarity(); + short[] query = randomNormalizedFloat16Vector(dim); + + try (BaseDirectoryWrapper dir = newDirectory()) { + dir.setCheckIndexOnClose(false); // raw .vec is deliberately emptied below + + try (IndexWriter w = + new IndexWriter( + dir, + new IndexWriterConfig() + .setMaxBufferedDocs(numVectors + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE) + .setUseCompoundFile(false) + .setCodec(getCodecForFloatVectorFallbackTest()))) { + for (int i = 0; i < numVectors; i++) { + Document doc = new Document(); + doc.add( + new KnnFloat16VectorField( + vectorFieldName, randomNormalizedFloat16Vector(dim), similarityFunction)); + w.addDocument(doc); + } + } + + // Scores while the raw float16 vectors are still present. + Map expectedScores = scoreAllFloat16Docs(dir, vectorFieldName, query); + assertEquals("expected every document to be scored", numVectors, expectedScores.size()); + + simulateEmptyRawVectors(dir); + + // Both reads are expected to score against the same quantized vectors, so the scores must + // match exactly rather than merely within a quantization-error tolerance. + Map actualScores = scoreAllFloat16Docs(dir, vectorFieldName, query); + assertEquals(expectedScores.keySet(), actualScores.keySet()); + for (Map.Entry entry : expectedScores.entrySet()) { + assertEquals( + "score changed for doc " + entry.getKey() + " after dropping raw vectors", + entry.getValue(), + actualScores.get(entry.getKey()), + 0f); + } + } + } + + /** + * Scores every document holding a float16 vector for {@code field} against {@code query}, keyed + * by doc id. Keying by doc id rather than ordinal keeps the comparison meaningful even when a + * different {@link Float16VectorValues} implementation backs the iterator. + */ + private Map scoreAllFloat16Docs(Directory dir, String field, short[] query) + throws IOException { + Map scores = new HashMap<>(); + try (IndexReader reader = DirectoryReader.open(dir)) { + LeafReader leafReader = getOnlyLeafReader(reader); + if (leafReader instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = + codecReader.getVectorReader().unwrapReaderForField(field); + Float16VectorValues float16VectorValues = knnVectorsReader.getFloat16VectorValues(field); + assertNotNull(float16VectorValues); + VectorScorer scorer = float16VectorValues.scorer(query); + assertNotNull(scorer); + DocIdSetIterator iterator = scorer.iterator(); + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + scores.put(doc, scorer.score()); + } + } else { + fail("reader is not CodecReader"); + } + } + return scores; + } + @Override protected boolean supportsFloatVectorFallback() { return true; diff --git a/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java b/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java index 510008a4b313..917413bdd17a 100644 --- a/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java +++ b/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java @@ -78,6 +78,44 @@ public void testQuantizationQuality() { } } + public void testDeQuantizeFloat16MatchesFloat() { + int dims = 16; + int numVectors = 32; + float[][] vectors = new float[numVectors][]; + float[] centroid = new float[dims]; + for (int i = 0; i < numVectors; ++i) { + vectors[i] = new float[dims]; + for (int j = 0; j < dims; ++j) { + vectors[i][j] = randomFloat(); + centroid[j] += vectors[i][j]; + } + } + for (int j = 0; j < dims; ++j) { + centroid[j] /= numVectors; + } + OptimizedScalarQuantizer osq = + new OptimizedScalarQuantizer(VectorSimilarityFunction.DOT_PRODUCT); + float[] scratch = new float[dims]; + for (byte bit : ALL_BITS) { + byte[] destination = new byte[dims]; + for (int i = 0; i < numVectors; ++i) { + System.arraycopy(vectors[i], 0, scratch, 0, dims); + OptimizedScalarQuantizer.QuantizationResult result = + osq.scalarQuantize(scratch, destination, bit, centroid); + float[] floatDeq = new float[dims]; + deQuantize( + destination, floatDeq, bit, result.lowerInterval(), result.upperInterval(), centroid); + short[] shortDeq = new short[dims]; + deQuantize( + destination, shortDeq, bit, result.lowerInterval(), result.upperInterval(), centroid); + // The short (fp16) overload must equal the float overload rounded to fp16, component-wise. + for (int k = 0; k < dims; ++k) { + assertEquals(Float.floatToFloat16(floatDeq[k]), shortDeq[k]); + } + } + } + } + public void testAbusiveEdgeCases() { // large zero array for (VectorSimilarityFunction vectorSimilarityFunction : VectorSimilarityFunction.values()) {