Use bit scale for both index and queries

Author: carlosdelest

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

@@ -39,9 +39,11 @@
 public class ES910BinaryFlatVectorsScorer implements FlatVectorsScorer {
     private final FlatVectorsScorer nonQuantizedDelegate;
     private final byte queryBits;
+    private final byte indexBits;
 
-    public ES910BinaryFlatVectorsScorer(FlatVectorsScorer nonQuantizedDelegate, byte queryBits) {
+    public ES910BinaryFlatVectorsScorer(FlatVectorsScorer nonQuantizedDelegate, byte indexBits, byte queryBits) {
         this.nonQuantizedDelegate = nonQuantizedDelegate;
+        this.indexBits = indexBits;
         this.queryBits = queryBits;
     }
 
@@ -85,15 +87,20 @@ public float score(int i) throws IOException {
                         queryCorrections,
                         binarizedVectors.vectorValue(i),
                         binarizedVectors.getCorrectiveTerms(i),
-                        getBitsScale()
+                        getIndexBitsScale(),
+                        getQueryBitsScale()
                     );
                 }
             };
         }
         return nonQuantizedDelegate.getRandomVectorScorer(similarityFunction, vectorValues, target);
     }
 
-    private float getBitsScale() {
+    private float getIndexBitsScale() {
+        return 1f / ((1 << indexBits) - 1);
+    }
+
+    private float getQueryBitsScale() {
         return 1f / ((1 << queryBits) - 1);
     }
 
@@ -111,7 +118,7 @@ RandomVectorScorerSupplier getRandomVectorScorerSupplier(
         ES910BinaryQuantizedVectorsWriter.OffHeapBinarizedQueryVectorValues scoringVectors,
         BinarizedByteVectorValues targetVectors
     ) {
-        return new BinarizedRandomVectorScorerSupplier(scoringVectors, targetVectors, similarityFunction, queryBits);
+        return new BinarizedRandomVectorScorerSupplier(scoringVectors, targetVectors, similarityFunction, indexBits, queryBits);
     }
 
     @Override
@@ -124,28 +131,31 @@ static class BinarizedRandomVectorScorerSupplier implements RandomVectorScorerSu
         private final ES910BinaryQuantizedVectorsWriter.OffHeapBinarizedQueryVectorValues queryVectors;
         private final BinarizedByteVectorValues targetVectors;
         private final VectorSimilarityFunction similarityFunction;
+        private final byte indexBits;
         private final byte queryBits;
 
         BinarizedRandomVectorScorerSupplier(
             ES910BinaryQuantizedVectorsWriter.OffHeapBinarizedQueryVectorValues queryVectors,
             BinarizedByteVectorValues targetVectors,
             VectorSimilarityFunction similarityFunction,
+            byte indexBits,
             byte queryBits
         ) {
             this.queryVectors = queryVectors;
             this.targetVectors = targetVectors;
             this.similarityFunction = similarityFunction;
+            this.indexBits = indexBits;
             this.queryBits = queryBits;
         }
 
         @Override
         public BinarizedRandomVectorScorer scorer() throws IOException {
-            return new BinarizedRandomVectorScorer(queryVectors.copy(), targetVectors.copy(), similarityFunction, queryBits);
+            return new BinarizedRandomVectorScorer(queryVectors.copy(), targetVectors.copy(), similarityFunction, indexBits, queryBits);
         }
 
         @Override
         public RandomVectorScorerSupplier copy() throws IOException {
-            return new BinarizedRandomVectorScorerSupplier(queryVectors, targetVectors, similarityFunction, queryBits);
+            return new BinarizedRandomVectorScorerSupplier(queryVectors, targetVectors, similarityFunction, indexBits, queryBits);
         }
     }
 
@@ -157,20 +167,23 @@ public static class BinarizedRandomVectorScorer extends UpdateableRandomVectorSc
         private final byte[] quantizedQuery;
         private OptimizedScalarQuantizer.QuantizationResult queryCorrections = null;
         private int currentOrdinal = -1;
-        private final float bitScale;
+        private final float queryBitsScale;
+        private final float indexBitsScale;
 
         BinarizedRandomVectorScorer(
             ES910BinaryQuantizedVectorsWriter.OffHeapBinarizedQueryVectorValues queryVectors,
             BinarizedByteVectorValues targetVectors,
             VectorSimilarityFunction similarityFunction,
+            byte indexBits,
             byte queryBits
         ) {
             super(targetVectors);
             this.queryVectors = queryVectors;
             this.quantizedQuery = new byte[queryVectors.dimension()];
             this.targetVectors = targetVectors;
             this.similarityFunction = similarityFunction;
-            bitScale = 1.0F / (float) ((1 << queryBits) - 1);
+            this.indexBitsScale = 1.0F / (float) ((1 << indexBits) - 1);
+            this.queryBitsScale = 1.0F / (float) ((1 << queryBits) - 1);
         }
 
         @Override
@@ -186,7 +199,8 @@ public float score(int targetOrd) throws IOException {
                 queryCorrections,
                 targetVectors.vectorValue(targetOrd),
                 targetVectors.getCorrectiveTerms(targetOrd),
-                bitScale
+                indexBitsScale,
+                queryBitsScale
             );
         }
 
@@ -209,15 +223,15 @@ private static float quantizedScore(
         OptimizedScalarQuantizer.QuantizationResult queryCorrections,
         byte[] d,
         OptimizedScalarQuantizer.QuantizationResult indexCorrections,
-        float bitsScale
+        float indexBitsScale,
+        float queryBitsScale
     ) {
         float qcDist = VectorUtil.dotProduct(q, d);
         float x1 = indexCorrections.quantizedComponentSum();
         float ax = indexCorrections.lowerInterval();
-        // Here we assume `lx` is simply bit vectors, so the scaling isn't necessary
-        float lx = indexCorrections.upperInterval() - ax;
+        float lx = (indexCorrections.upperInterval() - ax) * indexBitsScale;
         float ay = queryCorrections.lowerInterval();
-        float ly = (queryCorrections.upperInterval() - ay) * bitsScale;
+        float ly = (queryCorrections.upperInterval() - ay) * queryBitsScale;
         float y1 = queryCorrections.quantizedComponentSum();
         float score = ax * ay * dims + ay * lx * x1 + ax * ly * y1 + lx * ly * qcDist;
         // For euclidean, we need to invert the score and apply the additional correction, which is

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

@@ -122,7 +122,7 @@ public ES910BinaryQuantizedVectorsFormat(byte indexBits, byte queryBits) {
         // 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(), queryBits);
+        this.scorer = new ES910BinaryFlatVectorsScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer(), indexBits, queryBits);
     }
 
     @Override