Create int bfloat16 vector formats

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;
 }

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<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
+            return reader.getOffHeapByteSize(fieldInfo);
+        }
+
+        @Override
+        public void close() throws IOException {
+            reader.close();
+        }
+    }
+}

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

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));
+    }
+}

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);
+                }
+            }
+        }
+    }
+
+}