From c10fd7605e4d7efa9671312d44401e93fcd0bf6b Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Wed, 15 Oct 2025 15:11:04 +0100 Subject: [PATCH 01/15] Add HNSW scalar quantized bfloat16 implementation --- server/src/main/java/module-info.java | 3 +- .../ES93HnswScalarQuantizedVectorsFormat.java | 64 +++++ .../ES93ScalarQuantizedVectorsFormat.java | 192 +++++++++++++++ .../org.apache.lucene.codecs.KnnVectorsFormat | 1 + ...arQuantizedBFloat16VectorsFormatTests.java | 103 ++++++++ ...HnswScalarQuantizedVectorsFormatTests.java | 220 ++++++++++++++++++ 6 files changed, 582 insertions(+), 1 deletion(-) create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 2987b3849e663..234d43315a9d4 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -464,6 +464,7 @@ org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; @@ -494,6 +495,6 @@ exports org.elasticsearch.inference.telemetry; exports org.elasticsearch.index.codec.vectors.diskbbq to org.elasticsearch.test.knn; exports org.elasticsearch.index.codec.vectors.cluster to org.elasticsearch.test.knn; - exports org.elasticsearch.index.codec.vectors.es93 to org.elasticsearch.test.knn; exports org.elasticsearch.search.crossproject; + exports org.elasticsearch.index.codec.vectors.es93 to org.elasticsearch.gpu, org.elasticsearch.test.knn; } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java new file mode 100644 index 0000000000000..fe5e2c5696371 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -0,0 +1,64 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.elasticsearch.index.codec.vectors.AbstractHnswVectorsFormat; + +import java.io.IOException; + +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; + +public class ES93HnswScalarQuantizedVectorsFormat extends AbstractHnswVectorsFormat { + + static final String NAME = "ES93HnswScalarQuantizedVectorsFormat"; + + /** The format for storing, reading, merging vectors on disk */ + private final FlatVectorsFormat flatVectorsFormat; + + public ES93HnswScalarQuantizedVectorsFormat() { + this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, false, null, 7, false, false); + } + + public ES93HnswScalarQuantizedVectorsFormat( + int maxConn, + int beamWidth, + boolean useBFloat16, + Float confidenceInterval, + int bits, + boolean compress, + boolean useDirectIO + ) { + super(NAME, maxConn, beamWidth); + this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat(useBFloat16, confidenceInterval, bits, compress, useDirectIO); + } + + @Override + protected FlatVectorsFormat flatVectorsFormat() { + return flatVectorsFormat; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return new Lucene99HnswVectorsWriter(state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), numMergeWorkers, mergeExec); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state)); + } +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java new file mode 100644 index 0000000000000..f1dd1a283f88f --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -0,0 +1,192 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +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.FlatVectorsScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.elasticsearch.simdvec.VectorScorerFactory; +import org.elasticsearch.simdvec.VectorSimilarityType; + +import java.io.IOException; + +import static org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat.DYNAMIC_CONFIDENCE_INTERVAL; +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; + +public class ES93ScalarQuantizedVectorsFormat extends FlatVectorsFormat { + + static final String NAME = "ES93ScalarQuantizedVectorsFormat"; + private static final int ALLOWED_BITS = (1 << 8) | (1 << 7) | (1 << 4); + + static final FlatVectorsScorer flatVectorScorer = new ESFlatVectorsScorer( + new ScalarQuantizedVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()) + ); + + /** The minimum confidence interval */ + private static final float MINIMUM_CONFIDENCE_INTERVAL = 0.9f; + + /** The maximum confidence interval */ + private static final float MAXIMUM_CONFIDENCE_INTERVAL = 1f; + + private final FlatVectorsFormat rawVectorFormat; + + /** + * Controls the confidence interval used to scalar quantize the vectors the default value is + * calculated as `1-1/(vector_dimensions + 1)` + */ + public final Float confidenceInterval; + + private final byte bits; + private final boolean compress; + + public ES93ScalarQuantizedVectorsFormat( + boolean useBFloat16, + Float confidenceInterval, + int bits, + boolean compress, + boolean useDirectIO + ) { + super(NAME); + if (confidenceInterval != null + && confidenceInterval != DYNAMIC_CONFIDENCE_INTERVAL + && (confidenceInterval < MINIMUM_CONFIDENCE_INTERVAL || confidenceInterval > MAXIMUM_CONFIDENCE_INTERVAL)) { + throw new IllegalArgumentException( + "confidenceInterval must be between " + + MINIMUM_CONFIDENCE_INTERVAL + + " and " + + MAXIMUM_CONFIDENCE_INTERVAL + + "; confidenceInterval=" + + confidenceInterval + ); + } + if (bits < 1 || bits > 8 || (ALLOWED_BITS & (1 << bits)) == 0) { + throw new IllegalArgumentException("bits must be one of: 4, 7, 8; bits=" + bits); + } + this.confidenceInterval = confidenceInterval; + this.bits = (byte) bits; + this.compress = compress; + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(useBFloat16, useDirectIO); + } + + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMS_COUNT; + } + + @Override + public String toString() { + return NAME + + "(name=" + + NAME + + ", confidenceInterval=" + + confidenceInterval + + ", bits=" + + bits + + ", compressed=" + + compress + + ", flatVectorScorer=" + + flatVectorScorer + + ", rawVectorFormat=" + + rawVectorFormat + + ")"; + } + + @Override + public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return new Lucene99ScalarQuantizedVectorsWriter( + state, + confidenceInterval, + bits, + compress, + rawVectorFormat.fieldsWriter(state), + flatVectorScorer + ); + } + + @Override + public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new Lucene99ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); + } + + static final class ESFlatVectorsScorer implements FlatVectorsScorer { + + final FlatVectorsScorer delegate; + final VectorScorerFactory factory; + + ESFlatVectorsScorer(FlatVectorsScorer delegate) { + this.delegate = delegate; + factory = VectorScorerFactory.instance().orElse(null); + } + + @Override + public String toString() { + return "ESFlatVectorsScorer(" + "delegate=" + delegate + ", factory=" + factory + ')'; + } + + @Override + public RandomVectorScorerSupplier getRandomVectorScorerSupplier(VectorSimilarityFunction sim, KnnVectorValues values) + throws IOException { + if (values instanceof QuantizedByteVectorValues qValues && qValues.getSlice() != null) { + // TODO: optimize int4 quantization + if (qValues.getScalarQuantizer().getBits() != 7) { + return delegate.getRandomVectorScorerSupplier(sim, values); + } + if (factory != null) { + var scorer = factory.getInt7SQVectorScorerSupplier( + VectorSimilarityType.of(sim), + qValues.getSlice(), + qValues, + qValues.getScalarQuantizer().getConstantMultiplier() + ); + if (scorer.isPresent()) { + return scorer.get(); + } + } + } + return delegate.getRandomVectorScorerSupplier(sim, values); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(VectorSimilarityFunction sim, KnnVectorValues values, float[] query) + throws IOException { + if (values instanceof QuantizedByteVectorValues qValues && qValues.getSlice() != null) { + // TODO: optimize int4 quantization + if (qValues.getScalarQuantizer().getBits() != 7) { + return delegate.getRandomVectorScorer(sim, values, query); + } + if (factory != null) { + var scorer = factory.getInt7SQVectorScorer(sim, qValues, query); + if (scorer.isPresent()) { + return scorer.get(); + } + } + } + return delegate.getRandomVectorScorer(sim, values, query); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(VectorSimilarityFunction sim, KnnVectorValues values, byte[] query) + throws IOException { + return delegate.getRandomVectorScorer(sim, values, query); + } + } +} diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 6c21437d71d28..8e8e17f4681ec 100644 --- a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -9,5 +9,6 @@ org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat +org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java new file mode 100644 index 0000000000000..5688ede94e1ec --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java @@ -0,0 +1,103 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.index.VectorEncoding; + +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import static org.hamcrest.Matchers.closeTo; + +public class ES93HnswScalarQuantizedBFloat16VectorsFormatTests extends ES93HnswScalarQuantizedVectorsFormatTests { + @Override + boolean useBFloat16() { + return true; + } + + @Override + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Override + public void testEmptyByteVectorData() throws Exception { + // no bytes + } + + @Override + public void testMergingWithDifferentByteKnnFields() throws Exception { + // no bytes + } + + @Override + public void testByteVectorScorerIteration() throws Exception { + // no bytes + } + + @Override + public void testSortedIndexBytes() throws Exception { + // no bytes + } + + @Override + public void testMismatchedFields() throws Exception { + // no bytes + } + + @Override + public void testRandomBytes() throws Exception { + // no bytes + } + + @Override + public void testWriterRamEstimate() throws Exception { + // estimate is different due to bfloat16 + } + + @Override + public void testRandom() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testRandom); + assertFloatsWithinBounds(err); + } + + @Override + public void testRandomWithUpdatesAndGraph() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph); + assertFloatsWithinBounds(err); + } + + @Override + public void testSparseVectors() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors); + assertFloatsWithinBounds(err); + } + + @Override + public void testVectorValuesReportCorrectDocs() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs); + assertFloatsWithinBounds(err); + } + + private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>"); + + private static void assertFloatsWithinBounds(AssertionError error) { + Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage()); + if (m.matches() == false) { + throw error; // nothing to do with us, just rethrow + } + + // numbers just need to be in the same vicinity + double expected = Double.parseDouble(m.group(1)); + double actual = Double.parseDouble(m.group(2)); + double allowedError = expected * 0.01; // within 1% + assertThat(error.getMessage(), actual, closeTo(expected, allowedError)); + } +} diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java new file mode 100644 index 0000000000000..a6c12548821b9 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java @@ -0,0 +1,220 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.StoredFields; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.MMapDirectory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.elasticsearch.index.codec.vectors.BFloat16; + +import java.io.IOException; +import java.nio.file.Path; + +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +public class ES93HnswScalarQuantizedVectorsFormatTests extends BaseKnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized + } + + boolean useBFloat16() { + return false; + } + + @Override + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat( + new ES93HnswScalarQuantizedVectorsFormat(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, useBFloat16(), null, 7, false, false) + ); + } + + // The following test scenarios are similar to their superclass namesakes, + // but here we ensure that the Directory implementation is a FSDirectory + // which helps test the native code vector distance implementation + + public void testAddIndexesDirectory0FS() throws Exception { + Path root = createTempDir(); + String fieldName = "field"; + Document doc = new Document(); + doc.add(new KnnFloatVectorField(fieldName, new float[4], VectorSimilarityFunction.DOT_PRODUCT)); + try (Directory dir = new MMapDirectory(root.resolve("dir1")); Directory dir2 = new MMapDirectory(root.resolve("dir2"))) { + try (IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + w.addDocument(doc); + } + try (IndexWriter w2 = new IndexWriter(dir2, newIndexWriterConfig())) { + w2.addIndexes(dir); + w2.forceMerge(1); + try (IndexReader reader = DirectoryReader.open(w2)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues vectorValues = r.getFloatVectorValues(fieldName); + KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator(); + assertEquals(0, iterator.nextDoc()); + assertEquals(0, vectorValues.vectorValue(iterator.index())[0], 0); + assertEquals(NO_MORE_DOCS, iterator.nextDoc()); + } + } + } + } + + public void testAddIndexesDirectory01FSCosine() throws Exception { + testAddIndexesDirectory01FS(VectorSimilarityFunction.COSINE); + } + + public void testAddIndexesDirectory01FSDot() throws Exception { + testAddIndexesDirectory01FS(VectorSimilarityFunction.DOT_PRODUCT); + } + + public void testAddIndexesDirectory01FSEuclidean() throws Exception { + testAddIndexesDirectory01FS(VectorSimilarityFunction.EUCLIDEAN); + } + + public void testAddIndexesDirectory01FSMaxIP() throws Exception { + testAddIndexesDirectory01FS(VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT); + } + + private void testAddIndexesDirectory01FS(VectorSimilarityFunction similarityFunction) throws Exception { + Path root = createTempDir(); + String fieldName = "field"; + float[] vector = new float[] { 1f }; + Document doc = new Document(); + doc.add(new KnnFloatVectorField(fieldName, vector, similarityFunction)); + try (Directory dir = new MMapDirectory(root.resolve("dir1")); Directory dir2 = new MMapDirectory(root.resolve("dir2"))) { + try (IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + w.addDocument(doc); + } + try (IndexWriter w2 = new IndexWriter(dir2, newIndexWriterConfig())) { + vector[0] = 2f; + w2.addDocument(doc); + w2.addIndexes(dir); + w2.forceMerge(1); + try (IndexReader reader = DirectoryReader.open(w2)) { + LeafReader r = getOnlyLeafReader(reader); + FloatVectorValues vectorValues = r.getFloatVectorValues(fieldName); + KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator(); + assertEquals(0, iterator.nextDoc()); + // The merge order is randomized, we might get 1 first, or 2 + float value = vectorValues.vectorValue(iterator.index())[0]; + assertTrue(value == 1 || value == 2); + assertEquals(1, iterator.nextDoc()); + value += vectorValues.vectorValue(iterator.index())[0]; + assertEquals(3f, value, 0); + } + } + } + } + + public void testSingleVectorPerSegmentCosine() throws Exception { + testSingleVectorPerSegment(VectorSimilarityFunction.COSINE); + } + + public void testSingleVectorPerSegmentDot() throws Exception { + testSingleVectorPerSegment(VectorSimilarityFunction.DOT_PRODUCT); + } + + public void testSingleVectorPerSegmentEuclidean() throws Exception { + testSingleVectorPerSegment(VectorSimilarityFunction.EUCLIDEAN); + } + + public void testSingleVectorPerSegmentMIP() throws Exception { + testSingleVectorPerSegment(VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT); + } + + private void testSingleVectorPerSegment(VectorSimilarityFunction sim) throws Exception { + var codec = getCodec(); + try (Directory dir = new MMapDirectory(createTempDir().resolve("dir1"))) { + try (IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig().setCodec(codec))) { + Document doc2 = new Document(); + doc2.add(new KnnFloatVectorField("field", new float[] { 0.8f, 0.6f }, sim)); + doc2.add(newTextField("id", "A", Field.Store.YES)); + writer.addDocument(doc2); + writer.commit(); + + Document doc1 = new Document(); + doc1.add(new KnnFloatVectorField("field", new float[] { 0.6f, 0.8f }, sim)); + doc1.add(newTextField("id", "B", Field.Store.YES)); + writer.addDocument(doc1); + writer.commit(); + + Document doc3 = new Document(); + doc3.add(new KnnFloatVectorField("field", new float[] { -0.6f, -0.8f }, sim)); + doc3.add(newTextField("id", "C", Field.Store.YES)); + writer.addDocument(doc3); + writer.commit(); + + writer.forceMerge(1); + } + try (DirectoryReader reader = DirectoryReader.open(dir)) { + LeafReader leafReader = getOnlyLeafReader(reader); + StoredFields storedFields = reader.storedFields(); + float[] queryVector = new float[] { 0.6f, 0.8f }; + var hits = leafReader.searchNearestVectors( + "field", + queryVector, + 3, + AcceptDocs.fromLiveDocs(leafReader.getLiveDocs(), leafReader.maxDoc()), + 100 + ); + assertEquals(hits.scoreDocs.length, 3); + assertEquals("B", storedFields.document(hits.scoreDocs[0].doc).get("id")); + assertEquals("A", storedFields.document(hits.scoreDocs[1].doc).get("id")); + assertEquals("C", storedFields.document(hits.scoreDocs[2].doc).get("id")); + } + } + } + + public void testSimpleOffHeapSize() throws IOException { + float[] vector = randomVector(random().nextInt(12, 500)); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT)); + w.addDocument(doc); + w.commit(); + try (IndexReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnVectorsReader = fieldsReader.getFieldReader("f"); + } + var fieldInfo = r.getFieldInfos().fieldInfo("f"); + var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); + assertEquals(3, offHeap.size()); + int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES; + assertEquals(vector.length * bytes, (long) offHeap.get("vec")); + assertEquals(1L, (long) offHeap.get("vex")); + assertTrue(offHeap.get("veq") > 0L); + } + } + } + } +} From 3a6f7fccdeeb90fc2b98646ca06992333f275013 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Wed, 15 Oct 2025 15:26:40 +0100 Subject: [PATCH 02/15] Add flat format --- server/src/main/java/module-info.java | 1 + .../vectors/es93/ES93FlatVectorFormat.java | 124 ++++++++++++++++++ .../org.apache.lucene.codecs.KnnVectorsFormat | 1 + .../ES93FlatBFloat16VectorFormatTests.java | 103 +++++++++++++++ .../es93/ES93FlatVectorFormatTests.java | 75 +++++++++++ 5 files changed, 304 insertions(+) create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 234d43315a9d4..1d121c38d3c6c 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -464,6 +464,7 @@ org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java new file mode 100644 index 0000000000000..bd421abf167f0 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java @@ -0,0 +1,124 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; +import org.apache.lucene.util.hnsw.RandomVectorScorer; + +import java.io.IOException; +import java.util.Map; + +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; + +public class ES93FlatVectorFormat extends KnnVectorsFormat { + + static final String NAME = "ES93FlatVectorFormat"; + + private final FlatVectorsFormat format; + + /** + * Sole constructor + */ + public ES93FlatVectorFormat() { + super(NAME); + format = new ES93GenericFlatVectorsFormat(); + } + + public ES93FlatVectorFormat(boolean useBFloat16) { + super(NAME); + format = new ES93GenericFlatVectorsFormat(useBFloat16, false); + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return format.fieldsWriter(state); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new ES93FlatVectorReader(format.fieldsReader(state)); + } + + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMS_COUNT; + } + + static class ES93FlatVectorReader extends KnnVectorsReader { + + private final FlatVectorsReader reader; + + ES93FlatVectorReader(FlatVectorsReader reader) { + super(); + this.reader = reader; + } + + @Override + public void checkIntegrity() throws IOException { + reader.checkIntegrity(); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return reader.getFloatVectorValues(field); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return reader.getByteVectorValues(field); + } + + @Override + public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) + throws IOException { + OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); + Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); + for (int i = 0; i < scorer.maxOrd(); i++) { + if (acceptedOrds == null || acceptedOrds.get(i)) { + collector.collect(i, scorer.score(i)); + collector.incVisitedCount(1); + } + } + assert collector.earlyTerminated() == false; + } + + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + return reader.getOffHeapByteSize(fieldInfo); + } + + @Override + public void close() throws IOException { + reader.close(); + } + } +} diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 8e8e17f4681ec..8ca17b76098c4 100644 --- a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -9,6 +9,7 @@ org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat +org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java new file mode 100644 index 0000000000000..4d017daa2724c --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java @@ -0,0 +1,103 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.index.VectorEncoding; + +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import static org.hamcrest.Matchers.closeTo; + +public class ES93FlatBFloat16VectorFormatTests extends ES93FlatVectorFormatTests { + @Override + boolean useBFloat16() { + return true; + } + + @Override + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Override + public void testEmptyByteVectorData() throws Exception { + // no bytes + } + + @Override + public void testMergingWithDifferentByteKnnFields() throws Exception { + // no bytes + } + + @Override + public void testByteVectorScorerIteration() throws Exception { + // no bytes + } + + @Override + public void testSortedIndexBytes() throws Exception { + // no bytes + } + + @Override + public void testMismatchedFields() throws Exception { + // no bytes + } + + @Override + public void testRandomBytes() throws Exception { + // no bytes + } + + @Override + public void testWriterRamEstimate() throws Exception { + // estimate is different due to bfloat16 + } + + @Override + public void testRandom() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testRandom); + assertFloatsWithinBounds(err); + } + + @Override + public void testRandomWithUpdatesAndGraph() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph); + assertFloatsWithinBounds(err); + } + + @Override + public void testSparseVectors() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors); + assertFloatsWithinBounds(err); + } + + @Override + public void testVectorValuesReportCorrectDocs() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs); + assertFloatsWithinBounds(err); + } + + private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>"); + + private static void assertFloatsWithinBounds(AssertionError error) { + Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage()); + if (m.matches() == false) { + throw error; // nothing to do with us, just rethrow + } + + // numbers just need to be in the same vicinity + double expected = Double.parseDouble(m.group(1)); + double actual = Double.parseDouble(m.group(2)); + double allowedError = expected * 0.01; // within 1% + assertThat(error.getMessage(), actual, closeTo(expected, allowedError)); + } +} diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java new file mode 100644 index 0000000000000..f3faa98124a66 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java @@ -0,0 +1,75 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.elasticsearch.index.codec.vectors.BFloat16; + +import java.io.IOException; + +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; + +public class ES93FlatVectorFormatTests extends BaseKnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized + } + + boolean useBFloat16() { + return false; + } + + @Override + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(useBFloat16())); + } + + public void testSearchWithVisitedLimit() { + // requires graph-based vector codec + } + + public void testSimpleOffHeapSize() throws IOException { + float[] vector = randomVector(random().nextInt(12, 500)); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT)); + w.addDocument(doc); + w.commit(); + try (IndexReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnVectorsReader = fieldsReader.getFieldReader("f"); + } + var fieldInfo = r.getFieldInfos().fieldInfo("f"); + var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); + int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES; + assertEquals(vector.length * bytes, (long) offHeap.get("vec")); + assertEquals(1, offHeap.size()); + } + } + } + } +} From a85b7ca8faa7ab91a95775ed45b6e959f2a12481 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 16 Oct 2025 09:21:23 +0100 Subject: [PATCH 03/15] Add int8 implementation --- server/src/main/java/module-info.java | 3 +- .../es93/ES93Int8FlatVectorFormat.java | 129 ++++++++++++++++++ .../org.apache.lucene.codecs.KnnVectorsFormat | 1 + ...ES93Int8FlatBFloat16VectorFormatTests.java | 103 ++++++++++++++ .../es93/ES93Int8FlatVectorFormatTests.java | 77 +++++++++++ 5 files changed, 312 insertions(+), 1 deletion(-) create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 1d121c38d3c6c..a9dc2b8d645f6 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -465,6 +465,7 @@ org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat, + org.elasticsearch.index.codec.vectors.es93.ES93Int8FlatVectorFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; @@ -496,6 +497,6 @@ exports org.elasticsearch.inference.telemetry; exports org.elasticsearch.index.codec.vectors.diskbbq to org.elasticsearch.test.knn; exports org.elasticsearch.index.codec.vectors.cluster to org.elasticsearch.test.knn; + exports org.elasticsearch.index.codec.vectors.es93 to org.elasticsearch.test.knn; exports org.elasticsearch.search.crossproject; - exports org.elasticsearch.index.codec.vectors.es93 to org.elasticsearch.gpu, org.elasticsearch.test.knn; } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java new file mode 100644 index 0000000000000..43cd21fd04b87 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java @@ -0,0 +1,129 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; +import org.apache.lucene.util.hnsw.RandomVectorScorer; + +import java.io.IOException; +import java.util.Map; + +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; + +public class ES93Int8FlatVectorFormat extends KnnVectorsFormat { + + static final String NAME = "ES93Int8FlatVectorFormat"; + + private final FlatVectorsFormat format; + + public ES93Int8FlatVectorFormat() { + this(false, null, 7, false); + } + + public ES93Int8FlatVectorFormat(boolean useBFloat16) { + this(useBFloat16, null, 7, false); + } + + public ES93Int8FlatVectorFormat(boolean useBFloat16, Float confidenceInterval, int bits, boolean compress) { + super(NAME); + this.format = new ES93ScalarQuantizedVectorsFormat(useBFloat16, confidenceInterval, bits, compress, false); + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return format.fieldsWriter(state); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new ES813FlatVectorReader(format.fieldsReader(state)); + } + + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMS_COUNT; + } + + @Override + public String toString() { + return NAME + "(name=" + NAME + ", innerFormat=" + format + ")"; + } + + public static class ES813FlatVectorReader extends KnnVectorsReader { + + private final FlatVectorsReader reader; + + public ES813FlatVectorReader(FlatVectorsReader reader) { + super(); + this.reader = reader; + } + + @Override + public void checkIntegrity() throws IOException { + reader.checkIntegrity(); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return reader.getFloatVectorValues(field); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return reader.getByteVectorValues(field); + } + + @Override + public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) + throws IOException { + OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); + Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); + for (int i = 0; i < scorer.maxOrd(); i++) { + if (acceptedOrds == null || acceptedOrds.get(i)) { + collector.collect(i, scorer.score(i)); + collector.incVisitedCount(1); + } + } + assert collector.earlyTerminated() == false; + } + + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + return reader.getOffHeapByteSize(fieldInfo); + } + + @Override + public void close() throws IOException { + reader.close(); + } + } +} diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 8ca17b76098c4..4944b734ea018 100644 --- a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -10,6 +10,7 @@ org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsForma org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat +org.elasticsearch.index.codec.vectors.es93.ES93Int8FlatVectorFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java new file mode 100644 index 0000000000000..98d60e4723e28 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java @@ -0,0 +1,103 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.index.VectorEncoding; + +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import static org.hamcrest.Matchers.closeTo; + +public class ES93Int8FlatBFloat16VectorFormatTests extends ES93Int8FlatVectorFormatTests { + @Override + boolean useBFloat16() { + return true; + } + + @Override + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Override + public void testEmptyByteVectorData() throws Exception { + // no bytes + } + + @Override + public void testMergingWithDifferentByteKnnFields() throws Exception { + // no bytes + } + + @Override + public void testByteVectorScorerIteration() throws Exception { + // no bytes + } + + @Override + public void testSortedIndexBytes() throws Exception { + // no bytes + } + + @Override + public void testMismatchedFields() throws Exception { + // no bytes + } + + @Override + public void testRandomBytes() throws Exception { + // no bytes + } + + @Override + public void testWriterRamEstimate() throws Exception { + // estimate is different due to bfloat16 + } + + @Override + public void testRandom() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testRandom); + assertFloatsWithinBounds(err); + } + + @Override + public void testRandomWithUpdatesAndGraph() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph); + assertFloatsWithinBounds(err); + } + + @Override + public void testSparseVectors() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors); + assertFloatsWithinBounds(err); + } + + @Override + public void testVectorValuesReportCorrectDocs() throws Exception { + AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs); + assertFloatsWithinBounds(err); + } + + private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>"); + + private static void assertFloatsWithinBounds(AssertionError error) { + Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage()); + if (m.matches() == false) { + throw error; // nothing to do with us, just rethrow + } + + // numbers just need to be in the same vicinity + double expected = Double.parseDouble(m.group(1)); + double actual = Double.parseDouble(m.group(2)); + double allowedError = expected * 0.01; // within 1% + assertThat(error.getMessage(), actual, closeTo(expected, allowedError)); + } +} diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java new file mode 100644 index 0000000000000..a120b1be50556 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java @@ -0,0 +1,77 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.elasticsearch.index.codec.vectors.BFloat16; + +import java.io.IOException; + +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; + +public class ES93Int8FlatVectorFormatTests extends BaseKnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized + } + + boolean useBFloat16() { + return false; + } + + @Override + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(new ES93Int8FlatVectorFormat(useBFloat16())); + } + + public void testSearchWithVisitedLimit() { + // requires graph vector codec + } + + public void testSimpleOffHeapSize() throws IOException { + float[] vector = randomVector(random().nextInt(12, 500)); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT)); + w.addDocument(doc); + w.commit(); + try (IndexReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnVectorsReader = fieldsReader.getFieldReader("f"); + } + var fieldInfo = r.getFieldInfos().fieldInfo("f"); + var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); + assertEquals(2, offHeap.size()); + int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES; + assertEquals(vector.length * bytes, (long) offHeap.get("vec")); + assertTrue(offHeap.get("veq") > 0L); + } + } + } + } + +} From a638b5993d1e5056b5f02818b1a04fd25991d55b Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Mon, 20 Oct 2025 11:37:33 +0100 Subject: [PATCH 04/15] Rename class --- server/src/main/java/module-info.java | 3 ++- ...=> ES93ScalarQuantizedFlatVectorsFormat.java} | 16 ++++++++-------- .../org.apache.lucene.codecs.KnnVectorsFormat | 2 +- ...rQuantizedFlatBFloat16VectorFormatTests.java} | 2 +- ...93ScalarQuantizedFlatVectorsFormatTests.java} | 4 ++-- 5 files changed, 14 insertions(+), 13 deletions(-) rename server/src/main/java/org/elasticsearch/index/codec/vectors/es93/{ES93Int8FlatVectorFormat.java => ES93ScalarQuantizedFlatVectorsFormat.java} (87%) rename server/src/test/java/org/elasticsearch/index/codec/vectors/es93/{ES93Int8FlatBFloat16VectorFormatTests.java => ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java} (96%) rename server/src/test/java/org/elasticsearch/index/codec/vectors/es93/{ES93Int8FlatVectorFormatTests.java => ES93ScalarQuantizedFlatVectorsFormatTests.java} (93%) diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index a9dc2b8d645f6..fb4ec153ba635 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -7,6 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ +import org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedFlatVectorsFormat; import org.elasticsearch.plugins.internal.RestExtension; import org.elasticsearch.reservedstate.ReservedStateHandlerProvider; @@ -465,7 +466,7 @@ org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat, - org.elasticsearch.index.codec.vectors.es93.ES93Int8FlatVectorFormat, + ES93ScalarQuantizedFlatVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java similarity index 87% rename from server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java rename to server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java index 43cd21fd04b87..4bd80005bfa12 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java @@ -30,21 +30,21 @@ import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; -public class ES93Int8FlatVectorFormat extends KnnVectorsFormat { +public class ES93ScalarQuantizedFlatVectorsFormat extends KnnVectorsFormat { - static final String NAME = "ES93Int8FlatVectorFormat"; + static final String NAME = "ES93ScalarQuantizedFlatVectorsFormat"; private final FlatVectorsFormat format; - public ES93Int8FlatVectorFormat() { + public ES93ScalarQuantizedFlatVectorsFormat() { this(false, null, 7, false); } - public ES93Int8FlatVectorFormat(boolean useBFloat16) { + public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16) { this(useBFloat16, null, 7, false); } - public ES93Int8FlatVectorFormat(boolean useBFloat16, Float confidenceInterval, int bits, boolean compress) { + public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16, Float confidenceInterval, int bits, boolean compress) { super(NAME); this.format = new ES93ScalarQuantizedVectorsFormat(useBFloat16, confidenceInterval, bits, compress, false); } @@ -56,7 +56,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new ES813FlatVectorReader(format.fieldsReader(state)); + return new ES93FlatVectorsReader(format.fieldsReader(state)); } @Override @@ -69,11 +69,11 @@ public String toString() { return NAME + "(name=" + NAME + ", innerFormat=" + format + ")"; } - public static class ES813FlatVectorReader extends KnnVectorsReader { + public static class ES93FlatVectorsReader extends KnnVectorsReader { private final FlatVectorsReader reader; - public ES813FlatVectorReader(FlatVectorsReader reader) { + public ES93FlatVectorsReader(FlatVectorsReader reader) { super(); this.reader = reader; } diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 4944b734ea018..2dcbd21c86436 100644 --- a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -10,7 +10,7 @@ org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsForma org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat -org.elasticsearch.index.codec.vectors.es93.ES93Int8FlatVectorFormat +org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedFlatVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java similarity index 96% rename from server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java rename to server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java index 98d60e4723e28..3681d0a1299eb 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatBFloat16VectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java @@ -16,7 +16,7 @@ import static org.hamcrest.Matchers.closeTo; -public class ES93Int8FlatBFloat16VectorFormatTests extends ES93Int8FlatVectorFormatTests { +public class ES93ScalarQuantizedFlatBFloat16VectorFormatTests extends ES93ScalarQuantizedFlatVectorsFormatTests { @Override boolean useBFloat16() { return true; diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java similarity index 93% rename from server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java rename to server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java index a120b1be50556..a55147f6e5d9f 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93Int8FlatVectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java @@ -29,7 +29,7 @@ import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; -public class ES93Int8FlatVectorFormatTests extends BaseKnnVectorsFormatTestCase { +public class ES93ScalarQuantizedFlatVectorsFormatTests extends BaseKnnVectorsFormatTestCase { static { LogConfigurator.loadLog4jPlugins(); @@ -42,7 +42,7 @@ boolean useBFloat16() { @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93Int8FlatVectorFormat(useBFloat16())); + return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedFlatVectorsFormat(useBFloat16())); } public void testSearchWithVisitedLimit() { From 419032d4360884ece4a157d80c8cbe99a441059a Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Mon, 20 Oct 2025 11:48:11 +0100 Subject: [PATCH 05/15] Improve tests --- .../ES93ScalarQuantizedFlatVectorsFormatTests.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java index a55147f6e5d9f..d05f33941f7b3 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java @@ -28,6 +28,10 @@ import java.io.IOException; import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.hasEntry; public class ES93ScalarQuantizedFlatVectorsFormatTests extends BaseKnnVectorsFormatTestCase { @@ -65,10 +69,9 @@ public void testSimpleOffHeapSize() throws IOException { } var fieldInfo = r.getFieldInfos().fieldInfo("f"); var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); - assertEquals(2, offHeap.size()); - int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES; - assertEquals(vector.length * bytes, (long) offHeap.get("vec")); - assertTrue(offHeap.get("veq") > 0L); + assertThat(offHeap, aMapWithSize(2)); + assertThat(offHeap, hasEntry("vec", vector.length * (useBFloat16() ? BFloat16.BYTES : Float.BYTES))); + assertThat(offHeap, hasEntry(equalTo("veq"), greaterThan(0L))); } } } From 1287a0b1d09d94cd017fa82543d6bffd47de7739 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Mon, 20 Oct 2025 12:16:50 +0100 Subject: [PATCH 06/15] Fix module reference --- server/src/main/java/module-info.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 6cbf55407f174..7fec6ce36874d 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -7,7 +7,6 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedFlatVectorsFormat; import org.elasticsearch.plugins.internal.RestExtension; import org.elasticsearch.reservedstate.ReservedStateHandlerProvider; @@ -467,7 +466,7 @@ org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat, - ES93ScalarQuantizedFlatVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedFlatVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; From 7c343c56c07271108e4ac405a46b637b5c730e3a Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 23 Oct 2025 09:52:32 +0100 Subject: [PATCH 07/15] Update to use new Lucene104 format --- .../ES814ScalarQuantizedVectorsFormat.java | 1 + .../Lucene99ScalarQuantizedVectorsReader.java | 455 ------------------ .../ES93HnswScalarQuantizedVectorsFormat.java | 31 +- .../ES93ScalarQuantizedFlatVectorsFormat.java | 9 +- .../ES93ScalarQuantizedVectorsFormat.java | 143 +----- ...arQuantizedBFloat16VectorsFormatTests.java | 131 +++-- ...HnswScalarQuantizedVectorsFormatTests.java | 237 ++------- 7 files changed, 147 insertions(+), 860 deletions(-) delete mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.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 692ccdfa9222c..dd125b68d292a 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 @@ -9,6 +9,7 @@ package org.elasticsearch.index.codec.vectors; +import org.apache.lucene.backward_codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; 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 deleted file mode 100644 index 9c0e855faf6ad..0000000000000 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.java +++ /dev/null @@ -1,455 +0,0 @@ -/* - * @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/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java index fe5e2c5696371..a2e14df6099ec 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -12,6 +12,7 @@ import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; import org.apache.lucene.index.SegmentReadState; @@ -19,9 +20,7 @@ import org.elasticsearch.index.codec.vectors.AbstractHnswVectorsFormat; import java.io.IOException; - -import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; -import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import java.util.concurrent.ExecutorService; public class ES93HnswScalarQuantizedVectorsFormat extends AbstractHnswVectorsFormat { @@ -31,20 +30,36 @@ public class ES93HnswScalarQuantizedVectorsFormat extends AbstractHnswVectorsFor private final FlatVectorsFormat flatVectorsFormat; public ES93HnswScalarQuantizedVectorsFormat() { - this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, false, null, 7, false, false); + super(NAME); + this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat( + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + false, + false + ); } public ES93HnswScalarQuantizedVectorsFormat( int maxConn, int beamWidth, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, boolean useBFloat16, - Float confidenceInterval, - int bits, - boolean compress, boolean useDirectIO ) { super(NAME, maxConn, beamWidth); - this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat(useBFloat16, confidenceInterval, bits, compress, useDirectIO); + this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat(encoding, useBFloat16, useDirectIO); + } + + public ES93HnswScalarQuantizedVectorsFormat( + int maxConn, + int beamWidth, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, + boolean useBFloat16, + boolean useDirectIO, + int numMergeWorkers, + ExecutorService mergeExec + ) { + super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec); + this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat(encoding, useBFloat16, useDirectIO); } @Override diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java index 4bd80005bfa12..f1375cb5469e7 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java @@ -14,6 +14,7 @@ import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FloatVectorValues; @@ -37,16 +38,16 @@ public class ES93ScalarQuantizedFlatVectorsFormat extends KnnVectorsFormat { private final FlatVectorsFormat format; public ES93ScalarQuantizedFlatVectorsFormat() { - this(false, null, 7, false); + this(false, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); } public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16) { - this(useBFloat16, null, 7, false); + this(useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); } - public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16, Float confidenceInterval, int bits, boolean compress) { + public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding) { super(NAME); - this.format = new ES93ScalarQuantizedVectorsFormat(useBFloat16, confidenceInterval, bits, compress, false); + this.format = new ES93ScalarQuantizedVectorsFormat(encoding, useBFloat16, false); } @Override diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java index f1dd1a283f88f..b393d91dd0c72 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -12,78 +12,36 @@ 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.FlatVectorsScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; -import org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer; -import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; -import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter; -import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.util.hnsw.RandomVectorScorer; -import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; -import org.apache.lucene.util.quantization.QuantizedByteVectorValues; -import org.elasticsearch.simdvec.VectorScorerFactory; -import org.elasticsearch.simdvec.VectorSimilarityType; import java.io.IOException; -import static org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat.DYNAMIC_CONFIDENCE_INTERVAL; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; public class ES93ScalarQuantizedVectorsFormat extends FlatVectorsFormat { static final String NAME = "ES93ScalarQuantizedVectorsFormat"; - private static final int ALLOWED_BITS = (1 << 8) | (1 << 7) | (1 << 4); - static final FlatVectorsScorer flatVectorScorer = new ESFlatVectorsScorer( - new ScalarQuantizedVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()) + static final Lucene104ScalarQuantizedVectorScorer flatVectorScorer = new Lucene104ScalarQuantizedVectorScorer( + FlatVectorScorerUtil.getLucene99FlatVectorsScorer() ); - /** The minimum confidence interval */ - private static final float MINIMUM_CONFIDENCE_INTERVAL = 0.9f; - - /** The maximum confidence interval */ - private static final float MAXIMUM_CONFIDENCE_INTERVAL = 1f; - + private final Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding; private final FlatVectorsFormat rawVectorFormat; - /** - * Controls the confidence interval used to scalar quantize the vectors the default value is - * calculated as `1-1/(vector_dimensions + 1)` - */ - public final Float confidenceInterval; - - private final byte bits; - private final boolean compress; - public ES93ScalarQuantizedVectorsFormat( + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, boolean useBFloat16, - Float confidenceInterval, - int bits, - boolean compress, boolean useDirectIO ) { super(NAME); - if (confidenceInterval != null - && confidenceInterval != DYNAMIC_CONFIDENCE_INTERVAL - && (confidenceInterval < MINIMUM_CONFIDENCE_INTERVAL || confidenceInterval > MAXIMUM_CONFIDENCE_INTERVAL)) { - throw new IllegalArgumentException( - "confidenceInterval must be between " - + MINIMUM_CONFIDENCE_INTERVAL - + " and " - + MAXIMUM_CONFIDENCE_INTERVAL - + "; confidenceInterval=" - + confidenceInterval - ); - } - if (bits < 1 || bits > 8 || (ALLOWED_BITS & (1 << bits)) == 0) { - throw new IllegalArgumentException("bits must be one of: 4, 7, 8; bits=" + bits); - } - this.confidenceInterval = confidenceInterval; - this.bits = (byte) bits; - this.compress = compress; + this.encoding = encoding; this.rawVectorFormat = new ES93GenericFlatVectorsFormat(useBFloat16, useDirectIO); } @@ -97,12 +55,8 @@ public String toString() { return NAME + "(name=" + NAME - + ", confidenceInterval=" - + confidenceInterval - + ", bits=" - + bits - + ", compressed=" - + compress + + ", encoding=" + + encoding + ", flatVectorScorer=" + flatVectorScorer + ", rawVectorFormat=" @@ -112,81 +66,12 @@ public String toString() { @Override public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new Lucene99ScalarQuantizedVectorsWriter( - state, - confidenceInterval, - bits, - compress, - rawVectorFormat.fieldsWriter(state), - flatVectorScorer - ); + return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { + }; } @Override public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new Lucene99ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); - } - - static final class ESFlatVectorsScorer implements FlatVectorsScorer { - - final FlatVectorsScorer delegate; - final VectorScorerFactory factory; - - ESFlatVectorsScorer(FlatVectorsScorer delegate) { - this.delegate = delegate; - factory = VectorScorerFactory.instance().orElse(null); - } - - @Override - public String toString() { - return "ESFlatVectorsScorer(" + "delegate=" + delegate + ", factory=" + factory + ')'; - } - - @Override - public RandomVectorScorerSupplier getRandomVectorScorerSupplier(VectorSimilarityFunction sim, KnnVectorValues values) - throws IOException { - if (values instanceof QuantizedByteVectorValues qValues && qValues.getSlice() != null) { - // TODO: optimize int4 quantization - if (qValues.getScalarQuantizer().getBits() != 7) { - return delegate.getRandomVectorScorerSupplier(sim, values); - } - if (factory != null) { - var scorer = factory.getInt7SQVectorScorerSupplier( - VectorSimilarityType.of(sim), - qValues.getSlice(), - qValues, - qValues.getScalarQuantizer().getConstantMultiplier() - ); - if (scorer.isPresent()) { - return scorer.get(); - } - } - } - return delegate.getRandomVectorScorerSupplier(sim, values); - } - - @Override - public RandomVectorScorer getRandomVectorScorer(VectorSimilarityFunction sim, KnnVectorValues values, float[] query) - throws IOException { - if (values instanceof QuantizedByteVectorValues qValues && qValues.getSlice() != null) { - // TODO: optimize int4 quantization - if (qValues.getScalarQuantizer().getBits() != 7) { - return delegate.getRandomVectorScorer(sim, values, query); - } - if (factory != null) { - var scorer = factory.getInt7SQVectorScorer(sim, qValues, query); - if (scorer.isPresent()) { - return scorer.get(); - } - } - } - return delegate.getRandomVectorScorer(sim, values, query); - } - - @Override - public RandomVectorScorer getRandomVectorScorer(VectorSimilarityFunction sim, KnnVectorValues values, byte[] query) - throws IOException { - return delegate.getRandomVectorScorer(sim, values, query); - } + return new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java index 5688ede94e1ec..e2bc3970398c5 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java @@ -9,95 +9,74 @@ package org.elasticsearch.index.codec.vectors.es93; -import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; +import org.apache.lucene.store.Directory; +import org.elasticsearch.index.codec.vectors.BFloat16; +import org.elasticsearch.index.codec.vectors.BaseHnswBFloat16VectorsFormatTestCase; -import java.util.regex.Matcher; -import java.util.regex.Pattern; +import java.io.IOException; +import java.util.concurrent.ExecutorService; -import static org.hamcrest.Matchers.closeTo; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.allOf; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.hasEntry; -public class ES93HnswScalarQuantizedBFloat16VectorsFormatTests extends ES93HnswScalarQuantizedVectorsFormatTests { - @Override - boolean useBFloat16() { - return true; - } - - @Override - protected VectorEncoding randomVectorEncoding() { - return VectorEncoding.FLOAT32; - } - - @Override - public void testEmptyByteVectorData() throws Exception { - // no bytes - } +public class ES93HnswScalarQuantizedBFloat16VectorsFormatTests extends BaseHnswBFloat16VectorsFormatTestCase { @Override - public void testMergingWithDifferentByteKnnFields() throws Exception { - // no bytes + protected KnnVectorsFormat createFormat() { + return new ES93HnswScalarQuantizedVectorsFormat( + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + true, + random().nextBoolean() + ); } @Override - public void testByteVectorScorerIteration() throws Exception { - // no bytes + protected KnnVectorsFormat createFormat(int maxConn, int beamWidth) { + return new ES93HnswScalarQuantizedVectorsFormat( + maxConn, + beamWidth, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + true, + random().nextBoolean() + ); } @Override - public void testSortedIndexBytes() throws Exception { - // no bytes + protected KnnVectorsFormat createFormat(int maxConn, int beamWidth, int numMergeWorkers, ExecutorService service) { + return new ES93HnswScalarQuantizedVectorsFormat( + maxConn, + beamWidth, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + true, + random().nextBoolean(), + numMergeWorkers, + service + ); } - @Override - public void testMismatchedFields() throws Exception { - // no bytes - } - - @Override - public void testRandomBytes() throws Exception { - // no bytes - } - - @Override - public void testWriterRamEstimate() throws Exception { - // estimate is different due to bfloat16 - } - - @Override - public void testRandom() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testRandom); - assertFloatsWithinBounds(err); - } - - @Override - public void testRandomWithUpdatesAndGraph() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph); - assertFloatsWithinBounds(err); - } - - @Override - public void testSparseVectors() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors); - assertFloatsWithinBounds(err); - } - - @Override - public void testVectorValuesReportCorrectDocs() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs); - assertFloatsWithinBounds(err); - } - - private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>"); - - private static void assertFloatsWithinBounds(AssertionError error) { - Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage()); - if (m.matches() == false) { - throw error; // nothing to do with us, just rethrow + public void testSimpleOffHeapSize() throws IOException { + float[] vector = randomVector(random().nextInt(12, 500)); + try (Directory dir = newDirectory()) { + testSimpleOffHeapSize( + dir, + newIndexWriterConfig(), + vector, + allOf( + aMapWithSize(3), + hasEntry("vec", (long) vector.length * BFloat16.BYTES), + hasEntry("vex", 1L), + hasEntry(equalTo("veq"), greaterThan(0L)) + ) + ); } - - // numbers just need to be in the same vicinity - double expected = Double.parseDouble(m.group(1)); - double actual = Double.parseDouble(m.group(2)); - double allowedError = expected * 0.01; // within 1% - assertThat(error.getMessage(), actual, closeTo(expected, allowedError)); } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java index a6c12548821b9..2d2564056afca 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java @@ -9,212 +9,73 @@ package org.elasticsearch.index.codec.vectors.es93; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.index.CodecReader; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.KnnVectorValues; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.index.StoredFields; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; import org.apache.lucene.store.Directory; -import org.apache.lucene.store.MMapDirectory; -import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; -import org.apache.lucene.tests.util.TestUtil; -import org.elasticsearch.common.logging.LogConfigurator; -import org.elasticsearch.index.codec.vectors.BFloat16; +import org.elasticsearch.index.codec.vectors.BaseHnswVectorsFormatTestCase; import java.io.IOException; -import java.nio.file.Path; +import java.util.concurrent.ExecutorService; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; -import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.allOf; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.hasEntry; -public class ES93HnswScalarQuantizedVectorsFormatTests extends BaseKnnVectorsFormatTestCase { - - static { - LogConfigurator.loadLog4jPlugins(); - LogConfigurator.configureESLogging(); // native access requires logging to be initialized - } - - boolean useBFloat16() { - return false; - } +public class ES93HnswScalarQuantizedVectorsFormatTests extends BaseHnswVectorsFormatTestCase { @Override - protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat( - new ES93HnswScalarQuantizedVectorsFormat(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, useBFloat16(), null, 7, false, false) + protected KnnVectorsFormat createFormat() { + return new ES93HnswScalarQuantizedVectorsFormat( + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + false, + random().nextBoolean() ); } - // The following test scenarios are similar to their superclass namesakes, - // but here we ensure that the Directory implementation is a FSDirectory - // which helps test the native code vector distance implementation - - public void testAddIndexesDirectory0FS() throws Exception { - Path root = createTempDir(); - String fieldName = "field"; - Document doc = new Document(); - doc.add(new KnnFloatVectorField(fieldName, new float[4], VectorSimilarityFunction.DOT_PRODUCT)); - try (Directory dir = new MMapDirectory(root.resolve("dir1")); Directory dir2 = new MMapDirectory(root.resolve("dir2"))) { - try (IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - w.addDocument(doc); - } - try (IndexWriter w2 = new IndexWriter(dir2, newIndexWriterConfig())) { - w2.addIndexes(dir); - w2.forceMerge(1); - try (IndexReader reader = DirectoryReader.open(w2)) { - LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues vectorValues = r.getFloatVectorValues(fieldName); - KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator(); - assertEquals(0, iterator.nextDoc()); - assertEquals(0, vectorValues.vectorValue(iterator.index())[0], 0); - assertEquals(NO_MORE_DOCS, iterator.nextDoc()); - } - } - } - } - - public void testAddIndexesDirectory01FSCosine() throws Exception { - testAddIndexesDirectory01FS(VectorSimilarityFunction.COSINE); - } - - public void testAddIndexesDirectory01FSDot() throws Exception { - testAddIndexesDirectory01FS(VectorSimilarityFunction.DOT_PRODUCT); - } - - public void testAddIndexesDirectory01FSEuclidean() throws Exception { - testAddIndexesDirectory01FS(VectorSimilarityFunction.EUCLIDEAN); - } - - public void testAddIndexesDirectory01FSMaxIP() throws Exception { - testAddIndexesDirectory01FS(VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT); - } - - private void testAddIndexesDirectory01FS(VectorSimilarityFunction similarityFunction) throws Exception { - Path root = createTempDir(); - String fieldName = "field"; - float[] vector = new float[] { 1f }; - Document doc = new Document(); - doc.add(new KnnFloatVectorField(fieldName, vector, similarityFunction)); - try (Directory dir = new MMapDirectory(root.resolve("dir1")); Directory dir2 = new MMapDirectory(root.resolve("dir2"))) { - try (IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - w.addDocument(doc); - } - try (IndexWriter w2 = new IndexWriter(dir2, newIndexWriterConfig())) { - vector[0] = 2f; - w2.addDocument(doc); - w2.addIndexes(dir); - w2.forceMerge(1); - try (IndexReader reader = DirectoryReader.open(w2)) { - LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues vectorValues = r.getFloatVectorValues(fieldName); - KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator(); - assertEquals(0, iterator.nextDoc()); - // The merge order is randomized, we might get 1 first, or 2 - float value = vectorValues.vectorValue(iterator.index())[0]; - assertTrue(value == 1 || value == 2); - assertEquals(1, iterator.nextDoc()); - value += vectorValues.vectorValue(iterator.index())[0]; - assertEquals(3f, value, 0); - } - } - } - } - - public void testSingleVectorPerSegmentCosine() throws Exception { - testSingleVectorPerSegment(VectorSimilarityFunction.COSINE); - } - - public void testSingleVectorPerSegmentDot() throws Exception { - testSingleVectorPerSegment(VectorSimilarityFunction.DOT_PRODUCT); - } - - public void testSingleVectorPerSegmentEuclidean() throws Exception { - testSingleVectorPerSegment(VectorSimilarityFunction.EUCLIDEAN); - } - - public void testSingleVectorPerSegmentMIP() throws Exception { - testSingleVectorPerSegment(VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT); + @Override + protected KnnVectorsFormat createFormat(int maxConn, int beamWidth) { + return new ES93HnswScalarQuantizedVectorsFormat( + maxConn, + beamWidth, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + false, + random().nextBoolean() + ); } - private void testSingleVectorPerSegment(VectorSimilarityFunction sim) throws Exception { - var codec = getCodec(); - try (Directory dir = new MMapDirectory(createTempDir().resolve("dir1"))) { - try (IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig().setCodec(codec))) { - Document doc2 = new Document(); - doc2.add(new KnnFloatVectorField("field", new float[] { 0.8f, 0.6f }, sim)); - doc2.add(newTextField("id", "A", Field.Store.YES)); - writer.addDocument(doc2); - writer.commit(); - - Document doc1 = new Document(); - doc1.add(new KnnFloatVectorField("field", new float[] { 0.6f, 0.8f }, sim)); - doc1.add(newTextField("id", "B", Field.Store.YES)); - writer.addDocument(doc1); - writer.commit(); - - Document doc3 = new Document(); - doc3.add(new KnnFloatVectorField("field", new float[] { -0.6f, -0.8f }, sim)); - doc3.add(newTextField("id", "C", Field.Store.YES)); - writer.addDocument(doc3); - writer.commit(); - - writer.forceMerge(1); - } - try (DirectoryReader reader = DirectoryReader.open(dir)) { - LeafReader leafReader = getOnlyLeafReader(reader); - StoredFields storedFields = reader.storedFields(); - float[] queryVector = new float[] { 0.6f, 0.8f }; - var hits = leafReader.searchNearestVectors( - "field", - queryVector, - 3, - AcceptDocs.fromLiveDocs(leafReader.getLiveDocs(), leafReader.maxDoc()), - 100 - ); - assertEquals(hits.scoreDocs.length, 3); - assertEquals("B", storedFields.document(hits.scoreDocs[0].doc).get("id")); - assertEquals("A", storedFields.document(hits.scoreDocs[1].doc).get("id")); - assertEquals("C", storedFields.document(hits.scoreDocs[2].doc).get("id")); - } - } + @Override + protected KnnVectorsFormat createFormat(int maxConn, int beamWidth, int numMergeWorkers, ExecutorService service) { + return new ES93HnswScalarQuantizedVectorsFormat( + maxConn, + beamWidth, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, + false, + random().nextBoolean(), + numMergeWorkers, + service + ); } public void testSimpleOffHeapSize() throws IOException { float[] vector = randomVector(random().nextInt(12, 500)); - try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc = new Document(); - doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT)); - w.addDocument(doc); - w.commit(); - try (IndexReader reader = DirectoryReader.open(w)) { - LeafReader r = getOnlyLeafReader(reader); - if (r instanceof CodecReader codecReader) { - KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); - if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { - knnVectorsReader = fieldsReader.getFieldReader("f"); - } - var fieldInfo = r.getFieldInfos().fieldInfo("f"); - var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); - assertEquals(3, offHeap.size()); - int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES; - assertEquals(vector.length * bytes, (long) offHeap.get("vec")); - assertEquals(1L, (long) offHeap.get("vex")); - assertTrue(offHeap.get("veq") > 0L); - } - } + try (Directory dir = newDirectory()) { + testSimpleOffHeapSize( + dir, + newIndexWriterConfig(), + vector, + allOf( + aMapWithSize(3), + hasEntry("vec", (long) vector.length * Float.BYTES), + hasEntry("vex", 1L), + hasEntry(equalTo("veq"), greaterThan(0L)) + ) + ); } } } From 4fa433808db5c25dadb085bae84b160c94f785e3 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 23 Oct 2025 10:20:02 +0100 Subject: [PATCH 08/15] Update more tests --- .../ES93FlatBFloat16VectorFormatTests.java | 134 +++++++---------- .../es93/ES93FlatVectorFormatTests.java | 17 +-- ...uantizedFlatBFloat16VectorFormatTests.java | 138 ++++++++---------- ...ScalarQuantizedFlatVectorsFormatTests.java | 13 +- 4 files changed, 123 insertions(+), 179 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java index 4d017daa2724c..41d777dde2cf9 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java @@ -9,95 +9,65 @@ package org.elasticsearch.index.codec.vectors.es93; -import org.apache.lucene.index.VectorEncoding; - -import java.util.regex.Matcher; -import java.util.regex.Pattern; - -import static org.hamcrest.Matchers.closeTo; - -public class ES93FlatBFloat16VectorFormatTests extends ES93FlatVectorFormatTests { - @Override - boolean useBFloat16() { - return true; - } - - @Override - protected VectorEncoding randomVectorEncoding() { - return VectorEncoding.FLOAT32; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.elasticsearch.index.codec.vectors.BFloat16; +import org.elasticsearch.index.codec.vectors.BaseBFloat16KnnVectorsFormatTestCase; +import org.junit.AssumptionViolatedException; + +import java.io.IOException; + +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.hasEntry; + +public class ES93FlatBFloat16VectorFormatTests extends BaseBFloat16KnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized } @Override - public void testEmptyByteVectorData() throws Exception { - // no bytes + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(true)); } - @Override - public void testMergingWithDifferentByteKnnFields() throws Exception { - // no bytes + public void testSearchWithVisitedLimit() { + throw new AssumptionViolatedException("requires graph-based vector codec"); } - @Override - public void testByteVectorScorerIteration() throws Exception { - // no bytes - } - - @Override - public void testSortedIndexBytes() throws Exception { - // no bytes - } - - @Override - public void testMismatchedFields() throws Exception { - // no bytes - } - - @Override - public void testRandomBytes() throws Exception { - // no bytes - } - - @Override - public void testWriterRamEstimate() throws Exception { - // estimate is different due to bfloat16 - } - - @Override - public void testRandom() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testRandom); - assertFloatsWithinBounds(err); - } - - @Override - public void testRandomWithUpdatesAndGraph() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph); - assertFloatsWithinBounds(err); - } - - @Override - public void testSparseVectors() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors); - assertFloatsWithinBounds(err); - } - - @Override - public void testVectorValuesReportCorrectDocs() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs); - assertFloatsWithinBounds(err); - } - - private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>"); - - private static void assertFloatsWithinBounds(AssertionError error) { - Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage()); - if (m.matches() == false) { - throw error; // nothing to do with us, just rethrow + public void testSimpleOffHeapSize() throws IOException { + float[] vector = randomVector(random().nextInt(12, 500)); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT)); + w.addDocument(doc); + w.commit(); + try (IndexReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnVectorsReader = fieldsReader.getFieldReader("f"); + } + var fieldInfo = r.getFieldInfos().fieldInfo("f"); + var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); + assertThat(offHeap, aMapWithSize(1)); + assertThat(offHeap, hasEntry("vec", (long) vector.length * BFloat16.BYTES)); + } + } } - - // numbers just need to be in the same vicinity - double expected = Double.parseDouble(m.group(1)); - double actual = Double.parseDouble(m.group(2)); - double allowedError = expected * 0.01; // within 1% - assertThat(error.getMessage(), actual, closeTo(expected, allowedError)); } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java index f3faa98124a66..18335758f82d8 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java @@ -23,11 +23,13 @@ import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.util.TestUtil; import org.elasticsearch.common.logging.LogConfigurator; -import org.elasticsearch.index.codec.vectors.BFloat16; +import org.junit.AssumptionViolatedException; import java.io.IOException; import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.hasEntry; public class ES93FlatVectorFormatTests extends BaseKnnVectorsFormatTestCase { @@ -36,17 +38,13 @@ public class ES93FlatVectorFormatTests extends BaseKnnVectorsFormatTestCase { LogConfigurator.configureESLogging(); // native access requires logging to be initialized } - boolean useBFloat16() { - return false; - } - @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(useBFloat16())); + return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(false)); } public void testSearchWithVisitedLimit() { - // requires graph-based vector codec + throw new AssumptionViolatedException("requires graph-based vector codec"); } public void testSimpleOffHeapSize() throws IOException { @@ -65,9 +63,8 @@ public void testSimpleOffHeapSize() throws IOException { } var fieldInfo = r.getFieldInfos().fieldInfo("f"); var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); - int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES; - assertEquals(vector.length * bytes, (long) offHeap.get("vec")); - assertEquals(1, offHeap.size()); + assertThat(offHeap, aMapWithSize(1)); + assertThat(offHeap, hasEntry("vec", (long) vector.length * Float.BYTES)); } } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java index 3681d0a1299eb..bd8376a96f102 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java @@ -9,95 +9,77 @@ package org.elasticsearch.index.codec.vectors.es93; -import org.apache.lucene.index.VectorEncoding; - -import java.util.regex.Matcher; -import java.util.regex.Pattern; - -import static org.hamcrest.Matchers.closeTo; - -public class ES93ScalarQuantizedFlatBFloat16VectorFormatTests extends ES93ScalarQuantizedFlatVectorsFormatTests { - @Override - boolean useBFloat16() { - return true; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.elasticsearch.index.codec.vectors.BFloat16; +import org.elasticsearch.index.codec.vectors.BaseBFloat16KnnVectorsFormatTestCase; +import org.junit.AssumptionViolatedException; + +import java.io.IOException; + +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.hasEntry; + +public class ES93ScalarQuantizedFlatBFloat16VectorFormatTests extends BaseBFloat16KnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized } - @Override - protected VectorEncoding randomVectorEncoding() { - return VectorEncoding.FLOAT32; - } + private KnnVectorsFormat format; @Override - public void testEmptyByteVectorData() throws Exception { - // no bytes + public void setUp() throws Exception { + format = new ES93ScalarQuantizedFlatVectorsFormat(true); + super.setUp(); } @Override - public void testMergingWithDifferentByteKnnFields() throws Exception { - // no bytes + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(format); } - @Override - public void testByteVectorScorerIteration() throws Exception { - // no bytes + public void testSearchWithVisitedLimit() { + throw new AssumptionViolatedException("requires graph vector codec"); } - @Override - public void testSortedIndexBytes() throws Exception { - // no bytes - } - - @Override - public void testMismatchedFields() throws Exception { - // no bytes - } - - @Override - public void testRandomBytes() throws Exception { - // no bytes - } - - @Override - public void testWriterRamEstimate() throws Exception { - // estimate is different due to bfloat16 - } - - @Override - public void testRandom() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testRandom); - assertFloatsWithinBounds(err); - } - - @Override - public void testRandomWithUpdatesAndGraph() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph); - assertFloatsWithinBounds(err); - } - - @Override - public void testSparseVectors() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors); - assertFloatsWithinBounds(err); - } - - @Override - public void testVectorValuesReportCorrectDocs() throws Exception { - AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs); - assertFloatsWithinBounds(err); - } - - private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>"); - - private static void assertFloatsWithinBounds(AssertionError error) { - Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage()); - if (m.matches() == false) { - throw error; // nothing to do with us, just rethrow + public void testSimpleOffHeapSize() throws IOException { + float[] vector = randomVector(random().nextInt(12, 500)); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, DOT_PRODUCT)); + w.addDocument(doc); + w.commit(); + try (IndexReader reader = DirectoryReader.open(w)) { + LeafReader r = getOnlyLeafReader(reader); + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnVectorsReader = fieldsReader.getFieldReader("f"); + } + var fieldInfo = r.getFieldInfos().fieldInfo("f"); + var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); + assertThat(offHeap, aMapWithSize(2)); + assertThat(offHeap, hasEntry("vec", (long) vector.length * BFloat16.BYTES)); + assertThat(offHeap, hasEntry(equalTo("veq"), greaterThan(0L))); + } + } } - - // numbers just need to be in the same vicinity - double expected = Double.parseDouble(m.group(1)); - double actual = Double.parseDouble(m.group(2)); - double allowedError = expected * 0.01; // within 1% - assertThat(error.getMessage(), actual, closeTo(expected, allowedError)); } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java index d05f33941f7b3..bb56b42d14400 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java @@ -23,7 +23,7 @@ import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.util.TestUtil; import org.elasticsearch.common.logging.LogConfigurator; -import org.elasticsearch.index.codec.vectors.BFloat16; +import org.junit.AssumptionViolatedException; import java.io.IOException; @@ -40,17 +40,13 @@ public class ES93ScalarQuantizedFlatVectorsFormatTests extends BaseKnnVectorsFor LogConfigurator.configureESLogging(); // native access requires logging to be initialized } - boolean useBFloat16() { - return false; - } - @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedFlatVectorsFormat(useBFloat16())); + return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedFlatVectorsFormat(false)); } public void testSearchWithVisitedLimit() { - // requires graph vector codec + throw new AssumptionViolatedException("requires graph vector codec"); } public void testSimpleOffHeapSize() throws IOException { @@ -70,11 +66,10 @@ public void testSimpleOffHeapSize() throws IOException { var fieldInfo = r.getFieldInfos().fieldInfo("f"); var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); assertThat(offHeap, aMapWithSize(2)); - assertThat(offHeap, hasEntry("vec", vector.length * (useBFloat16() ? BFloat16.BYTES : Float.BYTES))); + assertThat(offHeap, hasEntry("vec", (long) vector.length * Float.BYTES)); assertThat(offHeap, hasEntry(equalTo("veq"), greaterThan(0L))); } } } } - } From 4815bc2a373dadebc96f7a828bef6187155015fb Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 23 Oct 2025 10:53:16 +0100 Subject: [PATCH 09/15] Class renames --- server/src/main/java/module-info.java | 4 +- .../ES93HnswScalarQuantizedVectorsFormat.java | 8 +- .../ES93ScalarQuantizedFlatVectorsFormat.java | 125 +++++------------- .../ES93ScalarQuantizedVectorsFormat.java | 125 +++++++++++++----- .../org.apache.lucene.codecs.KnnVectorsFormat | 4 +- ...arQuantizedBFloat16VectorFormatTests.java} | 4 +- ...S93ScalarQuantizedVectorsFormatTests.java} | 4 +- 7 files changed, 137 insertions(+), 137 deletions(-) rename server/src/test/java/org/elasticsearch/index/codec/vectors/es93/{ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java => ES93ScalarQuantizedBFloat16VectorFormatTests.java} (95%) rename server/src/test/java/org/elasticsearch/index/codec/vectors/es93/{ES93ScalarQuantizedFlatVectorsFormatTests.java => ES93ScalarQuantizedVectorsFormatTests.java} (96%) diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index c56db1f377315..9ec9b42d34297 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -466,10 +466,10 @@ org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat, - org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedFlatVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, - org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; provides org.apache.lucene.codecs.Codec diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java index a2e14df6099ec..aec2120b97ace 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -31,7 +31,7 @@ public class ES93HnswScalarQuantizedVectorsFormat extends AbstractHnswVectorsFor public ES93HnswScalarQuantizedVectorsFormat() { super(NAME); - this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat( + this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat( Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, false, false @@ -46,7 +46,7 @@ public ES93HnswScalarQuantizedVectorsFormat( boolean useDirectIO ) { super(NAME, maxConn, beamWidth); - this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat(encoding, useBFloat16, useDirectIO); + this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, useBFloat16, useDirectIO); } public ES93HnswScalarQuantizedVectorsFormat( @@ -59,7 +59,7 @@ public ES93HnswScalarQuantizedVectorsFormat( ExecutorService mergeExec ) { super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec); - this.flatVectorsFormat = new ES93ScalarQuantizedVectorsFormat(encoding, useBFloat16, useDirectIO); + this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, useBFloat16, useDirectIO); } @Override @@ -69,7 +69,7 @@ protected FlatVectorsFormat flatVectorsFormat() { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new Lucene99HnswVectorsWriter(state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), numMergeWorkers, mergeExec); + return new Lucene99HnswVectorsWriter(state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), numMergeWorkers, mergeExec, 0); } @Override diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java index f1375cb5469e7..c4b99f75ee910 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java @@ -9,55 +9,40 @@ package org.elasticsearch.index.codec.vectors.es93; -import org.apache.lucene.codecs.KnnVectorsFormat; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.KnnVectorsWriter; +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.lucene104.Lucene104ScalarQuantizedVectorScorer; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; -import org.apache.lucene.index.ByteVectorValues; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.search.AcceptDocs; -import org.apache.lucene.search.KnnCollector; -import org.apache.lucene.util.Bits; -import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; -import org.apache.lucene.util.hnsw.RandomVectorScorer; import java.io.IOException; -import java.util.Map; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; -public class ES93ScalarQuantizedFlatVectorsFormat extends KnnVectorsFormat { +public class ES93ScalarQuantizedFlatVectorsFormat extends FlatVectorsFormat { static final String NAME = "ES93ScalarQuantizedFlatVectorsFormat"; - private final FlatVectorsFormat format; + static final Lucene104ScalarQuantizedVectorScorer flatVectorScorer = new Lucene104ScalarQuantizedVectorScorer( + FlatVectorScorerUtil.getLucene99FlatVectorsScorer() + ); - public ES93ScalarQuantizedFlatVectorsFormat() { - this(false, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); - } - - public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16) { - this(useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); - } + private final Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding; + private final FlatVectorsFormat rawVectorFormat; - public ES93ScalarQuantizedFlatVectorsFormat(boolean useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding) { + public ES93ScalarQuantizedFlatVectorsFormat( + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, + boolean useBFloat16, + boolean useDirectIO + ) { super(NAME); - this.format = new ES93ScalarQuantizedVectorsFormat(encoding, useBFloat16, false); - } - - @Override - public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return format.fieldsWriter(state); - } - - @Override - public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new ES93FlatVectorsReader(format.fieldsReader(state)); + this.encoding = encoding; + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(useBFloat16, useDirectIO); } @Override @@ -67,64 +52,26 @@ public int getMaxDimensions(String fieldName) { @Override public String toString() { - return NAME + "(name=" + NAME + ", innerFormat=" + format + ")"; + return NAME + + "(name=" + + NAME + + ", encoding=" + + encoding + + ", flatVectorScorer=" + + flatVectorScorer + + ", rawVectorFormat=" + + rawVectorFormat + + ")"; } - public static class ES93FlatVectorsReader extends KnnVectorsReader { - - private final FlatVectorsReader reader; - - public ES93FlatVectorsReader(FlatVectorsReader reader) { - super(); - this.reader = reader; - } - - @Override - public void checkIntegrity() throws IOException { - reader.checkIntegrity(); - } - - @Override - public FloatVectorValues getFloatVectorValues(String field) throws IOException { - return reader.getFloatVectorValues(field); - } - - @Override - public ByteVectorValues getByteVectorValues(String field) throws IOException { - return reader.getByteVectorValues(field); - } - - @Override - public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { - collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); - } - - private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) - throws IOException { - OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); - Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); - for (int i = 0; i < scorer.maxOrd(); i++) { - if (acceptedOrds == null || acceptedOrds.get(i)) { - collector.collect(i, scorer.score(i)); - collector.incVisitedCount(1); - } - } - assert collector.earlyTerminated() == false; - } - - @Override - public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { - collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); - } - - @Override - public Map getOffHeapByteSize(FieldInfo fieldInfo) { - return reader.getOffHeapByteSize(fieldInfo); - } + @Override + public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { + }; + } - @Override - public void close() throws IOException { - reader.close(); - } + @Override + public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); } } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java index b393d91dd0c72..175b4a79ecce5 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -9,69 +9,122 @@ package org.elasticsearch.index.codec.vectors.es93; -import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; 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.lucene104.Lucene104ScalarQuantizedVectorScorer; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; -import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; -import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; +import org.apache.lucene.util.hnsw.RandomVectorScorer; import java.io.IOException; +import java.util.Map; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; -public class ES93ScalarQuantizedVectorsFormat extends FlatVectorsFormat { +public class ES93ScalarQuantizedVectorsFormat extends KnnVectorsFormat { static final String NAME = "ES93ScalarQuantizedVectorsFormat"; - static final Lucene104ScalarQuantizedVectorScorer flatVectorScorer = new Lucene104ScalarQuantizedVectorScorer( - FlatVectorScorerUtil.getLucene99FlatVectorsScorer() - ); + private final FlatVectorsFormat format; - private final Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding; - private final FlatVectorsFormat rawVectorFormat; + public ES93ScalarQuantizedVectorsFormat() { + this(false, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); + } + + public ES93ScalarQuantizedVectorsFormat(boolean useBFloat16) { + this(useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); + } - public ES93ScalarQuantizedVectorsFormat( - Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, - boolean useBFloat16, - boolean useDirectIO - ) { + public ES93ScalarQuantizedVectorsFormat(boolean useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding) { super(NAME); - this.encoding = encoding; - this.rawVectorFormat = new ES93GenericFlatVectorsFormat(useBFloat16, useDirectIO); + this.format = new ES93ScalarQuantizedFlatVectorsFormat(encoding, useBFloat16, false); } @Override - public int getMaxDimensions(String fieldName) { - return MAX_DIMS_COUNT; + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return format.fieldsWriter(state); } @Override - public String toString() { - return NAME - + "(name=" - + NAME - + ", encoding=" - + encoding - + ", flatVectorScorer=" - + flatVectorScorer - + ", rawVectorFormat=" - + rawVectorFormat - + ")"; + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new ES93FlatVectorsReader(format.fieldsReader(state)); } @Override - public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { - }; + public int getMaxDimensions(String fieldName) { + return MAX_DIMS_COUNT; } @Override - public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); + public String toString() { + return NAME + "(name=" + NAME + ", innerFormat=" + format + ")"; + } + + public static class ES93FlatVectorsReader extends KnnVectorsReader { + + private final FlatVectorsReader reader; + + public ES93FlatVectorsReader(FlatVectorsReader reader) { + super(); + this.reader = reader; + } + + @Override + public void checkIntegrity() throws IOException { + reader.checkIntegrity(); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return reader.getFloatVectorValues(field); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return reader.getByteVectorValues(field); + } + + @Override + public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) + throws IOException { + OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); + Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); + for (int i = 0; i < scorer.maxOrd(); i++) { + if (acceptedOrds == null || acceptedOrds.get(i)) { + collector.collect(i, scorer.score(i)); + collector.incVisitedCount(1); + } + } + assert collector.earlyTerminated() == false; + } + + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + return reader.getOffHeapByteSize(fieldInfo); + } + + @Override + public void close() throws IOException { + reader.close(); + } } } diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index d6b114fb57555..0dc34ea2e808d 100644 --- a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -10,8 +10,8 @@ org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsForma org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat -org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedFlatVectorsFormat +org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat +org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat -org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java similarity index 95% rename from server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java rename to server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java index bd8376a96f102..7e69492f6e848 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatBFloat16VectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java @@ -35,7 +35,7 @@ import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.hasEntry; -public class ES93ScalarQuantizedFlatBFloat16VectorFormatTests extends BaseBFloat16KnnVectorsFormatTestCase { +public class ES93ScalarQuantizedBFloat16VectorFormatTests extends BaseBFloat16KnnVectorsFormatTestCase { static { LogConfigurator.loadLog4jPlugins(); @@ -46,7 +46,7 @@ public class ES93ScalarQuantizedFlatBFloat16VectorFormatTests extends BaseBFloat @Override public void setUp() throws Exception { - format = new ES93ScalarQuantizedFlatVectorsFormat(true); + format = new ES93ScalarQuantizedVectorsFormat(true); super.setUp(); } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java similarity index 96% rename from server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java rename to server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java index bb56b42d14400..7c8461a6d033d 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java @@ -33,7 +33,7 @@ import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.hasEntry; -public class ES93ScalarQuantizedFlatVectorsFormatTests extends BaseKnnVectorsFormatTestCase { +public class ES93ScalarQuantizedVectorsFormatTests extends BaseKnnVectorsFormatTestCase { static { LogConfigurator.loadLog4jPlugins(); @@ -42,7 +42,7 @@ public class ES93ScalarQuantizedFlatVectorsFormatTests extends BaseKnnVectorsFor @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedFlatVectorsFormat(false)); + return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedVectorsFormat(false)); } public void testSearchWithVisitedLimit() { From 94ef4f1768df6283f370d1457e78b1838e0f4286 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Thu, 23 Oct 2025 13:10:19 +0000 Subject: [PATCH 10/15] [CI] Auto commit changes from spotless --- .../es93/ES93HnswScalarQuantizedVectorsFormat.java | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java index aec2120b97ace..9bd1ff1fbc113 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -69,7 +69,15 @@ protected FlatVectorsFormat flatVectorsFormat() { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new Lucene99HnswVectorsWriter(state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), numMergeWorkers, mergeExec, 0); + return new Lucene99HnswVectorsWriter( + state, + maxConn, + beamWidth, + flatVectorsFormat.fieldsWriter(state), + numMergeWorkers, + mergeExec, + 0 + ); } @Override From 67203847e04d69c436b4d16dd94a6f591177bc24 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Thu, 23 Oct 2025 16:08:20 +0100 Subject: [PATCH 11/15] Update for ElementType change --- .../codec/vectors/es93/ES93FlatVectorFormat.java | 5 +++-- .../es93/ES93HnswScalarQuantizedVectorsFormat.java | 10 +++++----- .../es93/ES93ScalarQuantizedFlatVectorsFormat.java | 5 +++-- .../es93/ES93ScalarQuantizedVectorsFormat.java | 13 ++++++++----- .../es93/ES93FlatBFloat16VectorFormatTests.java | 2 +- .../vectors/es93/ES93FlatVectorFormatTests.java | 2 +- ...swScalarQuantizedBFloat16VectorsFormatTests.java | 6 +++--- .../ES93HnswScalarQuantizedVectorsFormatTests.java | 6 +++--- ...S93ScalarQuantizedBFloat16VectorFormatTests.java | 2 +- .../es93/ES93ScalarQuantizedVectorsFormatTests.java | 2 +- 10 files changed, 29 insertions(+), 24 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java index bd421abf167f0..bdad21596d479 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java @@ -44,9 +44,10 @@ public ES93FlatVectorFormat() { format = new ES93GenericFlatVectorsFormat(); } - public ES93FlatVectorFormat(boolean useBFloat16) { + public ES93FlatVectorFormat(ES93GenericFlatVectorsFormat.ElementType elementType) { super(NAME); - format = new ES93GenericFlatVectorsFormat(useBFloat16, false); + assert elementType != ES93GenericFlatVectorsFormat.ElementType.BIT : "ES815BitFlatVectorFormat should be used for bits"; + format = new ES93GenericFlatVectorsFormat(elementType, false); } @Override diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java index 9bd1ff1fbc113..f22b5d5f14c64 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -33,7 +33,7 @@ public ES93HnswScalarQuantizedVectorsFormat() { super(NAME); this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat( Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - false, + ES93GenericFlatVectorsFormat.ElementType.STANDARD, false ); } @@ -42,24 +42,24 @@ public ES93HnswScalarQuantizedVectorsFormat( int maxConn, int beamWidth, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, - boolean useBFloat16, + ES93GenericFlatVectorsFormat.ElementType elementType, boolean useDirectIO ) { super(NAME, maxConn, beamWidth); - this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, useBFloat16, useDirectIO); + this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, elementType, useDirectIO); } public ES93HnswScalarQuantizedVectorsFormat( int maxConn, int beamWidth, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, - boolean useBFloat16, + ES93GenericFlatVectorsFormat.ElementType elementType, boolean useDirectIO, int numMergeWorkers, ExecutorService mergeExec ) { super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec); - this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, useBFloat16, useDirectIO); + this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, elementType, useDirectIO); } @Override diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java index c4b99f75ee910..18967a39fe8b3 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java @@ -37,12 +37,13 @@ public class ES93ScalarQuantizedFlatVectorsFormat extends FlatVectorsFormat { public ES93ScalarQuantizedFlatVectorsFormat( Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, - boolean useBFloat16, + ES93GenericFlatVectorsFormat.ElementType elementType, boolean useDirectIO ) { super(NAME); + assert elementType != ES93GenericFlatVectorsFormat.ElementType.BIT : "BIT should not be used with scalar quantization"; this.encoding = encoding; - this.rawVectorFormat = new ES93GenericFlatVectorsFormat(useBFloat16, useDirectIO); + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, useDirectIO); } @Override diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java index 175b4a79ecce5..647144cfaa2de 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -38,16 +38,19 @@ public class ES93ScalarQuantizedVectorsFormat extends KnnVectorsFormat { private final FlatVectorsFormat format; public ES93ScalarQuantizedVectorsFormat() { - this(false, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); + this(ES93GenericFlatVectorsFormat.ElementType.STANDARD, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); } - public ES93ScalarQuantizedVectorsFormat(boolean useBFloat16) { - this(useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); + public ES93ScalarQuantizedVectorsFormat(ES93GenericFlatVectorsFormat.ElementType elementType) { + this(elementType, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); } - public ES93ScalarQuantizedVectorsFormat(boolean useBFloat16, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding) { + public ES93ScalarQuantizedVectorsFormat( + ES93GenericFlatVectorsFormat.ElementType elementType, + Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding + ) { super(NAME); - this.format = new ES93ScalarQuantizedFlatVectorsFormat(encoding, useBFloat16, false); + this.format = new ES93ScalarQuantizedFlatVectorsFormat(encoding, elementType, false); } @Override diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java index 41d777dde2cf9..91d4054ae94ed 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatBFloat16VectorFormatTests.java @@ -41,7 +41,7 @@ public class ES93FlatBFloat16VectorFormatTests extends BaseBFloat16KnnVectorsFor @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(true)); + return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(ES93GenericFlatVectorsFormat.ElementType.BFLOAT16)); } public void testSearchWithVisitedLimit() { diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java index 18335758f82d8..1ada03a70bed6 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormatTests.java @@ -40,7 +40,7 @@ public class ES93FlatVectorFormatTests extends BaseKnnVectorsFormatTestCase { @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(false)); + return TestUtil.alwaysKnnVectorsFormat(new ES93FlatVectorFormat(ES93GenericFlatVectorsFormat.ElementType.STANDARD)); } public void testSearchWithVisitedLimit() { diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java index e2bc3970398c5..a1bda3e4b2342 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedBFloat16VectorsFormatTests.java @@ -34,7 +34,7 @@ protected KnnVectorsFormat createFormat() { DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - true, + ES93GenericFlatVectorsFormat.ElementType.BFLOAT16, random().nextBoolean() ); } @@ -45,7 +45,7 @@ protected KnnVectorsFormat createFormat(int maxConn, int beamWidth) { maxConn, beamWidth, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - true, + ES93GenericFlatVectorsFormat.ElementType.BFLOAT16, random().nextBoolean() ); } @@ -56,7 +56,7 @@ protected KnnVectorsFormat createFormat(int maxConn, int beamWidth, int numMerge maxConn, beamWidth, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - true, + ES93GenericFlatVectorsFormat.ElementType.BFLOAT16, random().nextBoolean(), numMergeWorkers, service diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java index 2d2564056afca..c2bf9e6352f15 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormatTests.java @@ -33,7 +33,7 @@ protected KnnVectorsFormat createFormat() { DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - false, + ES93GenericFlatVectorsFormat.ElementType.STANDARD, random().nextBoolean() ); } @@ -44,7 +44,7 @@ protected KnnVectorsFormat createFormat(int maxConn, int beamWidth) { maxConn, beamWidth, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - false, + ES93GenericFlatVectorsFormat.ElementType.STANDARD, random().nextBoolean() ); } @@ -55,7 +55,7 @@ protected KnnVectorsFormat createFormat(int maxConn, int beamWidth, int numMerge maxConn, beamWidth, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - false, + ES93GenericFlatVectorsFormat.ElementType.STANDARD, random().nextBoolean(), numMergeWorkers, service diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java index 7e69492f6e848..57578097f2db4 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedBFloat16VectorFormatTests.java @@ -46,7 +46,7 @@ public class ES93ScalarQuantizedBFloat16VectorFormatTests extends BaseBFloat16Kn @Override public void setUp() throws Exception { - format = new ES93ScalarQuantizedVectorsFormat(true); + format = new ES93ScalarQuantizedVectorsFormat(ES93GenericFlatVectorsFormat.ElementType.BFLOAT16); super.setUp(); } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java index 7c8461a6d033d..a880852378d61 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormatTests.java @@ -42,7 +42,7 @@ public class ES93ScalarQuantizedVectorsFormatTests extends BaseKnnVectorsFormatT @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedVectorsFormat(false)); + return TestUtil.alwaysKnnVectorsFormat(new ES93ScalarQuantizedVectorsFormat(ES93GenericFlatVectorsFormat.ElementType.STANDARD)); } public void testSearchWithVisitedLimit() { From 115507aaca47a9ee65c3c63640339ec658a35189 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Fri, 24 Oct 2025 09:23:19 +0100 Subject: [PATCH 12/15] Revert "Use the reader in Lucene BWC" This reverts commit 064392fbf29afd634f54616b5a8db5bc7a71e96f. --- .../ES814ScalarQuantizedVectorsFormat.java | 1 - .../Lucene99ScalarQuantizedVectorsReader.java | 455 ++++++++++++++++++ 2 files changed, 455 insertions(+), 1 deletion(-) create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/Lucene99ScalarQuantizedVectorsReader.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 dd125b68d292a..692ccdfa9222c 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 @@ -9,7 +9,6 @@ package org.elasticsearch.index.codec.vectors; -import org.apache.lucene.backward_codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; 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(); + } + } +} From 44ecd39edd0b8be4df65bfd8aaf34b56560db971 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Fri, 24 Oct 2025 10:01:34 +0100 Subject: [PATCH 13/15] Remove intermediate class --- .../ES93HnswScalarQuantizedVectorsFormat.java | 35 ++++++--- .../ES93ScalarQuantizedFlatVectorsFormat.java | 78 ------------------- .../ES93ScalarQuantizedVectorsFormat.java | 33 ++++++-- 3 files changed, 51 insertions(+), 95 deletions(-) delete mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java index f22b5d5f14c64..6bd225123d425 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -11,8 +11,12 @@ import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; import org.apache.lucene.index.SegmentReadState; @@ -26,16 +30,17 @@ public class ES93HnswScalarQuantizedVectorsFormat extends AbstractHnswVectorsFor static final String NAME = "ES93HnswScalarQuantizedVectorsFormat"; - /** The format for storing, reading, merging vectors on disk */ - private final FlatVectorsFormat flatVectorsFormat; + static final Lucene104ScalarQuantizedVectorScorer flatVectorScorer = new Lucene104ScalarQuantizedVectorScorer( + FlatVectorScorerUtil.getLucene99FlatVectorsScorer() + ); + + private final Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding; + private final FlatVectorsFormat rawVectorFormat; public ES93HnswScalarQuantizedVectorsFormat() { super(NAME); - this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat( - Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT, - ES93GenericFlatVectorsFormat.ElementType.STANDARD, - false - ); + this.encoding = Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT; + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(ES93GenericFlatVectorsFormat.ElementType.STANDARD, false); } public ES93HnswScalarQuantizedVectorsFormat( @@ -46,7 +51,8 @@ public ES93HnswScalarQuantizedVectorsFormat( boolean useDirectIO ) { super(NAME, maxConn, beamWidth); - this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, elementType, useDirectIO); + this.encoding = encoding; + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, useDirectIO); } public ES93HnswScalarQuantizedVectorsFormat( @@ -59,12 +65,13 @@ public ES93HnswScalarQuantizedVectorsFormat( ExecutorService mergeExec ) { super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec); - this.flatVectorsFormat = new ES93ScalarQuantizedFlatVectorsFormat(encoding, elementType, useDirectIO); + this.encoding = encoding; + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, useDirectIO); } @Override protected FlatVectorsFormat flatVectorsFormat() { - return flatVectorsFormat; + return rawVectorFormat; } @Override @@ -73,7 +80,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, maxConn, beamWidth, - flatVectorsFormat.fieldsWriter(state), + new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { + }, numMergeWorkers, mergeExec, 0 @@ -82,6 +90,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state)); + return new Lucene99HnswVectorsReader( + state, + new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer) + ); } } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java deleted file mode 100644 index 18967a39fe8b3..0000000000000 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedFlatVectorsFormat.java +++ /dev/null @@ -1,78 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the "Elastic License - * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side - * Public License v 1"; you may not use this file except in compliance with, at - * your election, the "Elastic License 2.0", the "GNU Affero General Public - * License v3.0 only", or the "Server Side Public License, v 1". - */ - -package org.elasticsearch.index.codec.vectors.es93; - -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.lucene104.Lucene104ScalarQuantizedVectorScorer; -import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; -import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; -import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; -import org.apache.lucene.index.SegmentReadState; -import org.apache.lucene.index.SegmentWriteState; - -import java.io.IOException; - -import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; - -public class ES93ScalarQuantizedFlatVectorsFormat extends FlatVectorsFormat { - - static final String NAME = "ES93ScalarQuantizedFlatVectorsFormat"; - - static final Lucene104ScalarQuantizedVectorScorer flatVectorScorer = new Lucene104ScalarQuantizedVectorScorer( - FlatVectorScorerUtil.getLucene99FlatVectorsScorer() - ); - - private final Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding; - private final FlatVectorsFormat rawVectorFormat; - - public ES93ScalarQuantizedFlatVectorsFormat( - Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding, - ES93GenericFlatVectorsFormat.ElementType elementType, - boolean useDirectIO - ) { - super(NAME); - assert elementType != ES93GenericFlatVectorsFormat.ElementType.BIT : "BIT should not be used with scalar quantization"; - this.encoding = encoding; - this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, useDirectIO); - } - - @Override - public int getMaxDimensions(String fieldName) { - return MAX_DIMS_COUNT; - } - - @Override - public String toString() { - return NAME - + "(name=" - + NAME - + ", encoding=" - + encoding - + ", flatVectorScorer=" - + flatVectorScorer - + ", rawVectorFormat=" - + rawVectorFormat - + ")"; - } - - @Override - public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { - }; - } - - @Override - public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); - } -} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java index 647144cfaa2de..0d08f738f65a7 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -12,9 +12,13 @@ import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; +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.lucene104.Lucene104ScalarQuantizedVectorScorer; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FloatVectorValues; @@ -35,7 +39,12 @@ public class ES93ScalarQuantizedVectorsFormat extends KnnVectorsFormat { static final String NAME = "ES93ScalarQuantizedVectorsFormat"; - private final FlatVectorsFormat format; + static final Lucene104ScalarQuantizedVectorScorer flatVectorScorer = new Lucene104ScalarQuantizedVectorScorer( + FlatVectorScorerUtil.getLucene99FlatVectorsScorer() + ); + + private final Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding; + private final FlatVectorsFormat rawVectorFormat; public ES93ScalarQuantizedVectorsFormat() { this(ES93GenericFlatVectorsFormat.ElementType.STANDARD, Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding.SEVEN_BIT); @@ -50,17 +59,22 @@ public ES93ScalarQuantizedVectorsFormat( Lucene104ScalarQuantizedVectorsFormat.ScalarEncoding encoding ) { super(NAME); - this.format = new ES93ScalarQuantizedFlatVectorsFormat(encoding, elementType, false); + assert elementType != ES93GenericFlatVectorsFormat.ElementType.BIT : "BIT should not be used with scalar quantization"; + this.encoding = encoding; + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, false); } @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return format.fieldsWriter(state); + return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { + }; } @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new ES93FlatVectorsReader(format.fieldsReader(state)); + return new ES93FlatVectorsReader( + new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer) + ); } @Override @@ -70,7 +84,16 @@ public int getMaxDimensions(String fieldName) { @Override public String toString() { - return NAME + "(name=" + NAME + ", innerFormat=" + format + ")"; + return NAME + + "(name=" + + NAME + + ", encoding=" + + encoding + + ", flatVectorScorer=" + + flatVectorScorer + + ", rawVectorFormat=" + + rawVectorFormat + + ")"; } public static class ES93FlatVectorsReader extends KnnVectorsReader { From f6ee76915a1b7f0c035774b455fe567f0acd1065 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Fri, 31 Oct 2025 12:38:07 +0000 Subject: [PATCH 14/15] Use public constructor --- .../vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java | 3 +-- .../codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java index 6bd225123d425..4f47b82c3b5a6 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -80,8 +80,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, maxConn, beamWidth, - new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { - }, + new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer), numMergeWorkers, mergeExec, 0 diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java index 0d08f738f65a7..90cb3aec59ed8 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -66,8 +66,7 @@ public ES93ScalarQuantizedVectorsFormat( @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer) { - }; + return new Lucene104ScalarQuantizedVectorsWriter(state, encoding, rawVectorFormat.fieldsWriter(state), flatVectorScorer); } @Override From 9795d0f15bc9151c6b49c8c66711701201acd078 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Mon, 3 Nov 2025 15:43:02 +0000 Subject: [PATCH 15/15] We don't need a separate search impl here --- .../ES93ScalarQuantizedVectorsFormat.java | 72 +------------------ 1 file changed, 1 insertion(+), 71 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java index 90cb3aec59ed8..075c1728f1029 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -14,24 +14,14 @@ import org.apache.lucene.codecs.KnnVectorsWriter; 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.lucene104.Lucene104ScalarQuantizedVectorScorer; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsReader; import org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsWriter; -import org.apache.lucene.index.ByteVectorValues; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.search.AcceptDocs; -import org.apache.lucene.search.KnnCollector; -import org.apache.lucene.util.Bits; -import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; -import org.apache.lucene.util.hnsw.RandomVectorScorer; import java.io.IOException; -import java.util.Map; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; @@ -71,9 +61,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new ES93FlatVectorsReader( - new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer) - ); + return new Lucene104ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer); } @Override @@ -94,62 +82,4 @@ public String toString() { + rawVectorFormat + ")"; } - - public static class ES93FlatVectorsReader extends KnnVectorsReader { - - private final FlatVectorsReader reader; - - public ES93FlatVectorsReader(FlatVectorsReader reader) { - super(); - this.reader = reader; - } - - @Override - public void checkIntegrity() throws IOException { - reader.checkIntegrity(); - } - - @Override - public FloatVectorValues getFloatVectorValues(String field) throws IOException { - return reader.getFloatVectorValues(field); - } - - @Override - public ByteVectorValues getByteVectorValues(String field) throws IOException { - return reader.getByteVectorValues(field); - } - - @Override - public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { - collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); - } - - private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) - throws IOException { - OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); - Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); - for (int i = 0; i < scorer.maxOrd(); i++) { - if (acceptedOrds == null || acceptedOrds.get(i)) { - collector.collect(i, scorer.score(i)); - collector.incVisitedCount(1); - } - } - assert collector.earlyTerminated() == false; - } - - @Override - public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { - collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); - } - - @Override - public Map getOffHeapByteSize(FieldInfo fieldInfo) { - return reader.getOffHeapByteSize(fieldInfo); - } - - @Override - public void close() throws IOException { - reader.close(); - } - } }