DiskBBQ - Always block encode doc vectors

docs/changelog/141598.yaml

@@ -0,0 +1,5 @@
+area: Vector Search
+issues: []
+pr: 141598
+summary: DiskBBQ - Always block encode doc vectors
+type: enhancement

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESNextOSQVectorsScorer.java

@@ -236,6 +236,31 @@ public float scoreBulk(
         float centroidDp,
         float[] scores
     ) throws IOException {
+        return scoreBulk(
+            q,
+            queryLowerInterval,
+            queryUpperInterval,
+            queryComponentSum,
+            queryAdditionalCorrection,
+            similarityFunction,
+            centroidDp,
+            scores,
+            bulkSize
+        );
+    }
+
+    public float scoreBulk(
+        byte[] q,
+        float queryLowerInterval,
+        float queryUpperInterval,
+        int queryComponentSum,
+        float queryAdditionalCorrection,
+        VectorSimilarityFunction similarityFunction,
+        float centroidDp,
+        float[] scores,
+        int bulkSize
+    ) throws IOException {
+        assert bulkSize <= this.bulkSize : "supplied bulkSize > this scorer's bulkSize";
         quantizeScoreBulk(q, bulkSize, scores);
         in.readFloats(lowerIntervals, 0, bulkSize);
         in.readFloats(upperIntervals, 0, bulkSize);

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSBitToInt4ESNextOSQVectorsScorer.java

@@ -392,7 +392,8 @@ public float scoreBulk(
         float queryAdditionalCorrection,
         VectorSimilarityFunction similarityFunction,
         float centroidDp,
-        float[] scores
+        float[] scores,
+        int bulkSize
     ) throws IOException {
         assert q.length == length * 4;
         // 128 / 8 == 16
@@ -414,7 +415,8 @@ public float scoreBulk(
                         queryAdditionalCorrection,
                         similarityFunction,
                         centroidDp,
-                        scores
+                        scores,
+                        bulkSize
                     );
                 } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) {
                     return score128Bulk(
@@ -424,7 +426,8 @@ public float scoreBulk(
                         queryAdditionalCorrection,
                         similarityFunction,
                         centroidDp,
-                        scores
+                        scores,
+                        bulkSize
                     );
                 }
             }
@@ -439,7 +442,8 @@ private float score128Bulk(
         float queryAdditionalCorrection,
         VectorSimilarityFunction similarityFunction,
         float centroidDp,
-        float[] scores
+        float[] scores,
+        int bulkSize
     ) throws IOException {
         int limit = FLOAT_SPECIES_128.loopBound(bulkSize);
         int i = 0;
@@ -501,6 +505,21 @@ private float score128Bulk(
                 }
             }
         }
+        if (limit < bulkSize) {
+            maxScore = scoreTailIndividually(
+                queryAdditionalCorrection,
+                similarityFunction,
+                centroidDp,
+                scores,
+                bulkSize,
+                limit,
+                offset,
+                ay,
+                ly,
+                y1,
+                maxScore
+            );
+        }
         in.seek(offset + 16L * bulkSize);
         return maxScore;
     }
@@ -512,7 +531,8 @@ private float score256Bulk(
         float queryAdditionalCorrection,
         VectorSimilarityFunction similarityFunction,
         float centroidDp,
-        float[] scores
+        float[] scores,
+        int bulkSize
     ) throws IOException {
         int limit = FLOAT_SPECIES_256.loopBound(bulkSize);
         int i = 0;
@@ -574,7 +594,82 @@ private float score256Bulk(
                 }
             }
         }
+        if (limit < bulkSize) {
+            maxScore = scoreTailIndividually(
+                queryAdditionalCorrection,
+                similarityFunction,
+                centroidDp,
+                scores,
+                bulkSize,
+                limit,
+                offset,
+                ay,
+                ly,
+                y1,
+                maxScore
+            );
+        }
         in.seek(offset + 16L * bulkSize);
         return maxScore;
     }
+
+    float scoreTailIndividually(
+        float queryAdditionalCorrection,
+        VectorSimilarityFunction similarityFunction,
+        float centroidDp,
+        float[] scores,
+        int bulkSize,
+        int limit,
+        long offset,
+        float ay,
+        float ly,
+        float y1,
+        float maxScore
+    ) {
+        for (int j = limit; j < bulkSize; j++) {
+            float ax = memorySegment.get(
+                ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + (long) j * Float.BYTES
+            );
+
+            float lx = memorySegment.get(
+                ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + 4L * bulkSize + (long) j * Float.BYTES
+            ) - ax;
+
+            int targetComponentSum = memorySegment.get(
+                ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + 8L * bulkSize + (long) j * Integer.BYTES
+            );
+
+            float additionalCorrection = memorySegment.get(
+                ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + 12L * bulkSize + (long) j * Float.BYTES
+            );
+
+            float qcDist = scores[j];
+
+            float res = ax * ay * dimensions + lx * ay * targetComponentSum + ax * ly * y1 + lx * ly * qcDist;
+
+            if (similarityFunction == EUCLIDEAN) {
+                res = res * -2f + additionalCorrection + queryAdditionalCorrection + 1f;
+                res = Math.max(1f / res, 0f);
+                scores[j] = res;
+                maxScore = Math.max(maxScore, res);
+            } else {
+                res = res + queryAdditionalCorrection + additionalCorrection - centroidDp;
+
+                if (similarityFunction == MAXIMUM_INNER_PRODUCT) {
+                    res = VectorUtil.scaleMaxInnerProductScore(res);
+                    scores[j] = res;
+                    maxScore = Math.max(maxScore, res);
+                } else {
+                    res = Math.max((res + 1f) * 0.5f, 0f);
+                    scores[j] = res;
+                    maxScore = Math.max(maxScore, res);
+                }
+            }
+        }
+        return maxScore;
+    }
 }

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSDibitToInt4ESNextOSQVectorsScorer.java

@@ -284,7 +284,8 @@ public float scoreBulk(
         float queryAdditionalCorrection,
         VectorSimilarityFunction similarityFunction,
         float centroidDp,
-        float[] scores
+        float[] scores,
+        int bulkSize
     ) throws IOException {
         assert q.length == length * 2;
         // 128 / 8 == 16
@@ -306,7 +307,8 @@ public float scoreBulk(
                         queryAdditionalCorrection,
                         similarityFunction,
                         centroidDp,
-                        scores
+                        scores,
+                        bulkSize
                     );
                 } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) {
                     return score128Bulk(
@@ -316,7 +318,8 @@ public float scoreBulk(
                         queryAdditionalCorrection,
                         similarityFunction,
                         centroidDp,
-                        scores
+                        scores,
+                        bulkSize
                     );
                 }
             }
@@ -331,7 +334,8 @@ private float score128Bulk(
         float queryAdditionalCorrection,
         VectorSimilarityFunction similarityFunction,
         float centroidDp,
-        float[] scores
+        float[] scores,
+        int bulkSize
     ) throws IOException {
         int limit = FLOAT_SPECIES_128.loopBound(bulkSize);
         int i = 0;
@@ -393,6 +397,22 @@ private float score128Bulk(
                 }
             }
         }
+        if (limit < bulkSize) {
+            // missing vectors to score
+            maxScore = scoreTailIndividually(
+                queryAdditionalCorrection,
+                similarityFunction,
+                centroidDp,
+                scores,
+                bulkSize,
+                limit,
+                offset,
+                ay,
+                ly,
+                y1,
+                maxScore
+            );
+        }
         in.seek(offset + 16L * bulkSize);
         return maxScore;
     }
@@ -404,7 +424,8 @@ private float score256Bulk(
         float queryAdditionalCorrection,
         VectorSimilarityFunction similarityFunction,
         float centroidDp,
-        float[] scores
+        float[] scores,
+        int bulkSize
     ) throws IOException {
         int limit = FLOAT_SPECIES_256.loopBound(bulkSize);
         int i = 0;
@@ -466,7 +487,84 @@ private float score256Bulk(
                 }
             }
         }
+        if (limit < bulkSize) {
+            // missing vectors to score
+            maxScore = scoreTailIndividually(
+                queryAdditionalCorrection,
+                similarityFunction,
+                centroidDp,
+                scores,
+                bulkSize,
+                limit,
+                offset,
+                ay,
+                ly,
+                y1,
+                maxScore
+            );
+        }
         in.seek(offset + 16L * bulkSize);
         return maxScore;
     }
+
+    private float scoreTailIndividually(
+        float queryAdditionalCorrection,
+        VectorSimilarityFunction similarityFunction,
+        float centroidDp,
+        float[] scores,
+        int bulkSize,
+        int limit,
+        long offset,
+        float ay,
+        float ly,
+        float y1,
+        float maxScore
+    ) {
+        for (int j = limit; j < bulkSize; j++) {
+            float ax = memorySegment.get(
+                ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + (long) j * Float.BYTES
+            );
+
+            float lx = memorySegment.get(
+                ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + 4L * bulkSize + (long) j * Float.BYTES
+            );
+            lx = (lx - ax) * TWO_BIT_SCALE;
+
+            int targetComponentSum = memorySegment.get(
+                ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + 8L * bulkSize + (long) j * Integer.BYTES
+            );
+
+            float additionalCorrection = memorySegment.get(
+                ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
+                offset + 12L * bulkSize + (long) j * Float.BYTES
+            );
+
+            float qcDist = scores[j];
+
+            float res = ax * ay * dimensions + lx * ay * targetComponentSum + ax * ly * y1 + lx * ly * qcDist;
+
+            if (similarityFunction == EUCLIDEAN) {
+                res = res * -2f + additionalCorrection + queryAdditionalCorrection + 1f;
+                res = Math.max(1f / res, 0f);
+                scores[j] = res;
+                maxScore = Math.max(maxScore, res);
+            } else {
+                res = res + queryAdditionalCorrection + additionalCorrection - centroidDp;
+
+                if (similarityFunction == MAXIMUM_INNER_PRODUCT) {
+                    res = VectorUtil.scaleMaxInnerProductScore(res);
+                    scores[j] = res;
+                    maxScore = Math.max(maxScore, res);
+                } else {
+                    res = Math.max((res + 1f) * 0.5f, 0f);
+                    scores[j] = res;
+                    maxScore = Math.max(maxScore, res);
+                }
+            }
+        }
+        return maxScore;
+    }
 }