From b886faaa8a3dedd1830497f740c56f0ff48d736a Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 25 Sep 2025 10:59:29 +0100 Subject: [PATCH 1/2] Add classes removed in Lucene 10.4 --- .../ES814ScalarQuantizedVectorsFormat.java | 6 +- .../Lucene99ScalarQuantizedVectorsReader.java | 455 +++++++ .../Lucene99ScalarQuantizedVectorsWriter.java | 1210 +++++++++++++++++ .../OffHeapQuantizedByteVectorValues.java | 403 ++++++ 4 files changed, 2071 insertions(+), 3 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.java create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsWriter.java create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapQuantizedByteVectorValues.java diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/ES814ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/ES814ScalarQuantizedVectorsFormat.java index 56710d49b5a7a..0b316ede8dbae 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/ES814ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/ES814ScalarQuantizedVectorsFormat.java @@ -17,8 +17,6 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; import org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer; import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; -import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; -import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FloatVectorValues; @@ -40,7 +38,6 @@ import java.io.IOException; import java.util.Map; -import static org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat.DYNAMIC_CONFIDENCE_INTERVAL; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; public class ES814ScalarQuantizedVectorsFormat extends FlatVectorsFormat { @@ -60,6 +57,9 @@ public class ES814ScalarQuantizedVectorsFormat extends FlatVectorsFormat { /** The maximum confidence interval */ private static final float MAXIMUM_CONFIDENCE_INTERVAL = 1f; + /** Dynamic confidence interval */ + public static final float DYNAMIC_CONFIDENCE_INTERVAL = 0f; + /** * Controls the confidence interval used to scalar quantize the vectors the default value is * calculated as `1-1/(vector_dimensions + 1)` diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.java new file mode 100644 index 0000000000000..9c0e855faf6ad --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.java @@ -0,0 +1,455 @@ +/* + * @notice + * 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. + * + * Modifications copyright (C) 2024 Elasticsearch B.V. + */ + +package org.elasticsearch.index.codec.vectors; + +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.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntObjectHashMap; +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.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.apache.lucene.util.quantization.QuantizedVectorsReader; +import org.apache.lucene.util.quantization.ScalarQuantizer; +import org.elasticsearch.core.SuppressForbidden; + +import java.io.IOException; +import java.util.Map; + +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding; + +/** + * Copied from Lucene 10.3. + */ +@SuppressForbidden(reason = "Lucene classes") +final class Lucene99ScalarQuantizedVectorsReader extends FlatVectorsReader implements QuantizedVectorsReader { + + private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Lucene99ScalarQuantizedVectorsReader.class); + + static final int VERSION_START = 0; + static final int VERSION_ADD_BITS = 1; + static final int VERSION_CURRENT = VERSION_ADD_BITS; + static final String META_CODEC_NAME = "Lucene99ScalarQuantizedVectorsFormatMeta"; + static final String VECTOR_DATA_CODEC_NAME = "Lucene99ScalarQuantizedVectorsFormatData"; + static final String META_EXTENSION = "vemq"; + static final String VECTOR_DATA_EXTENSION = "veq"; + + /** Dynamic confidence interval */ + public static final float DYNAMIC_CONFIDENCE_INTERVAL = 0f; + + private final IntObjectHashMap fields = new IntObjectHashMap<>(); + private final IndexInput quantizedVectorData; + private final FlatVectorsReader rawVectorsReader; + private final FieldInfos fieldInfos; + + Lucene99ScalarQuantizedVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader, FlatVectorsScorer scorer) + throws IOException { + super(scorer); + this.rawVectorsReader = rawVectorsReader; + this.fieldInfos = state.fieldInfos; + int versionMeta = -1; + String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + boolean success = false; + try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName)) { + Throwable priorE = null; + try { + versionMeta = CodecUtil.checkIndexHeader( + meta, + META_CODEC_NAME, + VERSION_START, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix + ); + readFields(meta, versionMeta, state.fieldInfos); + } catch (Throwable exception) { + priorE = exception; + } finally { + CodecUtil.checkFooter(meta, priorE); + } + quantizedVectorData = openDataInput( + state, + versionMeta, + VECTOR_DATA_EXTENSION, + VECTOR_DATA_CODEC_NAME, + // Quantized vectors are accessed randomly from their node ID stored in the HNSW + // graph. + state.context.withHints(FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM) + ); + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + private void readFields(ChecksumIndexInput meta, int versionMeta, 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, versionMeta, info); + validateFieldEntry(info, fieldEntry); + fields.put(info.number, 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 + ); + } + + final long quantizedVectorBytes; + if (fieldEntry.bits <= 4 && fieldEntry.compress) { + // two dimensions -> one byte + quantizedVectorBytes = ((dimension + 1) >> 1) + Float.BYTES; + } else { + // one dimension -> one byte + quantizedVectorBytes = dimension + Float.BYTES; + } + long numQuantizedVectorBytes = Math.multiplyExact(quantizedVectorBytes, fieldEntry.size); + if (numQuantizedVectorBytes != fieldEntry.vectorDataLength) { + throw new IllegalStateException( + "Quantized vector data length " + + fieldEntry.vectorDataLength + + " not matching size=" + + fieldEntry.size + + " * (dim=" + + dimension + + " + 4)" + + " = " + + numQuantizedVectorBytes + ); + } + } + + @Override + public void checkIntegrity() throws IOException { + rawVectorsReader.checkIntegrity(); + CodecUtil.checksumEntireFile(quantizedVectorData); + } + + private FieldEntry getFieldEntry(String field) { + final FieldInfo info = fieldInfos.fieldInfo(field); + final FieldEntry fieldEntry; + if (info == null || (fieldEntry = fields.get(info.number)) == null) { + throw new IllegalArgumentException("field=\"" + field + "\" not found"); + } + if (fieldEntry.vectorEncoding != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException( + "field=\"" + field + "\" is encoded as: " + fieldEntry.vectorEncoding + " expected: " + VectorEncoding.FLOAT32 + ); + } + return fieldEntry; + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + final FieldEntry fieldEntry = getFieldEntry(field); + final FloatVectorValues rawVectorValues = rawVectorsReader.getFloatVectorValues(field); + OffHeapQuantizedByteVectorValues quantizedByteVectorValues = OffHeapQuantizedByteVectorValues.load( + fieldEntry.ordToDoc, + fieldEntry.dimension, + fieldEntry.size, + fieldEntry.scalarQuantizer, + fieldEntry.similarityFunction, + vectorScorer, + fieldEntry.compress, + fieldEntry.vectorDataOffset, + fieldEntry.vectorDataLength, + quantizedVectorData + ); + return new QuantizedVectorValues(rawVectorValues, quantizedByteVectorValues); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return rawVectorsReader.getByteVectorValues(field); + } + + 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); + boolean success = false; + try { + int versionVectorData = CodecUtil.checkIndexHeader( + in, + codecName, + VERSION_START, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix + ); + if (versionMeta != versionVectorData) { + throw new CorruptIndexException( + "Format versions mismatch: meta=" + versionMeta + ", " + codecName + "=" + versionVectorData, + in + ); + } + CodecUtil.retrieveChecksum(in); + success = true; + return in; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(in); + } + } + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException { + final FieldEntry fieldEntry = getFieldEntry(field); + if (fieldEntry.scalarQuantizer == null) { + return rawVectorsReader.getRandomVectorScorer(field, target); + } + OffHeapQuantizedByteVectorValues vectorValues = OffHeapQuantizedByteVectorValues.load( + fieldEntry.ordToDoc, + fieldEntry.dimension, + fieldEntry.size, + fieldEntry.scalarQuantizer, + fieldEntry.similarityFunction, + vectorScorer, + fieldEntry.compress, + fieldEntry.vectorDataOffset, + fieldEntry.vectorDataLength, + quantizedVectorData + ); + return vectorScorer.getRandomVectorScorer(fieldEntry.similarityFunction, vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException { + return rawVectorsReader.getRandomVectorScorer(field, target); + } + + @Override + public void close() throws IOException { + IOUtils.close(quantizedVectorData, rawVectorsReader); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + fields.ramBytesUsed() + rawVectorsReader.ramBytesUsed(); + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + var raw = rawVectorsReader.getOffHeapByteSize(fieldInfo); + var fieldEntry = fields.get(fieldInfo.number); + if (fieldEntry == null) { + assert fieldInfo.getVectorEncoding() == VectorEncoding.BYTE; + return raw; + } + var quant = Map.of(VECTOR_DATA_EXTENSION, fieldEntry.vectorDataLength()); + return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, quant); + } + + private FieldEntry readField(IndexInput input, int versionMeta, FieldInfo info) throws IOException { + VectorEncoding vectorEncoding = readVectorEncoding(input); + 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, versionMeta, vectorEncoding, info.getVectorSimilarityFunction()); + } + + @Override + public QuantizedByteVectorValues getQuantizedVectorValues(String field) throws IOException { + final FieldEntry fieldEntry = getFieldEntry(field); + return OffHeapQuantizedByteVectorValues.load( + fieldEntry.ordToDoc, + fieldEntry.dimension, + fieldEntry.size, + fieldEntry.scalarQuantizer, + fieldEntry.similarityFunction, + vectorScorer, + fieldEntry.compress, + fieldEntry.vectorDataOffset, + fieldEntry.vectorDataLength, + quantizedVectorData + ); + } + + @Override + public ScalarQuantizer getQuantizationState(String field) { + final FieldEntry fieldEntry = getFieldEntry(field); + return fieldEntry.scalarQuantizer; + } + + private record FieldEntry( + VectorSimilarityFunction similarityFunction, + VectorEncoding vectorEncoding, + int dimension, + long vectorDataOffset, + long vectorDataLength, + ScalarQuantizer scalarQuantizer, + int size, + byte bits, + boolean compress, + OrdToDocDISIReaderConfiguration ordToDoc + ) { + + static FieldEntry create( + IndexInput input, + int versionMeta, + VectorEncoding vectorEncoding, + VectorSimilarityFunction similarityFunction + ) throws IOException { + final var vectorDataOffset = input.readVLong(); + final var vectorDataLength = input.readVLong(); + final var dimension = input.readVInt(); + final var size = input.readInt(); + final ScalarQuantizer scalarQuantizer; + final byte bits; + final boolean compress; + if (size > 0) { + if (versionMeta < VERSION_ADD_BITS) { + int floatBits = input.readInt(); // confidenceInterval, unused + if (floatBits == -1) { // indicates a null confidence interval + throw new CorruptIndexException("Missing confidence interval for scalar quantizer", input); + } + float confidenceInterval = Float.intBitsToFloat(floatBits); + // indicates a dynamic interval, which shouldn't be provided in this version + if (confidenceInterval == DYNAMIC_CONFIDENCE_INTERVAL) { + throw new CorruptIndexException("Invalid confidence interval for scalar quantizer: " + confidenceInterval, input); + } + bits = (byte) 7; + compress = false; + float minQuantile = Float.intBitsToFloat(input.readInt()); + float maxQuantile = Float.intBitsToFloat(input.readInt()); + scalarQuantizer = new ScalarQuantizer(minQuantile, maxQuantile, (byte) 7); + } else { + input.readInt(); // confidenceInterval, unused + bits = input.readByte(); + compress = input.readByte() == 1; + float minQuantile = Float.intBitsToFloat(input.readInt()); + float maxQuantile = Float.intBitsToFloat(input.readInt()); + scalarQuantizer = new ScalarQuantizer(minQuantile, maxQuantile, bits); + } + } else { + scalarQuantizer = null; + bits = (byte) 7; + compress = false; + } + final var ordToDoc = OrdToDocDISIReaderConfiguration.fromStoredMeta(input, size); + return new FieldEntry( + similarityFunction, + vectorEncoding, + dimension, + vectorDataOffset, + vectorDataLength, + scalarQuantizer, + size, + bits, + compress, + ordToDoc + ); + } + } + + private static final class QuantizedVectorValues extends FloatVectorValues { + private final FloatVectorValues rawVectorValues; + private final QuantizedByteVectorValues quantizedVectorValues; + + QuantizedVectorValues(FloatVectorValues 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 float[] vectorValue(int ord) throws IOException { + return rawVectorValues.vectorValue(ord); + } + + @Override + public int ordToDoc(int ord) { + return rawVectorValues.ordToDoc(ord); + } + + @Override + public QuantizedVectorValues copy() throws IOException { + return new QuantizedVectorValues(rawVectorValues.copy(), quantizedVectorValues.copy()); + } + + @Override + public VectorScorer scorer(float[] query) throws IOException { + return quantizedVectorValues.scorer(query); + } + + @Override + public VectorScorer rescorer(float[] query) throws IOException { + return rawVectorValues.rescorer(query); + } + + @Override + public DocIndexIterator iterator() { + return rawVectorValues.iterator(); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsWriter.java new file mode 100644 index 0000000000000..c230c81feaa93 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsWriter.java @@ -0,0 +1,1210 @@ +/* + * @notice + * 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. + * + * Modifications copyright (C) 2024 Elasticsearch B.V. + */ + +package org.elasticsearch.index.codec.vectors; + +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +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.IntArrayList; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.VectorUtil; +import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.apache.lucene.util.quantization.QuantizedVectorsReader; +import org.apache.lucene.util.quantization.ScalarQuantizer; +import org.elasticsearch.core.SuppressForbidden; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; + +import static org.apache.lucene.codecs.KnnVectorsWriter.MergedVectorValues.hasVectorValues; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +/** + * Copied from Lucene 10.3. + */ +@SuppressForbidden(reason = "Lucene classes") +public final class Lucene99ScalarQuantizedVectorsWriter extends FlatVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(Lucene99ScalarQuantizedVectorsWriter.class); + + static final String QUANTIZED_VECTOR_COMPONENT = "QVEC"; + static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; + + static final int VERSION_START = 0; + static final int VERSION_ADD_BITS = 1; + static final String META_CODEC_NAME = "Lucene99ScalarQuantizedVectorsFormatMeta"; + static final String VECTOR_DATA_CODEC_NAME = "Lucene99ScalarQuantizedVectorsFormatData"; + static final String META_EXTENSION = "vemq"; + static final String VECTOR_DATA_EXTENSION = "veq"; + + private static final float MINIMUM_CONFIDENCE_INTERVAL = 0.9f; + + /** Dynamic confidence interval */ + public static final float DYNAMIC_CONFIDENCE_INTERVAL = 0f; + + static float calculateDefaultConfidenceInterval(int vectorDimension) { + return Math.max(MINIMUM_CONFIDENCE_INTERVAL, 1f - (1f / (vectorDimension + 1))); + } + + // Used for determining when merged quantiles shifted too far from individual segment quantiles. + // When merging quantiles from various segments, we need to ensure that the new quantiles + // are not exceptionally different from an individual segments quantiles. + // This would imply that the quantization buckets would shift too much + // for floating point values and justify recalculating the quantiles. This helps preserve + // accuracy of the calculated quantiles, even in adversarial cases such as vector clustering. + // This number was determined via empirical testing + private static final float QUANTILE_RECOMPUTE_LIMIT = 32; + // Used for determining if a new quantization state requires a re-quantization + // for a given segment. + // This ensures that in expectation 4/5 of the vector would be unchanged by requantization. + // Furthermore, only those values where the value is within 1/5 of the centre of a quantization + // bin will be changed. In these cases the error introduced by snapping one way or another + // is small compared to the error introduced by quantization in the first place. Furthermore, + // empirical testing showed that the relative error by not requantizing is small (compared to + // the quantization error) and the condition is sensitive enough to detect all adversarial cases, + // such as merging clustered data. + private static final float REQUANTIZATION_LIMIT = 0.2f; + private final SegmentWriteState segmentWriteState; + + private final List fields = new ArrayList<>(); + private final IndexOutput meta, quantizedVectorData; + private final Float confidenceInterval; + private final FlatVectorsWriter rawVectorDelegate; + private final byte bits; + private final boolean compress; + private final int version; + private boolean finished; + + public Lucene99ScalarQuantizedVectorsWriter( + SegmentWriteState state, + Float confidenceInterval, + FlatVectorsWriter rawVectorDelegate, + FlatVectorsScorer scorer + ) throws IOException { + this(state, VERSION_START, confidenceInterval, (byte) 7, false, rawVectorDelegate, scorer); + if (confidenceInterval != null && confidenceInterval == 0) { + throw new IllegalArgumentException("confidenceInterval cannot be set to zero"); + } + } + + public Lucene99ScalarQuantizedVectorsWriter( + SegmentWriteState state, + Float confidenceInterval, + byte bits, + boolean compress, + FlatVectorsWriter rawVectorDelegate, + FlatVectorsScorer scorer + ) throws IOException { + this(state, VERSION_ADD_BITS, confidenceInterval, bits, compress, rawVectorDelegate, scorer); + } + + private Lucene99ScalarQuantizedVectorsWriter( + SegmentWriteState state, + int version, + Float confidenceInterval, + byte bits, + boolean compress, + FlatVectorsWriter rawVectorDelegate, + FlatVectorsScorer scorer + ) throws IOException { + super(scorer); + this.confidenceInterval = confidenceInterval; + this.bits = bits; + this.compress = compress; + this.version = version; + segmentWriteState = state; + String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + + String quantizedVectorDataFileName = IndexFileNames.segmentFileName( + state.segmentInfo.name, + state.segmentSuffix, + VECTOR_DATA_EXTENSION + ); + this.rawVectorDelegate = rawVectorDelegate; + boolean success = false; + try { + meta = state.directory.createOutput(metaFileName, state.context); + quantizedVectorData = state.directory.createOutput(quantizedVectorDataFileName, state.context); + + CodecUtil.writeIndexHeader(meta, META_CODEC_NAME, version, state.segmentInfo.getId(), state.segmentSuffix); + CodecUtil.writeIndexHeader( + quantizedVectorData, + VECTOR_DATA_CODEC_NAME, + version, + state.segmentInfo.getId(), + state.segmentSuffix + ); + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + @Override + public FlatFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + FlatFieldVectorsWriter rawVectorDelegate = this.rawVectorDelegate.addField(fieldInfo); + if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) { + if (bits <= 4 && fieldInfo.getVectorDimension() % 2 != 0) { + throw new IllegalArgumentException( + "bits=" + bits + " is not supported for odd vector dimensions; vector dimension=" + fieldInfo.getVectorDimension() + ); + } + @SuppressWarnings("unchecked") + FieldWriter quantizedWriter = new FieldWriter( + confidenceInterval, + bits, + compress, + fieldInfo, + segmentWriteState.infoStream, + (FlatFieldVectorsWriter) rawVectorDelegate + ); + fields.add(quantizedWriter); + return quantizedWriter; + } + return rawVectorDelegate; + } + + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + rawVectorDelegate.mergeOneField(fieldInfo, mergeState); + // Since we know we will not be searching for additional indexing, we can just write the + // the vectors directly to the new segment. + // No need to use temporary file as we don't have to re-open for reading + if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) { + ScalarQuantizer mergedQuantizationState = mergeAndRecalculateQuantiles(mergeState, fieldInfo, confidenceInterval, bits); + MergedQuantizedVectorValues byteVectorValues = MergedQuantizedVectorValues.mergeQuantizedByteVectorValues( + fieldInfo, + mergeState, + mergedQuantizationState + ); + long vectorDataOffset = quantizedVectorData.alignFilePointer(Float.BYTES); + DocsWithFieldSet docsWithField = writeQuantizedVectorData(quantizedVectorData, byteVectorValues, bits, compress); + long vectorDataLength = quantizedVectorData.getFilePointer() - vectorDataOffset; + writeMeta( + fieldInfo, + segmentWriteState.segmentInfo.maxDoc(), + vectorDataOffset, + vectorDataLength, + confidenceInterval, + bits, + compress, + mergedQuantizationState.getLowerQuantile(), + mergedQuantizationState.getUpperQuantile(), + docsWithField + ); + } + } + + @Override + public CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) { + // Simply merge the underlying delegate, which just copies the raw vector data to a new + // segment file + rawVectorDelegate.mergeOneField(fieldInfo, mergeState); + ScalarQuantizer mergedQuantizationState = mergeAndRecalculateQuantiles(mergeState, fieldInfo, confidenceInterval, bits); + return mergeOneFieldToIndex(segmentWriteState, fieldInfo, mergeState, mergedQuantizationState); + } + // We only merge the delegate, since the field type isn't float32, quantization wasn't + // supported, so bypass it. + return rawVectorDelegate.mergeOneFieldToIndex(fieldInfo, mergeState); + } + + @Override + public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { + rawVectorDelegate.flush(maxDoc, sortMap); + for (FieldWriter field : fields) { + ScalarQuantizer quantizer = field.createQuantizer(); + if (sortMap == null) { + writeField(field, maxDoc, quantizer); + } else { + writeSortingField(field, maxDoc, sortMap, quantizer); + } + field.finish(); + } + } + + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + rawVectorDelegate.finish(); + if (meta != null) { + // write end of fields marker + meta.writeInt(-1); + CodecUtil.writeFooter(meta); + } + if (quantizedVectorData != null) { + CodecUtil.writeFooter(quantizedVectorData); + } + } + + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (FieldWriter field : fields) { + // the field tracks the delegate field usage + total += field.ramBytesUsed(); + } + return total; + } + + private void writeField(FieldWriter fieldData, int maxDoc, ScalarQuantizer scalarQuantizer) throws IOException { + // write vector values + long vectorDataOffset = quantizedVectorData.alignFilePointer(Float.BYTES); + writeQuantizedVectors(fieldData, scalarQuantizer); + long vectorDataLength = quantizedVectorData.getFilePointer() - vectorDataOffset; + + writeMeta( + fieldData.fieldInfo, + maxDoc, + vectorDataOffset, + vectorDataLength, + confidenceInterval, + bits, + compress, + scalarQuantizer.getLowerQuantile(), + scalarQuantizer.getUpperQuantile(), + fieldData.getDocsWithFieldSet() + ); + } + + private void writeMeta( + FieldInfo field, + int maxDoc, + long vectorDataOffset, + long vectorDataLength, + Float confidenceInterval, + byte bits, + boolean compress, + Float lowerQuantile, + Float upperQuantile, + DocsWithFieldSet docsWithField + ) throws IOException { + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(field.getVectorSimilarityFunction().ordinal()); + meta.writeVLong(vectorDataOffset); + meta.writeVLong(vectorDataLength); + meta.writeVInt(field.getVectorDimension()); + int count = docsWithField.cardinality(); + meta.writeInt(count); + if (count > 0) { + assert Float.isFinite(lowerQuantile) && Float.isFinite(upperQuantile); + if (version >= VERSION_ADD_BITS) { + meta.writeInt(confidenceInterval == null ? -1 : Float.floatToIntBits(confidenceInterval)); + meta.writeByte(bits); + meta.writeByte(compress ? (byte) 1 : (byte) 0); + } else { + assert confidenceInterval == null || confidenceInterval != DYNAMIC_CONFIDENCE_INTERVAL; + meta.writeInt( + Float.floatToIntBits( + confidenceInterval == null ? calculateDefaultConfidenceInterval(field.getVectorDimension()) : confidenceInterval + ) + ); + } + meta.writeInt(Float.floatToIntBits(lowerQuantile)); + meta.writeInt(Float.floatToIntBits(upperQuantile)); + } + // write docIDs + OrdToDocDISIReaderConfiguration.writeStoredMeta( + DIRECT_MONOTONIC_BLOCK_SHIFT, + meta, + quantizedVectorData, + count, + maxDoc, + docsWithField + ); + } + + private void writeQuantizedVectors(FieldWriter fieldData, ScalarQuantizer scalarQuantizer) throws IOException { + byte[] vector = new byte[fieldData.fieldInfo.getVectorDimension()]; + byte[] compressedVector = fieldData.compress + ? OffHeapQuantizedByteVectorValues.compressedArray(fieldData.fieldInfo.getVectorDimension(), bits) + : null; + final ByteBuffer offsetBuffer = ByteBuffer.allocate(Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + float[] copy = fieldData.normalize ? new float[fieldData.fieldInfo.getVectorDimension()] : null; + assert fieldData.getVectors().isEmpty() || scalarQuantizer != null; + for (float[] v : fieldData.getVectors()) { + if (fieldData.normalize) { + System.arraycopy(v, 0, copy, 0, copy.length); + VectorUtil.l2normalize(copy); + v = copy; + } + + float offsetCorrection = scalarQuantizer.quantize(v, vector, fieldData.fieldInfo.getVectorSimilarityFunction()); + if (compressedVector != null) { + OffHeapQuantizedByteVectorValues.compressBytes(vector, compressedVector); + quantizedVectorData.writeBytes(compressedVector, compressedVector.length); + } else { + quantizedVectorData.writeBytes(vector, vector.length); + } + offsetBuffer.putFloat(offsetCorrection); + quantizedVectorData.writeBytes(offsetBuffer.array(), offsetBuffer.array().length); + offsetBuffer.rewind(); + } + } + + private void writeSortingField(FieldWriter fieldData, int maxDoc, Sorter.DocMap sortMap, ScalarQuantizer scalarQuantizer) + throws IOException { + final int[] ordMap = new int[fieldData.getDocsWithFieldSet().cardinality()]; // new ord to old ord + + DocsWithFieldSet newDocsWithField = new DocsWithFieldSet(); + mapOldOrdToNewOrd(fieldData.getDocsWithFieldSet(), sortMap, null, ordMap, newDocsWithField); + + // write vector values + long vectorDataOffset = quantizedVectorData.alignFilePointer(Float.BYTES); + writeSortedQuantizedVectors(fieldData, ordMap, scalarQuantizer); + long quantizedVectorLength = quantizedVectorData.getFilePointer() - vectorDataOffset; + writeMeta( + fieldData.fieldInfo, + maxDoc, + vectorDataOffset, + quantizedVectorLength, + confidenceInterval, + bits, + compress, + scalarQuantizer.getLowerQuantile(), + scalarQuantizer.getUpperQuantile(), + newDocsWithField + ); + } + + private void writeSortedQuantizedVectors(FieldWriter fieldData, int[] ordMap, ScalarQuantizer scalarQuantizer) throws IOException { + byte[] vector = new byte[fieldData.fieldInfo.getVectorDimension()]; + byte[] compressedVector = fieldData.compress + ? OffHeapQuantizedByteVectorValues.compressedArray(fieldData.fieldInfo.getVectorDimension(), bits) + : null; + final ByteBuffer offsetBuffer = ByteBuffer.allocate(Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + float[] copy = fieldData.normalize ? new float[fieldData.fieldInfo.getVectorDimension()] : null; + for (int ordinal : ordMap) { + float[] v = fieldData.getVectors().get(ordinal); + if (fieldData.normalize) { + System.arraycopy(v, 0, copy, 0, copy.length); + VectorUtil.l2normalize(copy); + v = copy; + } + float offsetCorrection = scalarQuantizer.quantize(v, vector, fieldData.fieldInfo.getVectorSimilarityFunction()); + if (compressedVector != null) { + OffHeapQuantizedByteVectorValues.compressBytes(vector, compressedVector); + quantizedVectorData.writeBytes(compressedVector, compressedVector.length); + } else { + quantizedVectorData.writeBytes(vector, vector.length); + } + offsetBuffer.putFloat(offsetCorrection); + quantizedVectorData.writeBytes(offsetBuffer.array(), offsetBuffer.array().length); + offsetBuffer.rewind(); + } + } + + private ScalarQuantizedCloseableRandomVectorScorerSupplier mergeOneFieldToIndex( + SegmentWriteState segmentWriteState, + FieldInfo fieldInfo, + MergeState mergeState, + ScalarQuantizer mergedQuantizationState + ) throws IOException { + if (segmentWriteState.infoStream.isEnabled(QUANTIZED_VECTOR_COMPONENT)) { + segmentWriteState.infoStream.message( + QUANTIZED_VECTOR_COMPONENT, + "quantized field=" + + " confidenceInterval=" + + confidenceInterval + + " minQuantile=" + + mergedQuantizationState.getLowerQuantile() + + " maxQuantile=" + + mergedQuantizationState.getUpperQuantile() + ); + } + long vectorDataOffset = quantizedVectorData.alignFilePointer(Float.BYTES); + IndexOutput tempQuantizedVectorData = segmentWriteState.directory.createTempOutput( + quantizedVectorData.getName(), + "temp", + segmentWriteState.context + ); + IndexInput quantizationDataInput = null; + boolean success = false; + try { + MergedQuantizedVectorValues byteVectorValues = MergedQuantizedVectorValues.mergeQuantizedByteVectorValues( + fieldInfo, + mergeState, + mergedQuantizationState + ); + DocsWithFieldSet docsWithField = writeQuantizedVectorData(tempQuantizedVectorData, byteVectorValues, bits, compress); + CodecUtil.writeFooter(tempQuantizedVectorData); + IOUtils.close(tempQuantizedVectorData); + quantizationDataInput = segmentWriteState.directory.openInput(tempQuantizedVectorData.getName(), segmentWriteState.context); + quantizedVectorData.copyBytes(quantizationDataInput, quantizationDataInput.length() - CodecUtil.footerLength()); + long vectorDataLength = quantizedVectorData.getFilePointer() - vectorDataOffset; + CodecUtil.retrieveChecksum(quantizationDataInput); + writeMeta( + fieldInfo, + segmentWriteState.segmentInfo.maxDoc(), + vectorDataOffset, + vectorDataLength, + confidenceInterval, + bits, + compress, + mergedQuantizationState.getLowerQuantile(), + mergedQuantizationState.getUpperQuantile(), + docsWithField + ); + success = true; + final IndexInput finalQuantizationDataInput = quantizationDataInput; + return new ScalarQuantizedCloseableRandomVectorScorerSupplier(() -> { + IOUtils.close(finalQuantizationDataInput); + segmentWriteState.directory.deleteFile(tempQuantizedVectorData.getName()); + }, + docsWithField.cardinality(), + vectorsScorer.getRandomVectorScorerSupplier( + fieldInfo.getVectorSimilarityFunction(), + new OffHeapQuantizedByteVectorValues.DenseOffHeapVectorValues( + fieldInfo.getVectorDimension(), + docsWithField.cardinality(), + mergedQuantizationState, + compress, + fieldInfo.getVectorSimilarityFunction(), + vectorsScorer, + quantizationDataInput + ) + ) + ); + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(tempQuantizedVectorData, quantizationDataInput); + IOUtils.deleteFilesIgnoringExceptions(segmentWriteState.directory, tempQuantizedVectorData.getName()); + } + } + } + + static ScalarQuantizer mergeQuantiles(List quantizationStates, IntArrayList segmentSizes, byte bits) { + assert quantizationStates.size() == segmentSizes.size(); + if (quantizationStates.isEmpty()) { + return null; + } + float lowerQuantile = 0f; + float upperQuantile = 0f; + int totalCount = 0; + for (int i = 0; i < quantizationStates.size(); i++) { + if (quantizationStates.get(i) == null) { + return null; + } + lowerQuantile += quantizationStates.get(i).getLowerQuantile() * segmentSizes.get(i); + upperQuantile += quantizationStates.get(i).getUpperQuantile() * segmentSizes.get(i); + totalCount += segmentSizes.get(i); + if (quantizationStates.get(i).getBits() != bits) { + return null; + } + } + lowerQuantile /= totalCount; + upperQuantile /= totalCount; + return new ScalarQuantizer(lowerQuantile, upperQuantile, bits); + } + + /** + * Returns true if the quantiles of the merged state are too far from the quantiles of the + * individual states. + * + * @param mergedQuantizationState The merged quantization state + * @param quantizationStates The quantization states of the individual segments + * @return true if the quantiles should be recomputed + */ + static boolean shouldRecomputeQuantiles(ScalarQuantizer mergedQuantizationState, List quantizationStates) { + // calculate the limit for the quantiles to be considered too far apart + // We utilize upper & lower here to determine if the new upper and merged upper would + // drastically + // change the quantization buckets for floats + // This is a fairly conservative check. + float limit = (mergedQuantizationState.getUpperQuantile() - mergedQuantizationState.getLowerQuantile()) / QUANTILE_RECOMPUTE_LIMIT; + for (ScalarQuantizer quantizationState : quantizationStates) { + if (Math.abs(quantizationState.getUpperQuantile() - mergedQuantizationState.getUpperQuantile()) > limit) { + return true; + } + if (Math.abs(quantizationState.getLowerQuantile() - mergedQuantizationState.getLowerQuantile()) > limit) { + return true; + } + } + return false; + } + + private static QuantizedVectorsReader getQuantizedKnnVectorsReader(KnnVectorsReader vectorsReader, String fieldName) { + if (vectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader candidateReader) { + vectorsReader = candidateReader.getFieldReader(fieldName); + } + if (vectorsReader instanceof QuantizedVectorsReader reader) { + return reader; + } + return null; + } + + private static ScalarQuantizer getQuantizedState(KnnVectorsReader vectorsReader, String fieldName) { + QuantizedVectorsReader reader = getQuantizedKnnVectorsReader(vectorsReader, fieldName); + if (reader != null) { + return reader.getQuantizationState(fieldName); + } + return null; + } + + /** + * Merges the quantiles of the segments and recalculates the quantiles if necessary. + * + * @param mergeState The merge state + * @param fieldInfo The field info + * @param confidenceInterval The confidence interval + * @param bits The number of bits + * @return The merged quantiles + * @throws IOException If there is a low-level I/O error + */ + public static ScalarQuantizer mergeAndRecalculateQuantiles( + MergeState mergeState, + FieldInfo fieldInfo, + Float confidenceInterval, + byte bits + ) throws IOException { + assert fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32); + List quantizationStates = new ArrayList<>(mergeState.liveDocs.length); + IntArrayList segmentSizes = new IntArrayList(mergeState.liveDocs.length); + for (int i = 0; i < mergeState.liveDocs.length; i++) { + FloatVectorValues fvv; + if (hasVectorValues(mergeState.fieldInfos[i], fieldInfo.name) + && (fvv = mergeState.knnVectorsReaders[i].getFloatVectorValues(fieldInfo.name)) != null + && fvv.size() > 0) { + ScalarQuantizer quantizationState = getQuantizedState(mergeState.knnVectorsReaders[i], fieldInfo.name); + // If we have quantization state, we can utilize that to make merging cheaper + quantizationStates.add(quantizationState); + segmentSizes.add(fvv.size()); + } + } + ScalarQuantizer mergedQuantiles = mergeQuantiles(quantizationStates, segmentSizes, bits); + // Segments no providing quantization state indicates that their quantiles were never + // calculated. + // To be safe, we should always recalculate given a sample set over all the float vectors in the + // merged + // segment view + if (mergedQuantiles == null + // For smaller `bits` values, we should always recalculate the quantiles + // TODO: this is very conservative, could we reuse information for even int4 quantization? + || bits <= 4 + || shouldRecomputeQuantiles(mergedQuantiles, quantizationStates)) { + int numVectors = 0; + DocIdSetIterator iter = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState).iterator(); + // iterate vectorValues and increment numVectors + for (int doc = iter.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iter.nextDoc()) { + numVectors++; + } + return buildScalarQuantizer( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState), + numVectors, + fieldInfo.getVectorSimilarityFunction(), + confidenceInterval, + bits + ); + } + return mergedQuantiles; + } + + static ScalarQuantizer buildScalarQuantizer( + FloatVectorValues floatVectorValues, + int numVectors, + VectorSimilarityFunction vectorSimilarityFunction, + Float confidenceInterval, + byte bits + ) throws IOException { + if (vectorSimilarityFunction == VectorSimilarityFunction.COSINE) { + floatVectorValues = new NormalizedFloatVectorValues(floatVectorValues); + vectorSimilarityFunction = VectorSimilarityFunction.DOT_PRODUCT; + } + if (confidenceInterval != null && confidenceInterval == DYNAMIC_CONFIDENCE_INTERVAL) { + return ScalarQuantizer.fromVectorsAutoInterval(floatVectorValues, vectorSimilarityFunction, numVectors, bits); + } + return ScalarQuantizer.fromVectors( + floatVectorValues, + confidenceInterval == null ? calculateDefaultConfidenceInterval(floatVectorValues.dimension()) : confidenceInterval, + numVectors, + bits + ); + } + + /** + * Returns true if the quantiles of the new quantization state are too far from the quantiles of + * the existing quantization state. This would imply that floating point values would slightly + * shift quantization buckets. + * + * @param existingQuantiles The existing quantiles for a segment + * @param newQuantiles The new quantiles for a segment, could be merged, or fully re-calculated + * @return true if the floating point values should be requantized + */ + static boolean shouldRequantize(ScalarQuantizer existingQuantiles, ScalarQuantizer newQuantiles) { + float tol = REQUANTIZATION_LIMIT * (newQuantiles.getUpperQuantile() - newQuantiles.getLowerQuantile()) / 128f; + if (Math.abs(existingQuantiles.getUpperQuantile() - newQuantiles.getUpperQuantile()) > tol) { + return true; + } + return Math.abs(existingQuantiles.getLowerQuantile() - newQuantiles.getLowerQuantile()) > tol; + } + + /** + * Writes the vector values to the output and returns a set of documents that contains vectors. + */ + public static DocsWithFieldSet writeQuantizedVectorData( + IndexOutput output, + QuantizedByteVectorValues quantizedByteVectorValues, + byte bits, + boolean compress + ) throws IOException { + DocsWithFieldSet docsWithField = new DocsWithFieldSet(); + final byte[] compressedVector = compress + ? OffHeapQuantizedByteVectorValues.compressedArray(quantizedByteVectorValues.dimension(), bits) + : null; + KnnVectorValues.DocIndexIterator iter = quantizedByteVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + // write vector + byte[] binaryValue = quantizedByteVectorValues.vectorValue(iter.index()); + assert binaryValue.length == quantizedByteVectorValues.dimension() + : "dim=" + quantizedByteVectorValues.dimension() + " len=" + binaryValue.length; + if (compressedVector != null) { + OffHeapQuantizedByteVectorValues.compressBytes(binaryValue, compressedVector); + output.writeBytes(compressedVector, compressedVector.length); + } else { + output.writeBytes(binaryValue, binaryValue.length); + } + output.writeInt(Float.floatToIntBits(quantizedByteVectorValues.getScoreCorrectionConstant(iter.index()))); + docsWithField.add(docV); + } + return docsWithField; + } + + @Override + public void close() throws IOException { + IOUtils.close(meta, quantizedVectorData, rawVectorDelegate); + } + + static class FieldWriter extends FlatFieldVectorsWriter { + private static final long SHALLOW_SIZE = shallowSizeOfInstance(FieldWriter.class); + private final FieldInfo fieldInfo; + private final Float confidenceInterval; + private final byte bits; + private final boolean compress; + private final InfoStream infoStream; + private final boolean normalize; + private boolean finished; + private final FlatFieldVectorsWriter flatFieldVectorsWriter; + + FieldWriter( + Float confidenceInterval, + byte bits, + boolean compress, + FieldInfo fieldInfo, + InfoStream infoStream, + FlatFieldVectorsWriter indexWriter + ) { + super(); + this.confidenceInterval = confidenceInterval; + this.bits = bits; + this.fieldInfo = fieldInfo; + this.normalize = fieldInfo.getVectorSimilarityFunction() == VectorSimilarityFunction.COSINE; + this.infoStream = infoStream; + this.compress = compress; + this.flatFieldVectorsWriter = Objects.requireNonNull(indexWriter); + } + + @Override + public boolean isFinished() { + return finished && flatFieldVectorsWriter.isFinished(); + } + + @Override + public void finish() throws IOException { + if (finished) { + return; + } + assert flatFieldVectorsWriter.isFinished(); + finished = true; + } + + ScalarQuantizer createQuantizer() throws IOException { + assert flatFieldVectorsWriter.isFinished(); + List floatVectors = flatFieldVectorsWriter.getVectors(); + if (floatVectors.size() == 0) { + return new ScalarQuantizer(0, 0, bits); + } + ScalarQuantizer quantizer = buildScalarQuantizer( + new FloatVectorWrapper(floatVectors), + floatVectors.size(), + fieldInfo.getVectorSimilarityFunction(), + confidenceInterval, + bits + ); + if (infoStream.isEnabled(QUANTIZED_VECTOR_COMPONENT)) { + infoStream.message( + QUANTIZED_VECTOR_COMPONENT, + "quantized field=" + + " confidenceInterval=" + + confidenceInterval + + " bits=" + + bits + + " minQuantile=" + + quantizer.getLowerQuantile() + + " maxQuantile=" + + quantizer.getUpperQuantile() + ); + } + return quantizer; + } + + @Override + public long ramBytesUsed() { + long size = SHALLOW_SIZE; + size += flatFieldVectorsWriter.ramBytesUsed(); + return size; + } + + @Override + public void addValue(int docID, float[] vectorValue) throws IOException { + flatFieldVectorsWriter.addValue(docID, vectorValue); + } + + @Override + public float[] copyValue(float[] vectorValue) { + throw new UnsupportedOperationException(); + } + + @Override + public List getVectors() { + return flatFieldVectorsWriter.getVectors(); + } + + @Override + public DocsWithFieldSet getDocsWithFieldSet() { + return flatFieldVectorsWriter.getDocsWithFieldSet(); + } + } + + static class FloatVectorWrapper extends FloatVectorValues { + private final List vectorList; + + FloatVectorWrapper(List vectorList) { + this.vectorList = vectorList; + } + + @Override + public int dimension() { + return vectorList.get(0).length; + } + + @Override + public int size() { + return vectorList.size(); + } + + @Override + public FloatVectorValues copy() throws IOException { + return this; + } + + @Override + public float[] vectorValue(int ord) throws IOException { + if (ord < 0 || ord >= vectorList.size()) { + throw new IOException("vector ord " + ord + " out of bounds"); + } + return vectorList.get(ord); + } + + @Override + public DocIndexIterator iterator() { + return createDenseIterator(); + } + } + + static class QuantizedByteVectorValueSub extends DocIDMerger.Sub { + private final QuantizedByteVectorValues values; + private final KnnVectorValues.DocIndexIterator iterator; + + QuantizedByteVectorValueSub(MergeState.DocMap docMap, QuantizedByteVectorValues values) { + super(docMap); + this.values = values; + iterator = values.iterator(); + assert iterator.docID() == -1; + } + + @Override + public int nextDoc() throws IOException { + return iterator.nextDoc(); + } + + public int index() { + return iterator.index(); + } + } + + /** Returns a merged view over all the segment's {@link QuantizedByteVectorValues}. */ + static class MergedQuantizedVectorValues extends QuantizedByteVectorValues { + public static MergedQuantizedVectorValues mergeQuantizedByteVectorValues( + FieldInfo fieldInfo, + MergeState mergeState, + ScalarQuantizer scalarQuantizer + ) throws IOException { + assert fieldInfo != null && fieldInfo.hasVectorValues(); + + List subs = new ArrayList<>(); + for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { + if (hasVectorValues(mergeState.fieldInfos[i], fieldInfo.name)) { + QuantizedVectorsReader reader = getQuantizedKnnVectorsReader(mergeState.knnVectorsReaders[i], fieldInfo.name); + assert scalarQuantizer != null; + final QuantizedByteVectorValueSub sub; + // Either our quantization parameters are way different than the merged ones + // Or we have never been quantized. + if (reader == null || reader.getQuantizationState(fieldInfo.name) == null + // For smaller `bits` values, we should always recalculate the quantiles + // TODO: this is very conservative, could we reuse information for even int4 + // quantization? + || scalarQuantizer.getBits() <= 4 + || shouldRequantize(reader.getQuantizationState(fieldInfo.name), scalarQuantizer)) { + FloatVectorValues toQuantize = mergeState.knnVectorsReaders[i].getFloatVectorValues(fieldInfo.name); + if (fieldInfo.getVectorSimilarityFunction() == VectorSimilarityFunction.COSINE) { + toQuantize = new NormalizedFloatVectorValues(toQuantize); + } + sub = new QuantizedByteVectorValueSub( + mergeState.docMaps[i], + new QuantizedFloatVectorValues(toQuantize, fieldInfo.getVectorSimilarityFunction(), scalarQuantizer) + ); + } else { + sub = new QuantizedByteVectorValueSub( + mergeState.docMaps[i], + new OffsetCorrectedQuantizedByteVectorValues( + reader.getQuantizedVectorValues(fieldInfo.name), + fieldInfo.getVectorSimilarityFunction(), + scalarQuantizer, + reader.getQuantizationState(fieldInfo.name) + ) + ); + } + subs.add(sub); + } + } + return new MergedQuantizedVectorValues(subs, mergeState); + } + + private final List subs; + private final DocIDMerger docIdMerger; + private final int size; + + private QuantizedByteVectorValueSub current; + + private MergedQuantizedVectorValues(List subs, MergeState mergeState) throws IOException { + this.subs = subs; + docIdMerger = DocIDMerger.of(subs, mergeState.needsIndexSort); + int totalSize = 0; + for (QuantizedByteVectorValueSub sub : subs) { + totalSize += sub.values.size(); + } + size = totalSize; + } + + @Override + public byte[] vectorValue(int ord) throws IOException { + return current.values.vectorValue(current.index()); + } + + @Override + public DocIndexIterator iterator() { + return new CompositeIterator(); + } + + @Override + public int size() { + return size; + } + + @Override + public int dimension() { + return subs.get(0).values.dimension(); + } + + @Override + public float getScoreCorrectionConstant(int ord) throws IOException { + return current.values.getScoreCorrectionConstant(current.index()); + } + + private class CompositeIterator extends DocIndexIterator { + private int docId; + private int ord; + + CompositeIterator() { + docId = -1; + ord = -1; + } + + @Override + public int index() { + return ord; + } + + @Override + public int docID() { + return docId; + } + + @Override + public int nextDoc() throws IOException { + current = docIdMerger.next(); + if (current == null) { + docId = NO_MORE_DOCS; + ord = NO_MORE_DOCS; + } else { + docId = current.mappedDocID; + ++ord; + } + return docId; + } + + @Override + public int advance(int target) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public long cost() { + return size; + } + } + } + + static class QuantizedFloatVectorValues extends QuantizedByteVectorValues { + private final FloatVectorValues values; + private final ScalarQuantizer quantizer; + private final byte[] quantizedVector; + private int lastOrd = -1; + private float offsetValue = 0f; + + private final VectorSimilarityFunction vectorSimilarityFunction; + + QuantizedFloatVectorValues(FloatVectorValues values, VectorSimilarityFunction vectorSimilarityFunction, ScalarQuantizer quantizer) { + this.values = values; + this.quantizer = quantizer; + this.quantizedVector = new byte[values.dimension()]; + this.vectorSimilarityFunction = vectorSimilarityFunction; + } + + @Override + public float getScoreCorrectionConstant(int ord) { + if (ord != lastOrd) { + throw new IllegalStateException( + "attempt to retrieve score correction for different ord " + ord + " than the quantization was done for: " + lastOrd + ); + } + return offsetValue; + } + + @Override + public int dimension() { + return values.dimension(); + } + + @Override + public int size() { + return values.size(); + } + + @Override + public byte[] vectorValue(int ord) throws IOException { + if (ord != lastOrd) { + offsetValue = quantize(ord); + lastOrd = ord; + } + return quantizedVector; + } + + @Override + public VectorScorer scorer(float[] target) throws IOException { + throw new UnsupportedOperationException(); + } + + private float quantize(int ord) throws IOException { + return quantizer.quantize(values.vectorValue(ord), quantizedVector, vectorSimilarityFunction); + } + + @Override + public int ordToDoc(int ord) { + return values.ordToDoc(ord); + } + + @Override + public DocIndexIterator iterator() { + return values.iterator(); + } + } + + static final class ScalarQuantizedCloseableRandomVectorScorerSupplier implements CloseableRandomVectorScorerSupplier { + + private final RandomVectorScorerSupplier supplier; + private final Closeable onClose; + private final int numVectors; + + ScalarQuantizedCloseableRandomVectorScorerSupplier(Closeable onClose, int numVectors, RandomVectorScorerSupplier supplier) { + this.onClose = onClose; + this.supplier = supplier; + this.numVectors = numVectors; + } + + @Override + public UpdateableRandomVectorScorer scorer() throws IOException { + return supplier.scorer(); + } + + @Override + public RandomVectorScorerSupplier copy() throws IOException { + return supplier.copy(); + } + + @Override + public void close() throws IOException { + onClose.close(); + } + + @Override + public int totalVectorCount() { + return numVectors; + } + } + + static final class OffsetCorrectedQuantizedByteVectorValues extends QuantizedByteVectorValues { + + private final QuantizedByteVectorValues in; + private final VectorSimilarityFunction vectorSimilarityFunction; + private final ScalarQuantizer scalarQuantizer, oldScalarQuantizer; + + OffsetCorrectedQuantizedByteVectorValues( + QuantizedByteVectorValues in, + VectorSimilarityFunction vectorSimilarityFunction, + ScalarQuantizer scalarQuantizer, + ScalarQuantizer oldScalarQuantizer + ) { + this.in = in; + this.vectorSimilarityFunction = vectorSimilarityFunction; + this.scalarQuantizer = scalarQuantizer; + this.oldScalarQuantizer = oldScalarQuantizer; + } + + @Override + public float getScoreCorrectionConstant(int ord) throws IOException { + return scalarQuantizer.recalculateCorrectiveOffset(in.vectorValue(ord), oldScalarQuantizer, vectorSimilarityFunction); + } + + @Override + public int dimension() { + return in.dimension(); + } + + @Override + public int size() { + return in.size(); + } + + @Override + public byte[] vectorValue(int ord) throws IOException { + return in.vectorValue(ord); + } + + @Override + public int ordToDoc(int ord) { + return in.ordToDoc(ord); + } + + @Override + public DocIndexIterator iterator() { + return in.iterator(); + } + } + + static final class NormalizedFloatVectorValues extends FloatVectorValues { + private final FloatVectorValues values; + private final float[] normalizedVector; + + NormalizedFloatVectorValues(FloatVectorValues values) { + this.values = values; + this.normalizedVector = 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 { + System.arraycopy(values.vectorValue(ord), 0, normalizedVector, 0, normalizedVector.length); + VectorUtil.l2normalize(normalizedVector); + return normalizedVector; + } + + @Override + public DocIndexIterator iterator() { + return values.iterator(); + } + + @Override + public NormalizedFloatVectorValues copy() throws IOException { + return new NormalizedFloatVectorValues(values.copy()); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapQuantizedByteVectorValues.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapQuantizedByteVectorValues.java new file mode 100644 index 0000000000000..5f303fa60cee9 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapQuantizedByteVectorValues.java @@ -0,0 +1,403 @@ +/* + * @notice + * 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. + * + * Modifications copyright (C) 2024 Elasticsearch B.V. + */ + +package org.elasticsearch.index.codec.vectors; + +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene90.IndexedDISI; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +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.QuantizedByteVectorValues; +import org.apache.lucene.util.quantization.ScalarQuantizer; + +import java.io.IOException; +import java.nio.ByteBuffer; + +/** + * Copied from Lucene 10.3. + * Read the quantized vector values and their score correction values from the index input. This + * supports both iterated and random access. + */ +public abstract class OffHeapQuantizedByteVectorValues extends QuantizedByteVectorValues { + + final int dimension; + final int size; + final int numBytes; + final ScalarQuantizer scalarQuantizer; + final VectorSimilarityFunction similarityFunction; + final FlatVectorsScorer vectorsScorer; + final boolean compress; + + final IndexInput slice; + final byte[] binaryValue; + final ByteBuffer byteBuffer; + final int byteSize; + int lastOrd = -1; + final float[] scoreCorrectionConstant = new float[1]; + + static void decompressBytes(byte[] compressed, int numBytes) { + if (numBytes == compressed.length) { + return; + } + if (numBytes << 1 != compressed.length) { + throw new IllegalArgumentException("numBytes: " + numBytes + " does not match compressed length: " + compressed.length); + } + for (int i = 0; i < numBytes; ++i) { + compressed[numBytes + i] = (byte) (compressed[i] & 0x0F); + compressed[i] = (byte) ((compressed[i] & 0xFF) >> 4); + } + } + + static byte[] compressedArray(int dimension, byte bits) { + if (bits <= 4) { + return new byte[(dimension + 1) >> 1]; + } else { + return null; + } + } + + static void compressBytes(byte[] raw, byte[] compressed) { + if (compressed.length != ((raw.length + 1) >> 1)) { + throw new IllegalArgumentException("compressed length: " + compressed.length + " does not match raw length: " + raw.length); + } + for (int i = 0; i < compressed.length; ++i) { + int v = (raw[i] << 4) | raw[compressed.length + i]; + compressed[i] = (byte) v; + } + } + + OffHeapQuantizedByteVectorValues( + int dimension, + int size, + ScalarQuantizer scalarQuantizer, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + boolean compress, + IndexInput slice + ) { + this.dimension = dimension; + this.size = size; + this.slice = slice; + this.scalarQuantizer = scalarQuantizer; + this.compress = compress; + if (scalarQuantizer.getBits() <= 4 && compress) { + this.numBytes = (dimension + 1) >> 1; + } else { + this.numBytes = dimension; + } + this.byteSize = this.numBytes + Float.BYTES; + byteBuffer = ByteBuffer.allocate(dimension); + binaryValue = byteBuffer.array(); + this.similarityFunction = similarityFunction; + this.vectorsScorer = vectorsScorer; + } + + @Override + public ScalarQuantizer getScalarQuantizer() { + return scalarQuantizer; + } + + @Override + public int dimension() { + return dimension; + } + + @Override + public int size() { + return size; + } + + @Override + public byte[] vectorValue(int targetOrd) throws IOException { + if (lastOrd == targetOrd) { + return binaryValue; + } + slice.seek((long) targetOrd * byteSize); + slice.readBytes(byteBuffer.array(), byteBuffer.arrayOffset(), numBytes); + slice.readFloats(scoreCorrectionConstant, 0, 1); + decompressBytes(binaryValue, numBytes); + lastOrd = targetOrd; + return binaryValue; + } + + @Override + public float getScoreCorrectionConstant(int targetOrd) throws IOException { + if (lastOrd == targetOrd) { + return scoreCorrectionConstant[0]; + } + slice.seek(((long) targetOrd * byteSize) + numBytes); + slice.readFloats(scoreCorrectionConstant, 0, 1); + return scoreCorrectionConstant[0]; + } + + @Override + public IndexInput getSlice() { + return slice; + } + + @Override + public int getVectorByteLength() { + return numBytes; + } + + static OffHeapQuantizedByteVectorValues load( + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int size, + ScalarQuantizer scalarQuantizer, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + boolean compress, + long quantizedVectorDataOffset, + long quantizedVectorDataLength, + IndexInput vectorData + ) throws IOException { + if (configuration.isEmpty()) { + return new EmptyOffHeapVectorValues(dimension, similarityFunction, vectorsScorer); + } + IndexInput bytesSlice = vectorData.slice("quantized-vector-data", quantizedVectorDataOffset, quantizedVectorDataLength); + if (configuration.isDense()) { + return new DenseOffHeapVectorValues(dimension, size, scalarQuantizer, compress, similarityFunction, vectorsScorer, bytesSlice); + } else { + return new SparseOffHeapVectorValues( + configuration, + dimension, + size, + scalarQuantizer, + compress, + vectorData, + similarityFunction, + vectorsScorer, + bytesSlice + ); + } + } + + /** + * Dense vector values that are stored off-heap. This is the most common case when every doc has a + * vector. + */ + public static class DenseOffHeapVectorValues extends OffHeapQuantizedByteVectorValues { + + /** + * Create dense off-heap vector values + * + * @param dimension vector dimension + * @param size number of vectors + * @param scalarQuantizer the scalar quantizer + * @param compress whether the vectors are compressed + * @param similarityFunction the similarity function + * @param vectorsScorer the vectors scorer + * @param slice the index input slice containing the vector data + */ + public DenseOffHeapVectorValues( + int dimension, + int size, + ScalarQuantizer scalarQuantizer, + boolean compress, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + IndexInput slice + ) { + super(dimension, size, scalarQuantizer, similarityFunction, vectorsScorer, compress, slice); + } + + @Override + public DenseOffHeapVectorValues copy() throws IOException { + return new DenseOffHeapVectorValues( + dimension, + size, + scalarQuantizer, + compress, + similarityFunction, + vectorsScorer, + slice.clone() + ); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return acceptDocs; + } + + @Override + public VectorScorer scorer(float[] target) throws IOException { + DenseOffHeapVectorValues copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(similarityFunction, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return vectorScorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + }; + } + + @Override + public DocIndexIterator iterator() { + return createDenseIterator(); + } + } + + private static class SparseOffHeapVectorValues extends OffHeapQuantizedByteVectorValues { + private final DirectMonotonicReader ordToDoc; + private final IndexedDISI disi; + // dataIn was used to init a new IndexedDIS for #randomAccess() + private final IndexInput dataIn; + private final OrdToDocDISIReaderConfiguration configuration; + + SparseOffHeapVectorValues( + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int size, + ScalarQuantizer scalarQuantizer, + boolean compress, + IndexInput dataIn, + VectorSimilarityFunction similarityFunction, + FlatVectorsScorer vectorsScorer, + IndexInput slice + ) throws IOException { + super(dimension, size, scalarQuantizer, similarityFunction, vectorsScorer, compress, slice); + this.configuration = configuration; + this.dataIn = dataIn; + this.ordToDoc = configuration.getDirectMonotonicReader(dataIn); + this.disi = configuration.getIndexedDISI(dataIn); + } + + @Override + public DocIndexIterator iterator() { + return IndexedDISI.asDocIndexIterator(disi); + } + + @Override + public SparseOffHeapVectorValues copy() throws IOException { + return new SparseOffHeapVectorValues( + configuration, + dimension, + size, + scalarQuantizer, + compress, + 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 VectorScorer scorer(float[] target) throws IOException { + SparseOffHeapVectorValues copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(similarityFunction, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return vectorScorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + }; + } + } + + private static class EmptyOffHeapVectorValues extends OffHeapQuantizedByteVectorValues { + + EmptyOffHeapVectorValues(int dimension, VectorSimilarityFunction similarityFunction, FlatVectorsScorer vectorsScorer) { + super(dimension, 0, new ScalarQuantizer(-1, 1, (byte) 7), similarityFunction, vectorsScorer, false, null); + } + + @Override + public int dimension() { + return super.dimension(); + } + + @Override + public int size() { + return 0; + } + + @Override + public DocIndexIterator iterator() { + return createDenseIterator(); + } + + @Override + public EmptyOffHeapVectorValues copy() { + throw new UnsupportedOperationException(); + } + + @Override + public byte[] vectorValue(int targetOrd) { + throw new UnsupportedOperationException(); + } + + @Override + public int ordToDoc(int ord) { + throw new UnsupportedOperationException(); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return null; + } + + @Override + public VectorScorer scorer(float[] target) { + return null; + } + } +} From 457d16c11cd66196c8ab55c0ac0fa9e46e1365d6 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 25 Sep 2025 11:54:41 +0100 Subject: [PATCH 2/2] Package moves --- .../elasticsearch/benchmark/vector/Int7uScorerBenchmark.java | 2 +- .../elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmark.java b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmark.java index 7529e39522582..e4a1ea5a2f3c0 100644 --- a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmark.java +++ b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmark.java @@ -10,7 +10,6 @@ package org.elasticsearch.benchmark.vector; import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorScorer; -import org.apache.lucene.codecs.lucene99.OffHeapQuantizedByteVectorValues; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IOContext; @@ -24,6 +23,7 @@ import org.apache.lucene.util.quantization.ScalarQuantizer; import org.elasticsearch.common.logging.LogConfigurator; import org.elasticsearch.core.IOUtils; +import org.elasticsearch.index.codec.vectors.OffHeapQuantizedByteVectorValues; import org.elasticsearch.logging.LogManager; import org.elasticsearch.logging.Logger; import org.elasticsearch.simdvec.VectorScorerFactory; diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java index eca67582b7a16..761f7ebd3f554 100644 --- a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java +++ b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java @@ -12,7 +12,6 @@ import com.carrotsearch.randomizedtesting.generators.RandomNumbers; import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorScorer; -import org.apache.lucene.codecs.lucene99.OffHeapQuantizedByteVectorValues; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IOContext; @@ -24,6 +23,7 @@ import org.apache.lucene.util.quantization.QuantizedByteVectorValues; import org.apache.lucene.util.quantization.ScalarQuantizedVectorSimilarity; import org.apache.lucene.util.quantization.ScalarQuantizer; +import org.elasticsearch.index.codec.vectors.OffHeapQuantizedByteVectorValues; import java.io.IOException; import java.util.Arrays;