Fix setting the query bits / index bits statically and printing them out

Author: carlosdelest

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

@@ -187,7 +187,13 @@ public static void main(String[] args) throws Exception {
                 : new int[] { 0 };
             Results[] results = new Results[nProbes.length];
             for (int i = 0; i < nProbes.length; i++) {
-                results[i] = new Results(cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT), cmdLineArgs.numDocs());
+                results[i] = new Results(
+                    cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT),
+                    cmdLineArgs.numDocs(),
+                    cmdLineArgs.quantizeBits(),
+                    cmdLineArgs.quantizeQueryBits(),
+                    cmdLineArgs.overSamplingFactor()
+                );
             }
             logger.info("Running KNN index tester with arguments: " + cmdLineArgs);
             Codec codec = createCodec(cmdLineArgs);
@@ -247,7 +253,11 @@ public String toString() {
                 "avg_cpu_count",
                 "QPS",
                 "recall",
-                "visited" };
+                "visited",
+                "indexBits",
+                "queryBits",
+                "oversampling"
+            };
 
             // Calculate appropriate column widths based on headers and data
 
@@ -274,8 +284,12 @@ public String toString() {
                     String.format(Locale.ROOT, "%.2f", result.netCpuTimeMS),
                     String.format(Locale.ROOT, "%.2f", result.avgCpuCount),
                     String.format(Locale.ROOT, "%.2f", result.qps),
-                    String.format(Locale.ROOT, "%.2f", result.avgRecall),
-                    String.format(Locale.ROOT, "%.2f", result.averageVisited) };
+                    String.format(Locale.ROOT, "%.4f", result.avgRecall),
+                    String.format(Locale.ROOT, "%.2f", result.averageVisited),
+                    String.format(Locale.ROOT, "%d", result.indexBits),
+                    String.format(Locale.ROOT, "%d", result.queryBits),
+                    String.format(Locale.ROOT, "%.2f", result.oversampling)
+                };
 
             }
 
@@ -341,6 +355,8 @@ private int[] calculateColumnWidths(String[] headers, String[]... data) {
     static class Results {
         final String indexType;
         final int numDocs;
+        int indexBits;
+        int queryBits;
         long indexTimeMS;
         long forceMergeTimeMS;
         int numSegments;
@@ -351,10 +367,14 @@ static class Results {
         double averageVisited;
         double netCpuTimeMS;
         double avgCpuCount;
+        float oversampling;
 
-        Results(String indexType, int numDocs) {
+        Results(String indexType, int numDocs, int indexBits, int queryBits, float oversampling) {
             this.indexType = indexType;
             this.numDocs = numDocs;
+            this.indexBits = indexBits;
+            this.queryBits = queryBits;
+            this.oversampling = oversampling;
         }
     }
 

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

@@ -45,6 +45,7 @@
 import org.apache.lucene.store.MMapDirectory;
 import org.elasticsearch.common.io.Channels;
 import org.elasticsearch.core.PathUtils;
+import org.elasticsearch.index.codec.vectors.es910.ES910BinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
@@ -115,6 +116,13 @@ class KnnSearcher {
         this.nProbe = nProbe;
         this.indexType = cmdLineArgs.indexType();
         this.searchThreads = cmdLineArgs.searchThreads();
+        if (cmdLineArgs.useNewFlatVectorsFormat()) {
+            // use the query bits statically
+            ES910BinaryQuantizedVectorsFormat.setQuantizationBits(
+                (byte) cmdLineArgs.quantizeBits(),
+                (byte) cmdLineArgs.quantizeQueryBits()
+            );
+        }
     }
 
     void runSearch(KnnIndexTester.Results finalResults, boolean earlyTermination) throws IOException {

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

@@ -27,7 +27,6 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
-import org.elasticsearch.index.codec.vectors.es818.DirectIOLucene99FlatVectorsFormat;
 
 import java.io.IOException;
 
@@ -102,38 +101,37 @@ public class ES910BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
     private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
-    private static byte DEFAULT_INDEX_BITS = (byte) 1;
-    private static byte DEFAULT_QUERY_BITS = (byte) 4;
 
-    private final ES910BinaryFlatVectorsScorer scorer;
-
-    private final byte indexBits;
-    private final byte queryBits;
+    // index and query bits for quantization. They are static as we have no mapping, and thus can't create a format with specific bits.
+    // we use static setters to change the quantization bits and scorers.
+    private static byte INDEX_BITS = (byte) 1;
+    private static byte QUERY_BITS = (byte) 4;
+    private static ES910BinaryFlatVectorsScorer SCORER;
 
     public ES910BinaryQuantizedVectorsFormat() {
-        this(DEFAULT_INDEX_BITS, DEFAULT_QUERY_BITS); // Default to 4 bits for index and 2 bits for query
+        this(INDEX_BITS, QUERY_BITS); // Default to 4 bits for index and 2 bits for query
     }
 
     /** Creates a new instance with the default number of vectors per cluster. */
     public ES910BinaryQuantizedVectorsFormat(byte indexBits, byte queryBits) {
         super(NAME);
-        this.indexBits = indexBits;
-        this.queryBits = queryBits;
-        // Set the default bits for index and query vectors. I know, I know, this is a hack, but we
-        // don't have the possibility of doing a PerFieldMapperCodec yet on KnnSearcher
-        DEFAULT_QUERY_BITS = queryBits;
-        DEFAULT_INDEX_BITS = indexBits;
-        this.scorer = new ES910BinaryFlatVectorsScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer(), indexBits, queryBits);
+        setQuantizationBits(indexBits, queryBits);
+    }
+
+    public static void setQuantizationBits(byte indexBits, byte queryBits) {
+        INDEX_BITS = indexBits;
+        QUERY_BITS = queryBits;
+        SCORER = new ES910BinaryFlatVectorsScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer(), indexBits, queryBits);
     }
 
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-        return new ES910BinaryQuantizedVectorsWriter(scorer, rawVectorFormat.fieldsWriter(state), state, indexBits, queryBits);
+        return new ES910BinaryQuantizedVectorsWriter(SCORER, rawVectorFormat.fieldsWriter(state), state, INDEX_BITS, QUERY_BITS);
     }
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new ES910BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), scorer);
+        return new ES910BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), SCORER);
     }
 
     @Override
@@ -146,11 +144,11 @@ public String toString() {
         return "ES910BinaryQuantizedVectorsFormat(name="
             + NAME
             + ", flatVectorScorer="
-            + scorer
+            + SCORER
             + ", indexBits="
-            + indexBits
+            + INDEX_BITS
             + ", queryBits="
-            + queryBits
+            + QUERY_BITS
             + ")";
     }
 }