Add CLI args for KnnIndexTester

Author: carlosdelest

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

@@ -47,7 +47,9 @@ record CmdLineArgs(
     VectorSimilarityFunction vectorSpace,
     int quantizeBits,
     VectorEncoding vectorEncoding,
-    int dimensions
+    int dimensions,
+    boolean useNewFlatVectorsFormat,
+    int quantizeQueryBits
 ) implements ToXContentObject {
 
     static final ParseField DOC_VECTORS_FIELD = new ParseField("doc_vectors");
@@ -70,6 +72,9 @@ record CmdLineArgs(
     static final ParseField QUANTIZE_BITS_FIELD = new ParseField("quantize_bits");
     static final ParseField VECTOR_ENCODING_FIELD = new ParseField("vector_encoding");
     static final ParseField DIMENSIONS_FIELD = new ParseField("dimensions");
+    static final ParseField QUERY_BITS_FIELD = new ParseField("query_bits");
+    static final ParseField USE_NEW_FLAT_VECTORS_FORMAT_FIELD = new ParseField("use_new_flat_vectors_format");
+    static final ParseField QUANTIZE_QUERY_BITS_FIELD = new ParseField("quantize_query_bits");
 
     static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         Builder builder = PARSER.apply(parser, null);
@@ -99,6 +104,8 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareInt(Builder::setQuantizeBits, QUANTIZE_BITS_FIELD);
         PARSER.declareString(Builder::setVectorEncoding, VECTOR_ENCODING_FIELD);
         PARSER.declareInt(Builder::setDimensions, DIMENSIONS_FIELD);
+        PARSER.declareBoolean(Builder::setUseNewFlatVectorsFormat, USE_NEW_FLAT_VECTORS_FORMAT_FIELD);
+        PARSER.declareInt(Builder::setQuantizeQueryBits, QUANTIZE_QUERY_BITS_FIELD);
     }
 
     @Override
@@ -128,6 +135,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(QUANTIZE_BITS_FIELD.getPreferredName(), quantizeBits);
         builder.field(VECTOR_ENCODING_FIELD.getPreferredName(), vectorEncoding.name().toLowerCase(Locale.ROOT));
         builder.field(DIMENSIONS_FIELD.getPreferredName(), dimensions);
+        builder.field(USE_NEW_FLAT_VECTORS_FORMAT_FIELD.getPreferredName(), useNewFlatVectorsFormat);
+        builder.field(QUANTIZE_QUERY_BITS_FIELD.getPreferredName(), quantizeQueryBits);
         return builder.endObject();
     }
 
@@ -157,6 +166,8 @@ static class Builder {
         private int quantizeBits = 8;
         private VectorEncoding vectorEncoding = VectorEncoding.FLOAT32;
         private int dimensions;
+        private boolean useNewFlatVectorsFormat = false;
+        private int quantizeQueryBits = 8;
 
         public Builder setDocVectors(String docVectors) {
             this.docVectors = PathUtils.get(docVectors);
@@ -258,6 +269,16 @@ public Builder setDimensions(int dimensions) {
             return this;
         }
 
+        public Builder setUseNewFlatVectorsFormat(boolean useNewFlatVectorsFormat) {
+            this.useNewFlatVectorsFormat = useNewFlatVectorsFormat;
+            return this;
+        }
+
+        public Builder setQuantizeQueryBits(int quantizeQueryBits) {
+            this.quantizeQueryBits = quantizeQueryBits;
+            return this;
+        }
+
         public CmdLineArgs build() {
             if (docVectors == null) {
                 throw new IllegalArgumentException("Document vectors path must be provided");
@@ -285,7 +306,9 @@ public CmdLineArgs build() {
                 vectorSpace,
                 quantizeBits,
                 vectorEncoding,
-                dimensions
+                dimensions,
+                useNewFlatVectorsFormat,
+                quantizeQueryBits
             );
         }
     }

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
+import org.elasticsearch.index.codec.vectors.es910.ES910BinaryQuantizedVectorsFormat;
 import org.elasticsearch.logging.Level;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
@@ -91,7 +92,10 @@ static Codec createCodec(CmdLineArgs args) {
         if (args.indexType() == IndexType.IVF) {
             format = new IVFVectorsFormat(args.ivfClusterSize());
         } else {
-            if (args.quantizeBits() == 1) {
+            if (args.useNewFlatVectorsFormat() && args.indexType() == IndexType.FLAT) {
+                logger.warn("Using new flat vectors format for index type FLAT");
+                format = new ES910BinaryQuantizedVectorsFormat((byte) args.quantizeBits(), (byte) args.quantizeQueryBits());
+            } else if (args.quantizeBits() == 1) {
                 if (args.indexType() == IndexType.FLAT) {
                     format = new ES818BinaryQuantizedVectorsFormat();
                 } else {
@@ -113,6 +117,7 @@ static Codec createCodec(CmdLineArgs args) {
                 format = new Lucene99HnswVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
             }
         }
+        logger.info("Using vector format: " + format);
         return new Lucene101Codec() {
             @Override
             public KnnVectorsFormat getKnnVectorsFormatForField(String field) {

server/src/main/java/module-info.java

@@ -483,4 +483,5 @@
     exports org.elasticsearch.lucene.search;
     exports org.elasticsearch.index.codec.vectors to org.elasticsearch.test.knn;
     exports org.elasticsearch.index.codec.vectors.es818 to org.elasticsearch.test.knn;
+    exports org.elasticsearch.index.codec.vectors.es910;
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/OptimizedScalarQuantizer.java

@@ -132,6 +132,8 @@ public QuantizationResult scalarQuantize(float[] vector, byte[] destination, byt
             int assignment = Math.round((xi - a) / step);
             sumQuery += assignment;
             destination[h] = (byte) assignment;
+            assert assignment >= 0 && assignment < points
+                : "assignment: " + assignment + ", a: " + a + ", b: " + b + ", step: " + step + ", xi: " + xi;
         }
         return new QuantizationResult(
             intervalScratch[0],

server/src/main/java/org/elasticsearch/index/codec/vectors/es910/ES910BinaryFlatVectorsScorer.java

@@ -72,8 +72,9 @@ public RandomVectorScorer getRandomVectorScorer(
             if (similarityFunction == COSINE) {
                 VectorUtil.l2normalize(copy);
             }
-            byte[] quantized = new byte[copy.length];
-            OptimizedScalarQuantizer.QuantizationResult queryCorrections = quantizer.scalarQuantize(copy, quantized, queryBits, centroid);
+            target = copy;
+            byte[] quantized = new byte[target.length];
+            OptimizedScalarQuantizer.QuantizationResult queryCorrections = quantizer.scalarQuantize(target, quantized, queryBits, centroid);
             return new RandomVectorScorer.AbstractRandomVectorScorer(vectorValues) {
                 @Override
                 public float score(int i) throws IOException {
@@ -116,7 +117,7 @@ RandomVectorScorerSupplier getRandomVectorScorerSupplier(
 
     @Override
     public String toString() {
-        return "ES910BinaryFlatVectorsScorer(nonQuantizedDelegate=" + nonQuantizedDelegate + ")";
+        return "ES910BinaryFlatVectorsScorer(nonQuantizedDelegate=" + nonQuantizedDelegate + ", queryBits = " + queryBits + ")";
     }
 
     /** Vector scorer supplier over binarized vector values */
@@ -171,7 +172,6 @@ public static class BinarizedRandomVectorScorer extends UpdateableRandomVectorSc
             this.targetVectors = targetVectors;
             this.similarityFunction = similarityFunction;
             bitScale = 1.0F / (float) ((1 << queryBits) - 1);
-
         }
 
         @Override