Merge branch 'bfloat16-vector-format' into bbq_hnsw-bfloat16

Author: thecoop

qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java

@@ -49,6 +49,7 @@ record CmdLineArgs(
     float filterSelectivity,
     long seed,
     VectorSimilarityFunction vectorSpace,
+    int rawVectorSize,
     int quantizeBits,
     VectorEncoding vectorEncoding,
     int dimensions,
@@ -76,6 +77,7 @@ record CmdLineArgs(
     static final ParseField FORCE_MERGE_FIELD = new ParseField("force_merge");
     static final ParseField VECTOR_SPACE_FIELD = new ParseField("vector_space");
     static final ParseField QUANTIZE_BITS_FIELD = new ParseField("quantize_bits");
+    static final ParseField RAW_VECTOR_SIZE_FIELD = new ParseField("raw_vector_size");
     static final ParseField VECTOR_ENCODING_FIELD = new ParseField("vector_encoding");
     static final ParseField DIMENSIONS_FIELD = new ParseField("dimensions");
     static final ParseField EARLY_TERMINATION_FIELD = new ParseField("early_termination");
@@ -110,6 +112,7 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareBoolean(Builder::setReindex, REINDEX_FIELD);
         PARSER.declareBoolean(Builder::setForceMerge, FORCE_MERGE_FIELD);
         PARSER.declareString(Builder::setVectorSpace, VECTOR_SPACE_FIELD);
+        PARSER.declareInt(Builder::setRawVectorSize, RAW_VECTOR_SIZE_FIELD);
         PARSER.declareInt(Builder::setQuantizeBits, QUANTIZE_BITS_FIELD);
         PARSER.declareString(Builder::setVectorEncoding, VECTOR_ENCODING_FIELD);
         PARSER.declareInt(Builder::setDimensions, DIMENSIONS_FIELD);
@@ -146,6 +149,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(REINDEX_FIELD.getPreferredName(), reindex);
         builder.field(FORCE_MERGE_FIELD.getPreferredName(), forceMerge);
         builder.field(VECTOR_SPACE_FIELD.getPreferredName(), vectorSpace.name().toLowerCase(Locale.ROOT));
+        builder.field(RAW_VECTOR_SIZE_FIELD.getPreferredName(), rawVectorSize);
         builder.field(QUANTIZE_BITS_FIELD.getPreferredName(), quantizeBits);
         builder.field(VECTOR_ENCODING_FIELD.getPreferredName(), vectorEncoding.name().toLowerCase(Locale.ROOT));
         builder.field(DIMENSIONS_FIELD.getPreferredName(), dimensions);
@@ -179,6 +183,7 @@ static class Builder {
         private boolean reindex = false;
         private boolean forceMerge = false;
         private VectorSimilarityFunction vectorSpace = VectorSimilarityFunction.EUCLIDEAN;
+        private int rawVectorSize = 32;
         private int quantizeBits = 8;
         private VectorEncoding vectorEncoding = VectorEncoding.FLOAT32;
         private int dimensions;
@@ -281,6 +286,11 @@ public Builder setVectorSpace(String vectorSpace) {
             return this;
         }
 
+        public Builder setRawVectorSize(int rawVectorSize) {
+            this.rawVectorSize = rawVectorSize;
+            return this;
+        }
+
         public Builder setQuantizeBits(int quantizeBits) {
             this.quantizeBits = quantizeBits;
             return this;
@@ -346,6 +356,7 @@ public CmdLineArgs build() {
                 filterSelectivity,
                 seed,
                 vectorSpace,
+                rawVectorSize,
                 quantizeBits,
                 vectorEncoding,
                 dimensions,

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -33,6 +33,8 @@
 import org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
+import org.elasticsearch.index.codec.vectors.es92.ES92BinaryQuantizedBFloat16VectorsFormat;
+import org.elasticsearch.index.codec.vectors.es92.ES92HnswBinaryQuantizedBFloat16VectorsFormat;
 import org.elasticsearch.logging.Level;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
@@ -125,9 +127,17 @@ static Codec createCodec(CmdLineArgs args) {
         } else {
             if (args.quantizeBits() == 1) {
                 if (args.indexType() == IndexType.FLAT) {
-                    format = new ES818BinaryQuantizedVectorsFormat();
+                    if (args.rawVectorSize() == 16) {
+                        format = new ES92BinaryQuantizedBFloat16VectorsFormat();
+                    } else {
+                        format = new ES818BinaryQuantizedVectorsFormat();
+                    }
                 } else {
-                    format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
+                    if (args.rawVectorSize() == 16) {
+                        format = new ES92HnswBinaryQuantizedBFloat16VectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
+                    } else {
+                        format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
+                    }
                 }
             } else if (args.quantizeBits() < 32) {
                 if (args.indexType() == IndexType.FLAT) {