kaivalnp commented on code in PR #16030: URL: https://github.com/apache/lucene/pull/16030#discussion_r3905675929
########## lucene/core/src/java/org/apache/lucene/codecs/lucene106/Lucene106ScalarQuantizedVectorsWriter.java: ########## @@ -0,0 +1,1104 @@ +/* + * 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.lucene106; + +import static org.apache.lucene.codecs.lucene106.Lucene106ScalarQuantizedVectorsFormat.DIRECT_MONOTONIC_BLOCK_SHIFT; +import static org.apache.lucene.codecs.lucene106.Lucene106ScalarQuantizedVectorsFormat.QUANTIZED_VECTOR_COMPONENT; +import static org.apache.lucene.index.VectorSimilarityFunction.COSINE; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; +import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.transposeHalfByte; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.DocIDMerger; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergeState; +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.ArrayUtil; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; +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; + +/** + * Writes quantized vector values and metadata to index segments in the format for Lucene 10.5. + * + * @lucene.experimental + */ +public class Lucene106ScalarQuantizedVectorsWriter extends FlatVectorsWriter { + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(Lucene106ScalarQuantizedVectorsWriter.class); + + private final SegmentWriteState segmentWriteState; + private final List<FieldWriter> fields = new ArrayList<>(); + private final IndexOutput meta, vectorData; + private final ScalarEncoding encoding; + private final boolean enableCentering; + private final FlatVectorsWriter rawVectorDelegate; + private boolean finished; + + /** Sole constructor */ + public Lucene106ScalarQuantizedVectorsWriter( + SegmentWriteState state, + ScalarEncoding encoding, + boolean enableCentering, + FlatVectorsWriter rawVectorDelegate, + Lucene106ScalarQuantizedVectorScorer vectorsScorer) + throws IOException { + super(vectorsScorer); + this.encoding = encoding; + this.enableCentering = enableCentering; + this.segmentWriteState = state; + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, + state.segmentSuffix, + Lucene106ScalarQuantizedVectorsFormat.META_EXTENSION); + + String vectorDataFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, + state.segmentSuffix, + Lucene106ScalarQuantizedVectorsFormat.VECTOR_DATA_EXTENSION); + this.rawVectorDelegate = rawVectorDelegate; + try { + meta = state.directory.createOutput(metaFileName, state.context); + vectorData = state.directory.createOutput(vectorDataFileName, state.context); + + CodecUtil.writeIndexHeader( + meta, + Lucene106ScalarQuantizedVectorsFormat.META_CODEC_NAME, + Lucene106ScalarQuantizedVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + vectorData, + Lucene106ScalarQuantizedVectorsFormat.VECTOR_DATA_CODEC_NAME, + Lucene106ScalarQuantizedVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + } catch (Throwable t) { + IOUtils.closeWhileSuppressingExceptions(t, this); + throw t; + } + } + + @Override + public FlatFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOException { + if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) { Review Comment: #16473 recently added scalar-quantization support for `FP16` vectors, should we include that too? ########## lucene/core/src/java/org/apache/lucene/codecs/lucene106/Lucene106ScalarQuantizedVectorScorer.java: ########## @@ -0,0 +1,302 @@ +/* + * 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.lucene106; + +import static org.apache.lucene.index.VectorSimilarityFunction.COSINE; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; + +import java.io.IOException; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.VectorUtil; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer; +import org.apache.lucene.util.quantization.OptimizedScalarQuantizer; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding; + +/** + * Vector scorer over OptimizedScalarQuantized vectors + * + * @lucene.experimental + */ +public class Lucene106ScalarQuantizedVectorScorer implements FlatVectorsScorer { + private final FlatVectorsScorer nonQuantizedDelegate; + + public Lucene106ScalarQuantizedVectorScorer(FlatVectorsScorer nonQuantizedDelegate) { + this.nonQuantizedDelegate = nonQuantizedDelegate; + } + + @Override + public RandomVectorScorerSupplier getRandomVectorScorerSupplier( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues) + throws IOException { + if (vectorValues instanceof QuantizedByteVectorValues qv) { + return new ScalarQuantizedVectorScorerSupplier(qv, similarityFunction); + } + // It is possible to get to this branch during initial indexing and flush + return nonQuantizedDelegate.getRandomVectorScorerSupplier(similarityFunction, vectorValues); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, float[] 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)]; + } + // We make a copy as the quantization process mutates the input + float[] copy = ArrayUtil.copyOfSubArray(target, 0, target.length); + if (similarityFunction == COSINE) { + VectorUtil.l2normalize(copy); + } + target = copy; + var targetCorrectiveTerms = + quantizer.scalarQuantize( + target, 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); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, short[] target) + throws IOException { + FlatVectorsScorer.checkDimensions(target.length, vectorValues.dimension()); + return nonQuantizedDelegate.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, byte[] target) + throws IOException { + FlatVectorsScorer.checkDimensions(target.length, vectorValues.dimension()); + return nonQuantizedDelegate.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + public RandomVectorScorerSupplier getRandomVectorScorerSupplier( + VectorSimilarityFunction similarityFunction, + QuantizedByteVectorValues scoringVectors, + QuantizedByteVectorValues targetVectors) { + return new AsymmetricQuantizedRandomVectorScorerSupplier( + scoringVectors, targetVectors, similarityFunction); + } + + @Override + public String toString() { + return "Lucene106ScalarQuantizedVectorScorer(nonQuantizedDelegate=" + + nonQuantizedDelegate + + ")"; + } + + static class AsymmetricQuantizedRandomVectorScorerSupplier implements RandomVectorScorerSupplier { + private final QuantizedByteVectorValues queryVectors; + private final QuantizedByteVectorValues targetVectors; + private final VectorSimilarityFunction similarityFunction; + + AsymmetricQuantizedRandomVectorScorerSupplier( + QuantizedByteVectorValues queryVectors, + QuantizedByteVectorValues targetVectors, + VectorSimilarityFunction similarityFunction) { + assert targetVectors.getScalarEncoding().isAsymmetric(); + this.queryVectors = queryVectors; + this.targetVectors = targetVectors; + this.similarityFunction = similarityFunction; + } + + @Override + public UpdateableRandomVectorScorer scorer() throws IOException { + final QuantizedByteVectorValues targetVectors = this.targetVectors.copy(); + final QuantizedByteVectorValues queryVectors = this.queryVectors.copy(); + return new UpdateableRandomVectorScorer.AbstractUpdateableRandomVectorScorer(targetVectors) { + private OptimizedScalarQuantizer.QuantizationResult queryCorrections = null; + private byte[] vector = null; + + @Override + public void setScoringOrdinal(int node) throws IOException { + vector = queryVectors.vectorValue(node); + queryCorrections = queryVectors.getCorrectiveTerms(node); + } + + @Override + public float score(int node) throws IOException { + if (vector == null || queryCorrections == null) { + throw new IllegalStateException("setScoringOrdinal was not called"); + } + + return quantizedScore(vector, queryCorrections, targetVectors, node, similarityFunction); + } + }; + } + + @Override + public RandomVectorScorerSupplier copy() throws IOException { + return new AsymmetricQuantizedRandomVectorScorerSupplier( + queryVectors.copy(), targetVectors.copy(), similarityFunction); + } + } + + private static final class ScalarQuantizedVectorScorerSupplier + implements RandomVectorScorerSupplier { + private final QuantizedByteVectorValues targetValues; + private final QuantizedByteVectorValues values; + private final VectorSimilarityFunction similarity; + + public ScalarQuantizedVectorScorerSupplier( + QuantizedByteVectorValues values, VectorSimilarityFunction similarity) throws IOException { + assert values.getScalarEncoding().isAsymmetric() == false; + this.targetValues = values.copy(); + this.values = values; + this.similarity = similarity; + } + + @Override + public UpdateableRandomVectorScorer scorer() throws IOException { + return new UpdateableRandomVectorScorer.AbstractUpdateableRandomVectorScorer(values) { + private byte[] targetVector; + private OptimizedScalarQuantizer.QuantizationResult targetCorrectiveTerms; + + @Override + public float score(int node) throws IOException { + return quantizedScore(targetVector, targetCorrectiveTerms, values, node, similarity); + } + + @Override + public void setScoringOrdinal(int node) throws IOException { + var rawTargetVector = targetValues.vectorValue(node); + switch (values.getScalarEncoding()) { + case UNSIGNED_BYTE, SEVEN_BIT -> targetVector = rawTargetVector; + case PACKED_NIBBLE -> { + if (targetVector == null) { + targetVector = new byte[OptimizedScalarQuantizer.discretize(values.dimension(), 2)]; + } + OffHeapScalarQuantizedVectorValues.unpackNibbles(rawTargetVector, targetVector); + } + case SINGLE_BIT_QUERY_NIBBLE, DIBIT_QUERY_NIBBLE -> { + throw new IllegalStateException( + values.getScalarEncoding().name() + + " encoding is not supported for symmetric quantization"); + } + } + targetCorrectiveTerms = targetValues.getCorrectiveTerms(node); + } + }; + } + + @Override + public RandomVectorScorerSupplier copy() throws IOException { + return new ScalarQuantizedVectorScorerSupplier(values.copy(), similarity); + } + } + + private static final float[] SCALE_LUT = + new float[] { + 1f, + 1f / ((1 << 2) - 1), + 1f / ((1 << 3) - 1), + 1f / ((1 << 4) - 1), + 1f / ((1 << 5) - 1), + 1f / ((1 << 6) - 1), + 1f / ((1 << 7) - 1), + 1f / ((1 << 8) - 1), + }; + + private static float quantizedScore( + byte[] quantizedQuery, + OptimizedScalarQuantizer.QuantizationResult queryCorrections, + QuantizedByteVectorValues targetVectors, + int targetOrd, + VectorSimilarityFunction similarityFunction) + throws IOException { + var scalarEncoding = targetVectors.getScalarEncoding(); + byte[] quantizedDoc = targetVectors.vectorValue(targetOrd); + float qcDist = + switch (scalarEncoding) { + case UNSIGNED_BYTE -> VectorUtil.uint8DotProduct(quantizedQuery, quantizedDoc); + case SEVEN_BIT -> VectorUtil.dotProduct(quantizedQuery, quantizedDoc); + case PACKED_NIBBLE -> VectorUtil.int4DotProductSinglePacked(quantizedQuery, quantizedDoc); + case SINGLE_BIT_QUERY_NIBBLE -> + VectorUtil.int4BitDotProduct(quantizedQuery, quantizedDoc); + case DIBIT_QUERY_NIBBLE -> VectorUtil.int4DibitDotProduct(quantizedQuery, quantizedDoc); + }; + OptimizedScalarQuantizer.QuantizationResult indexCorrections = + targetVectors.getCorrectiveTerms(targetOrd); + float queryScale = SCALE_LUT[scalarEncoding.getQueryBits() - 1]; + float scale = SCALE_LUT[scalarEncoding.getBits() - 1]; + float x1 = indexCorrections.quantizedComponentSum(); + float ax = indexCorrections.lowerInterval(); + // Here we must scale according to the bits + float lx = (indexCorrections.upperInterval() - ax) * scale; + float ay = queryCorrections.lowerInterval(); + float ly = (queryCorrections.upperInterval() - ay) * queryScale; + float y1 = queryCorrections.quantizedComponentSum(); + float score = + ax * ay * targetVectors.dimension() + ay * lx * x1 + ax * ly * y1 + lx * ly * qcDist; + // For euclidean, we need to invert the score and apply the additional correction, which is + // assumed to be the squared l2norm of the centroid centered vectors. + if (similarityFunction == EUCLIDEAN) { + score = + queryCorrections.additionalCorrection() + + indexCorrections.additionalCorrection() + - 2 * score; + // Ensure that 'score' (the squared euclidean distance) is non-negative. The computed value + // may be negative as a result of quantization loss. + return 1 / (1f + Math.max(score, 0f)); Review Comment: Can we use [existing utility functions](https://github.com/apache/lucene/blob/d8ffa2a68668f027f697bcbe37bf4d1bbcf014bf/lucene/core/src/java/org/apache/lucene/util/VectorUtil.java#L453-L464) for this? Same for [dot product](https://github.com/apache/lucene/blob/d8ffa2a68668f027f697bcbe37bf4d1bbcf014bf/lucene/core/src/java/org/apache/lucene/util/VectorUtil.java#L439-L451). ########## lucene/core/src/java/module-info.java: ########## @@ -87,7 +88,9 @@ provides org.apache.lucene.codecs.KnnVectorsFormat with org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat, org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat, - org.apache.lucene.codecs.lucene104.Lucene104HnswScalarQuantizedVectorsFormat; Review Comment: Should we also delete the writer + move it to `backward-codecs`? ########## lucene/core/src/java/org/apache/lucene/codecs/lucene106/package-info.java: ########## @@ -0,0 +1,22 @@ +/* + * 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. + */ + +/** + * Lucene 10.5 scalar quantized vector format, extending 10.4 with data-blind mode ({@code Review Comment: `10.5` -> `10.6` ########## lucene/core/src/java/org/apache/lucene/codecs/lucene106/Lucene106ScalarQuantizedVectorsReader.java: ########## @@ -0,0 +1,738 @@ +/* + * 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.lucene106; + +import static org.apache.lucene.codecs.lucene106.Lucene106ScalarQuantizedVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.transposeHalfByte; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Stream; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.store.ChecksumIndexInput; +import org.apache.lucene.store.DataAccessHint; +import org.apache.lucene.store.FileDataHint; +import org.apache.lucene.store.FileTypeHint; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.quantization.OptimizedScalarQuantizer; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding; +import org.apache.lucene.util.quantization.QuantizedVectorsReader; +import org.apache.lucene.util.quantization.ScalarQuantizer; + +/** + * Reader for scalar quantized vectors in the Lucene 10.5 format. + * + * @lucene.experimental + */ +public class Lucene106ScalarQuantizedVectorsReader extends FlatVectorsReader + implements QuantizedVectorsReader { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Lucene106ScalarQuantizedVectorsReader.class); + + private final Map<String, FieldEntry> fields = new HashMap<>(); + private final IndexInput quantizedVectorData; + private final FlatVectorsReader rawVectorsReader; + private final Lucene106ScalarQuantizedVectorScorer vectorScorer; + public static final int EXHAUSTIVE_BULK_SCORE_ORDS = 64; + + public Lucene106ScalarQuantizedVectorsReader( + SegmentReadState state, + FlatVectorsReader rawVectorsReader, + Lucene106ScalarQuantizedVectorScorer vectorsScorer) + throws IOException { + // Quantized vectors are accessed randomly from their node ID stored in the HNSW + // graph. + this(state, rawVectorsReader, vectorsScorer, DataAccessHint.RANDOM); + } + + public Lucene106ScalarQuantizedVectorsReader( + SegmentReadState state, + FlatVectorsReader rawVectorsReader, + Lucene106ScalarQuantizedVectorScorer vectorsScorer, + DataAccessHint accessHint) + throws IOException { + this.vectorScorer = vectorsScorer; + this.rawVectorsReader = rawVectorsReader; + int versionMeta = -1; + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, + state.segmentSuffix, + Lucene106ScalarQuantizedVectorsFormat.META_EXTENSION); + try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName)) { + Throwable priorE = null; + try { + versionMeta = + CodecUtil.checkIndexHeader( + meta, + Lucene106ScalarQuantizedVectorsFormat.META_CODEC_NAME, + Lucene106ScalarQuantizedVectorsFormat.VERSION_START, + Lucene106ScalarQuantizedVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + readFields(meta, state.fieldInfos); + } catch (Throwable exception) { + priorE = exception; + } finally { + CodecUtil.checkFooter(meta, priorE); + } + + final IOContext.FileOpenHint[] hints = + Stream.of(FileTypeHint.DATA, FileDataHint.KNN_VECTORS, accessHint) + .filter(Objects::nonNull) + .toArray(IOContext.FileOpenHint[]::new); + quantizedVectorData = + openDataInput( + state, + versionMeta, + VECTOR_DATA_EXTENSION, + Lucene106ScalarQuantizedVectorsFormat.VECTOR_DATA_CODEC_NAME, + state.context.withHints(hints)); + } catch (Throwable t) { + IOUtils.closeWhileSuppressingExceptions(t, this); + throw t; + } + } + + private void readFields(ChecksumIndexInput meta, FieldInfos infos) throws IOException { + for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) { + FieldInfo info = infos.fieldInfo(fieldNumber); + if (info == null) { + throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta); + } + FieldEntry fieldEntry = readField(meta, info); + validateFieldEntry(info, fieldEntry); + fields.put(info.name, fieldEntry); + } + } + + static void validateFieldEntry(FieldInfo info, FieldEntry fieldEntry) { + int dimension = info.getVectorDimension(); + if (dimension != fieldEntry.dimension) { + throw new IllegalStateException( + "Inconsistent vector dimension for field=\"" + + info.name + + "\"; " + + dimension + + " != " + + fieldEntry.dimension); + } + + long numQuantizedVectorBytes = + Math.multiplyExact( + (fieldEntry.scalarEncoding.getDocPackedLength(dimension) + + (Float.BYTES * 3) + + Integer.BYTES), + (long) fieldEntry.size); + if (numQuantizedVectorBytes != fieldEntry.vectorDataLength) { + throw new IllegalStateException( + "vector data length " + + fieldEntry.vectorDataLength + + " not matching size = " + + fieldEntry.size + + " * (dims=" + + dimension + + " + 16" + + ") = " + + numQuantizedVectorBytes); + } + } + + @Override + public FlatVectorsScorer getFlatVectorScorer(String field) throws IOException { + return vectorScorer; + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException { + 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 + public RandomVectorScorer getRandomVectorScorer(String field, short[] target) throws IOException { + return rawVectorsReader.getRandomVectorScorer(field, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException { + return rawVectorsReader.getRandomVectorScorer(field, target); + } + + @Override + public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException { + rawVectorsReader.checkIntegrity(merge); + CodecUtil.checksumEntireFile(quantizedVectorData); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + FieldEntry fi = fields.get(field); + if (fi == null) { + return null; + } + if (fi.vectorEncoding != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException( + "field=\"" + + field + + "\" is encoded as: " + + fi.vectorEncoding + + " expected: " + + VectorEncoding.FLOAT32); + } + + FloatVectorValues rawFloatVectorValues = + fi.isDataBlind() ? null : rawVectorsReader.getFloatVectorValues(field); + + if (rawFloatVectorValues == null) { + // Data-blind mode: full-precision float vectors were never stored. Reconstruct floats from + // the quantized data. + return OffHeapScalarQuantizedFloatVectorValues.load( + fi.ordToDocDISIReaderConfiguration, + fi.dimension, + fi.size, + fi.scalarEncoding, + fi.similarityFunction, + vectorScorer, + fi.centroid, + fi.vectorDataOffset, + fi.vectorDataLength, + quantizedVectorData); + } + + 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 (rawFloatVectorValues.size() == 0) { + // Full-precision vectors were dropped after writing. Wrap the dequantizing read view with + // sqvv so scorer() stays quantized (as in the branch where raw vectors are present) while + // vectorValue()/rescorer() dequantize. Passing the dequantized view straight to a scorer + // would misroute to the non-quantized flat scorer over the quantized slice. + FloatVectorValues dequantizedRawVectorValues = + OffHeapScalarQuantizedFloatVectorValues.load( + fi.ordToDocDISIReaderConfiguration, + fi.dimension, + fi.size, + fi.scalarEncoding, + fi.similarityFunction, + vectorScorer, + fi.centroid, + fi.vectorDataOffset, + fi.vectorDataLength, + quantizedVectorData); + return new ScalarQuantizedVectorValues(dequantizedRawVectorValues, sqvv); + } + + return new ScalarQuantizedVectorValues(rawFloatVectorValues, sqvv); + } + + @Override + public Float16VectorValues getFloat16VectorValues(String field) throws IOException { + return rawVectorsReader.getFloat16VectorValues(field); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return rawVectorsReader.getByteVectorValues(field); + } + + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) + throws IOException { + rawVectorsReader.search(field, target, knnCollector, acceptDocs); + } + + @Override + public void search(String field, short[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) + throws IOException { + rawVectorsReader.search(field, target, knnCollector, acceptDocs); + } + + @Override + public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) + throws IOException { + if (knnCollector.k() == 0) return; + final RandomVectorScorer scorer = getRandomVectorScorer(field, target); + 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; + int numVectors = scorer.maxOrd(); + for (int i = 0; i < numVectors; i++) { + if (acceptedOrds == null || acceptedOrds.get(i)) { + if (knnCollector.earlyTerminated()) { + break; + } + ords[numOrds++] = i; + if (numOrds == ords.length) { + knnCollector.incVisitedCount(numOrds); + if (scorer.bulkScore(ords, scores, numOrds) > knnCollector.minCompetitiveSimilarity()) { + for (int j = 0; j < numOrds; j++) { + knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]); + } + } + numOrds = 0; + } + } + } + + if (numOrds > 0) { + knnCollector.incVisitedCount(numOrds); + if (scorer.bulkScore(ords, scores, numOrds) > knnCollector.minCompetitiveSimilarity()) { + for (int j = 0; j < numOrds; j++) { + knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]); + } + } + } + } + + @Override + public void close() throws IOException { + IOUtils.close(quantizedVectorData, rawVectorsReader); + } + + @Override + public long ramBytesUsed() { + long size = SHALLOW_SIZE; + size += + RamUsageEstimator.sizeOfMap( + fields, RamUsageEstimator.shallowSizeOfInstance(FieldEntry.class)); + size += rawVectorsReader.ramBytesUsed(); + return size; + } + + @Override + public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) { + Objects.requireNonNull(fieldInfo); + var raw = rawVectorsReader.getOffHeapByteSize(fieldInfo); + var fieldEntry = fields.get(fieldInfo.name); + if (fieldEntry == null) { + // 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; + return raw; + } + var quant = Map.of(VECTOR_DATA_EXTENSION, fieldEntry.vectorDataLength()); + return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, quant); + } + + public float[] getCentroid(String field) { + FieldEntry fieldEntry = fields.get(field); + if (fieldEntry != null) { + return fieldEntry.centroid; + } + return null; + } + + boolean hasRawFloatVectors(String field) throws IOException { + FieldEntry fi = fields.get(field); + if (fi == null || fi.isDataBlind()) { + return false; + } + FloatVectorValues raw = rawVectorsReader.getFloatVectorValues(field); + return raw != null && raw.size() > 0; + } + + private static IndexInput openDataInput( + SegmentReadState state, + int versionMeta, + String fileExtension, + String codecName, + IOContext context) + throws IOException { + String fileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, fileExtension); + IndexInput in = state.directory.openInput(fileName, context); + try { + int versionVectorData = + CodecUtil.checkIndexHeader( + in, + codecName, + Lucene106ScalarQuantizedVectorsFormat.VERSION_START, + Lucene106ScalarQuantizedVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + if (versionMeta != versionVectorData) { + throw new CorruptIndexException( + "Format versions mismatch: meta=" + + versionMeta + + ", " + + codecName + + "=" + + versionVectorData, + in); + } + CodecUtil.retrieveChecksum(in); + return in; + } catch (Throwable t) { + IOUtils.closeWhileSuppressingExceptions(t, in); + throw t; + } + } + + private FieldEntry readField(IndexInput input, FieldInfo info) throws IOException { + VectorEncoding vectorEncoding = readVectorEncoding(input); + VectorSimilarityFunction similarityFunction = readSimilarityFunction(input); + if (similarityFunction != info.getVectorSimilarityFunction()) { + throw new IllegalStateException( + "Inconsistent vector similarity function for field=\"" + + info.name + + "\"; " + + similarityFunction + + " != " + + info.getVectorSimilarityFunction()); + } + return FieldEntry.create(input, vectorEncoding, info.getVectorSimilarityFunction()); + } + + @Override + public QuantizedByteVectorValues getQuantizedVectorValues(String field) throws IOException { + FieldEntry fi = fields.get(field); + if (fi == null) { + return null; + } + if (fi.vectorEncoding != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException( + "field=\"" + + field + + "\" is encoded as: " + + fi.vectorEncoding + + " expected: " + + VectorEncoding.FLOAT32); + } + return 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); + } + + @Override + public ScalarQuantizer getQuantizationState(String fieldName) { + return null; + } + + @Override + public CloseableRandomVectorScorerSupplier getRandomVectorScorerSupplierForMerge( + FieldInfo fieldInfo, SegmentWriteState segmentWriteState) throws IOException { + FieldEntry fi = fields.get(fieldInfo.name); + if (fi == null) { + return null; + } + QuantizedByteVectorValues vectorValues = getQuantizedVectorValues(fieldInfo.name); + if (fi.scalarEncoding.isAsymmetric() == false) { + RandomVectorScorerSupplier supplier = + vectorScorer.getRandomVectorScorerSupplier( + fieldInfo.getVectorSimilarityFunction(), vectorValues); + return CloseableRandomVectorScorerSupplier.create(supplier, vectorValues.size(), () -> {}); + } + FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo.name); + OptimizedScalarQuantizer quantizer = + new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction()); + String tempScoreQuantizedVectorName = null; + DocsWithFieldSet docsWithField; + try (IndexOutput tempScoreQuantizedVector = + segmentWriteState.directory.createTempOutput( + segmentWriteState.segmentInfo.name, "queries", segmentWriteState.context)) { + tempScoreQuantizedVectorName = tempScoreQuantizedVector.getName(); + docsWithField = + writeBinarizedQueryData( + vectorValues, + fi.scalarEncoding, + tempScoreQuantizedVector, + floatVectorValues, + quantizer); + CodecUtil.writeFooter(tempScoreQuantizedVector); + } catch (Throwable t) { + if (tempScoreQuantizedVectorName != null) { + IOUtils.deleteFilesSuppressingExceptions( + t, segmentWriteState.directory, tempScoreQuantizedVectorName); + } + throw t; + } + IndexInput quantizedScoreDataInput = + segmentWriteState.directory.openInput( + tempScoreQuantizedVectorName, segmentWriteState.context); + try { + OffHeapScalarQuantizedVectorValues scoreVectorValues = + new OffHeapScalarQuantizedVectorValues.DenseOffHeapVectorValues( + true, + fieldInfo.getVectorDimension(), + docsWithField.cardinality(), + vectorValues.getCentroid(), + vectorValues.getCentroidDP(), + quantizer, + fi.scalarEncoding, + fieldInfo.getVectorSimilarityFunction(), + vectorScorer, + quantizedScoreDataInput); + RandomVectorScorerSupplier scorerSupplier = + vectorScorer.getRandomVectorScorerSupplier( + fieldInfo.getVectorSimilarityFunction(), scoreVectorValues, vectorValues); + final String finalTempScoreQuantizedVectorName = tempScoreQuantizedVectorName; + return CloseableRandomVectorScorerSupplier.create( + scorerSupplier, + vectorValues.size(), + () -> { + IOUtils.close(quantizedScoreDataInput); + IOUtils.deleteFilesIgnoringExceptions( + segmentWriteState.directory, finalTempScoreQuantizedVectorName); + }); + } catch (Throwable t) { + IOUtils.closeWhileSuppressingExceptions(t, quantizedScoreDataInput); + throw t; + } + } + + static DocsWithFieldSet writeBinarizedQueryData( + QuantizedByteVectorValues quantizedByteVectorValues, + ScalarEncoding encoding, + IndexOutput binarizedQueryData, + FloatVectorValues floatVectorValues, + OptimizedScalarQuantizer binaryQuantizer) + throws IOException { + if (encoding.isAsymmetric() == false) { + throw new IllegalArgumentException("encoding and queryEncoding must be different"); + } + DocsWithFieldSet docsWithField = new DocsWithFieldSet(); + int discretizedDims = encoding.getDiscreteDimensions(floatVectorValues.dimension()); + byte[] quantizationScratch = new byte[discretizedDims]; + byte[] toQuery = new byte[encoding.getQueryPackedLength(discretizedDims)]; + KnnVectorValues.DocIndexIterator iterator = floatVectorValues.iterator(); + for (int docV = iterator.nextDoc(); docV != NO_MORE_DOCS; docV = iterator.nextDoc()) { + // write index vector + OptimizedScalarQuantizer.QuantizationResult r = + binaryQuantizer.scalarQuantize( + floatVectorValues.vectorValue(iterator.index()), + quantizationScratch, + encoding.getQueryBits(), + quantizedByteVectorValues.getCentroid()); + docsWithField.add(docV); + // pack and store the 4bit query vector + transposeHalfByte(quantizationScratch, toQuery); + binarizedQueryData.writeBytes(toQuery, toQuery.length); + binarizedQueryData.writeInt(Float.floatToIntBits(r.lowerInterval())); + binarizedQueryData.writeInt(Float.floatToIntBits(r.upperInterval())); + binarizedQueryData.writeInt(Float.floatToIntBits(r.additionalCorrection())); + binarizedQueryData.writeInt(r.quantizedComponentSum()); + } + return docsWithField; + } + + private record FieldEntry( + VectorSimilarityFunction similarityFunction, + VectorEncoding vectorEncoding, + int dimension, + long vectorDataOffset, + long vectorDataLength, + int size, + ScalarEncoding scalarEncoding, + float[] centroid, + float centroidDP, + OrdToDocDISIReaderConfiguration ordToDocDISIReaderConfiguration) { + + boolean isDataBlind() { + if (centroid == null) return false; + for (float v : centroid) { + if (v != 0f) return false; + } Review Comment: Q: what if the centroid happens to be zero, even if `enableCentering = true` was used? Would this cause the returned `FloatVectorValues` to be built from dequantized bytes, even when the full floats were available? Should `enableCentering` be persisted in the index, instead of inferring from the centroid? ########## lucene/core/src/java/org/apache/lucene/codecs/lucene106/Lucene106ScalarQuantizedVectorsFormat.java: ########## @@ -0,0 +1,170 @@ +/* + * 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.lucene106; + +import java.io.IOException; +import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding; + +/** + * The quantization format used here is a per-vector optimized scalar quantization. These ideas are + * evolutions of LVQ proposed in <a href="https://arxiv.org/abs/2304.04759";>Similarity search in the + * blink of an eye with compressed indices</a> by Cecilia Aguerrebere et al., the previous work on + * globally optimized scalar quantization in Apache Lucene, and <a + * href="https://arxiv.org/abs/1908.10396";>Accelerating Large-Scale Inference with Anisotropic + * Vector Quantization </a> by Ruiqi Guo et. al. Also see {@link + * org.apache.lucene.util.quantization.OptimizedScalarQuantizer}. Some of key features are: + * + * <ul> + * <li>Estimating the distance between two vectors using their centroid centered distance. This + * requires some additional corrective factors, but allows for centroid centering to occur. + * <li>Optimized scalar quantization to single bit level of centroid centered vectors. + * <li>Asymmetric quantization of vectors, where query vectors are quantized to half-byte (4 bits) + * precision (normalized to the centroid) and then compared directly against the single bit + * quantized vectors in the index. + * <li>Transforming the half-byte quantized query vectors in such a way that the comparison with + * single bit vectors can be done with bit arithmetic. + * </ul> + * + * A previous work related to improvements over regular LVQ is <a + * href="https://arxiv.org/abs/2409.09913";>Practical and Asymptotically Optimal Quantization of + * High-Dimensional Vectors in Euclidean Space for Approximate Nearest Neighbor Search </a> by + * Jianyang Gao, et. al. + * + * <p>The format is stored within two files: + * + * <h2>.veq (vector data) file</h2> + * + * <p>Stores the quantized vectors in a flat format. Additionally, it stores each vector's + * corrective factors. At the end of the file, additional information is stored for vector ordinal + * to centroid ordinal mapping and sparse vector information. + * + * <ul> + * <li>For each vector: + * <ul> + * <li><b>[byte]</b> the quantized values. Each dimension may be up to 8 bits, and multiple + * dimensions may be packed into a single byte. + * <li><b>[float]</b> the optimized quantiles and an additional similarity dependent + * corrective factor. + * <li><b>[int]</b> the sum of the quantized components + * </ul> + * <li>After the vectors, sparse vector information keeping track of monotonic blocks. + * </ul> + * + * <h2>.vemq (vector metadata) file</h2> + * + * <p>Stores the metadata for the vectors. This includes the number of vectors, the number of + * dimensions, and file offset information. + * + * <ul> + * <li><b>int</b> the field number + * <li><b>int</b> the vector encoding ordinal + * <li><b>int</b> the vector similarity ordinal + * <li><b>vint</b> the vector dimensions + * <li><b>vlong</b> the offset to the vector data in the .veq file + * <li><b>vlong</b> the length of the vector data in the .veq file + * <li><b>vint</b> the number of vectors + * <li><b>vint</b> the wire number for ScalarEncoding + * <li><b>[float]</b> the centroid + * <li><b>float</b> the centroid square magnitude + * <li>The sparse vector information, if required, mapping vector ordinal to doc ID Review Comment: This appears to be the [same](https://github.com/apache/lucene/blob/d8ffa2a68668f027f697bcbe37bf4d1bbcf014bf/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java#L79-L91) as the `Lucene104` codec? Is there no change to the on-disk representation? If so, could / should we add the data-blind option to the `Lucene104` class instead of creating a new format? ########## lucene/core/src/java/org/apache/lucene/codecs/lucene106/Lucene106ScalarQuantizedVectorsReader.java: ########## @@ -0,0 +1,738 @@ +/* + * 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.lucene106; + +import static org.apache.lucene.codecs.lucene106.Lucene106ScalarQuantizedVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.transposeHalfByte; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Stream; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.store.ChecksumIndexInput; +import org.apache.lucene.store.DataAccessHint; +import org.apache.lucene.store.FileDataHint; +import org.apache.lucene.store.FileTypeHint; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.quantization.OptimizedScalarQuantizer; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding; +import org.apache.lucene.util.quantization.QuantizedVectorsReader; +import org.apache.lucene.util.quantization.ScalarQuantizer; + +/** + * Reader for scalar quantized vectors in the Lucene 10.5 format. + * + * @lucene.experimental + */ +public class Lucene106ScalarQuantizedVectorsReader extends FlatVectorsReader + implements QuantizedVectorsReader { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Lucene106ScalarQuantizedVectorsReader.class); + + private final Map<String, FieldEntry> fields = new HashMap<>(); + private final IndexInput quantizedVectorData; + private final FlatVectorsReader rawVectorsReader; + private final Lucene106ScalarQuantizedVectorScorer vectorScorer; + public static final int EXHAUSTIVE_BULK_SCORE_ORDS = 64; + + public Lucene106ScalarQuantizedVectorsReader( + SegmentReadState state, + FlatVectorsReader rawVectorsReader, + Lucene106ScalarQuantizedVectorScorer vectorsScorer) + throws IOException { + // Quantized vectors are accessed randomly from their node ID stored in the HNSW + // graph. + this(state, rawVectorsReader, vectorsScorer, DataAccessHint.RANDOM); + } + + public Lucene106ScalarQuantizedVectorsReader( + SegmentReadState state, + FlatVectorsReader rawVectorsReader, + Lucene106ScalarQuantizedVectorScorer vectorsScorer, + DataAccessHint accessHint) + throws IOException { + this.vectorScorer = vectorsScorer; + this.rawVectorsReader = rawVectorsReader; + int versionMeta = -1; + String metaFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, + state.segmentSuffix, + Lucene106ScalarQuantizedVectorsFormat.META_EXTENSION); + try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName)) { + Throwable priorE = null; + try { + versionMeta = + CodecUtil.checkIndexHeader( + meta, + Lucene106ScalarQuantizedVectorsFormat.META_CODEC_NAME, + Lucene106ScalarQuantizedVectorsFormat.VERSION_START, + Lucene106ScalarQuantizedVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + readFields(meta, state.fieldInfos); + } catch (Throwable exception) { + priorE = exception; + } finally { + CodecUtil.checkFooter(meta, priorE); + } + + final IOContext.FileOpenHint[] hints = + Stream.of(FileTypeHint.DATA, FileDataHint.KNN_VECTORS, accessHint) + .filter(Objects::nonNull) + .toArray(IOContext.FileOpenHint[]::new); + quantizedVectorData = + openDataInput( + state, + versionMeta, + VECTOR_DATA_EXTENSION, + Lucene106ScalarQuantizedVectorsFormat.VECTOR_DATA_CODEC_NAME, + state.context.withHints(hints)); + } catch (Throwable t) { + IOUtils.closeWhileSuppressingExceptions(t, this); + throw t; + } + } + + private void readFields(ChecksumIndexInput meta, FieldInfos infos) throws IOException { + for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) { + FieldInfo info = infos.fieldInfo(fieldNumber); + if (info == null) { + throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta); + } + FieldEntry fieldEntry = readField(meta, info); + validateFieldEntry(info, fieldEntry); + fields.put(info.name, fieldEntry); + } + } + + static void validateFieldEntry(FieldInfo info, FieldEntry fieldEntry) { + int dimension = info.getVectorDimension(); + if (dimension != fieldEntry.dimension) { + throw new IllegalStateException( + "Inconsistent vector dimension for field=\"" + + info.name + + "\"; " + + dimension + + " != " + + fieldEntry.dimension); + } + + long numQuantizedVectorBytes = + Math.multiplyExact( + (fieldEntry.scalarEncoding.getDocPackedLength(dimension) + + (Float.BYTES * 3) + + Integer.BYTES), + (long) fieldEntry.size); + if (numQuantizedVectorBytes != fieldEntry.vectorDataLength) { + throw new IllegalStateException( + "vector data length " + + fieldEntry.vectorDataLength + + " not matching size = " + + fieldEntry.size + + " * (dims=" + + dimension + + " + 16" + + ") = " + + numQuantizedVectorBytes); + } + } + + @Override + public FlatVectorsScorer getFlatVectorScorer(String field) throws IOException { + return vectorScorer; + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException { + 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 + public RandomVectorScorer getRandomVectorScorer(String field, short[] target) throws IOException { + return rawVectorsReader.getRandomVectorScorer(field, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException { + return rawVectorsReader.getRandomVectorScorer(field, target); + } + + @Override + public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException { + rawVectorsReader.checkIntegrity(merge); + CodecUtil.checksumEntireFile(quantizedVectorData); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + FieldEntry fi = fields.get(field); + if (fi == null) { + return null; + } + if (fi.vectorEncoding != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException( + "field=\"" + + field + + "\" is encoded as: " + + fi.vectorEncoding + + " expected: " + + VectorEncoding.FLOAT32); + } + + FloatVectorValues rawFloatVectorValues = + fi.isDataBlind() ? null : rawVectorsReader.getFloatVectorValues(field); + + if (rawFloatVectorValues == null) { + // Data-blind mode: full-precision float vectors were never stored. Reconstruct floats from Review Comment: [DISCUSS] I think Lucene should disallow operations on data it dropped (i.e. throw an error when trying to retrieve floats / rescore using floats that were **not** stored in the index, instead of using a lossy value). Does Lucene have other use-cases that support indexing some data (and specific operations to search it), but not retain / support returning the original data? One example: if a user wrote some segments using `enableCentering = false`, then re-opens the index and attempts to merge with `enableCentering = true`, the dequantized vectors would be used instead of the original ones, leading to further loss of information? -- 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]
