DiskBBQ - Adapt visited_ratio based on query - segment affinity in multi segment scenario

docs/changelog/132396.yaml

@@ -0,0 +1,6 @@
+pr: 132396
+summary: DiskBBQ - Adapt `nProbe` based on query - segment affinity in multi segment
+  scenario
+area: Vector Search
+type: enhancement
+issues: []

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

@@ -63,7 +63,6 @@ record CmdLineArgs(
     static final ParseField INDEX_TYPE_FIELD = new ParseField("index_type");
     static final ParseField NUM_CANDIDATES_FIELD = new ParseField("num_candidates");
     static final ParseField K_FIELD = new ParseField("k");
-    // static final ParseField N_PROBE_FIELD = new ParseField("n_probe");
     static final ParseField VISIT_PERCENTAGE_FIELD = new ParseField("visit_percentage");
     static final ParseField IVF_CLUSTER_SIZE_FIELD = new ParseField("ivf_cluster_size");
     static final ParseField OVER_SAMPLING_FACTOR_FIELD = new ParseField("over_sampling_factor");
@@ -98,7 +97,6 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareString(Builder::setIndexType, INDEX_TYPE_FIELD);
         PARSER.declareInt(Builder::setNumCandidates, NUM_CANDIDATES_FIELD);
         PARSER.declareInt(Builder::setK, K_FIELD);
-        // PARSER.declareIntArray(Builder::setNProbe, N_PROBE_FIELD);
         PARSER.declareDoubleArray(Builder::setVisitPercentages, VISIT_PERCENTAGE_FIELD);
         PARSER.declareInt(Builder::setIvfClusterSize, IVF_CLUSTER_SIZE_FIELD);
         PARSER.declareInt(Builder::setOverSamplingFactor, OVER_SAMPLING_FACTOR_FIELD);
@@ -134,7 +132,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(INDEX_TYPE_FIELD.getPreferredName(), indexType.name().toLowerCase(Locale.ROOT));
         builder.field(NUM_CANDIDATES_FIELD.getPreferredName(), numCandidates);
         builder.field(K_FIELD.getPreferredName(), k);
-        // builder.field(N_PROBE_FIELD.getPreferredName(), nProbes);
         builder.field(VISIT_PERCENTAGE_FIELD.getPreferredName(), visitPercentages);
         builder.field(IVF_CLUSTER_SIZE_FIELD.getPreferredName(), ivfClusterSize);
         builder.field(OVER_SAMPLING_FACTOR_FIELD.getPreferredName(), overSamplingFactor);

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

@@ -86,6 +86,63 @@ CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, Inde
         return getCentroidIteratorNoParent(fieldInfo, centroids, numCentroids, scorer, quantized, queryParams, globalCentroidDp);
     }
 
+    @Override
+    public float[] getCentroidsScores(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] targetQuery, boolean parents)
+        throws IOException {
+        final FieldEntry fieldEntry = fields.get(fieldInfo.number);
+        final float globalCentroidDp = fieldEntry.globalCentroidDp();
+        final OptimizedScalarQuantizer scalarQuantizer = new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
+        final int[] scratch = new int[targetQuery.length];
+        float[] targetQueryCopy = ArrayUtil.copyArray(targetQuery);
+        if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
+            VectorUtil.l2normalize(targetQueryCopy);
+        }
+        final OptimizedScalarQuantizer.QuantizationResult queryParams = scalarQuantizer.scalarQuantize(
+            targetQueryCopy,
+            scratch,
+            (byte) 7,
+            fieldEntry.globalCentroid()
+        );
+        final byte[] quantizedQuery = new byte[targetQuery.length];
+        for (int i = 0; i < quantizedQuery.length; i++) {
+            quantizedQuery[i] = (byte) scratch[i];
+        }
+        final ES92Int7VectorsScorer scorer = ESVectorUtil.getES92Int7VectorsScorer(centroids, fieldInfo.getVectorDimension());
+        centroids.seek(0L);
+        // final scores
+        final float[] scores = new float[ES92Int7VectorsScorer.BULK_SIZE];
+
+        int numParents = centroids.readVInt();
+        if (parents && numParents > 0) {
+            final NeighborQueue parentsQueue = new NeighborQueue(numParents, true);
+            score(
+                parentsQueue,
+                numParents,
+                0,
+                scorer,
+                quantizedQuery,
+                queryParams,
+                globalCentroidDp,
+                fieldInfo.getVectorSimilarityFunction(),
+                scores
+            );
+        } else {
+            final NeighborQueue neighborQueue = new NeighborQueue(numCentroids, true);
+            score(
+                neighborQueue,
+                numCentroids,
+                0,
+                scorer,
+                quantizedQuery,
+                queryParams,
+                globalCentroidDp,
+                fieldInfo.getVectorSimilarityFunction(),
+                scores
+            );
+        }
+        return scores;
+    }
+
     private static CentroidIterator getCentroidIteratorNoParent(
         FieldInfo fieldInfo,
         IndexInput centroids,

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

@@ -88,6 +88,9 @@ protected IVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsR
     abstract CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target)
         throws IOException;
 
+    public abstract float[] getCentroidsScores(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target, boolean parents)
+        throws IOException;
+
     private static IndexInput openDataInput(
         SegmentReadState state,
         int versionMeta,
@@ -258,6 +261,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             // is enough?
             expectedDocs += scorer.resetPostingsScorer(offset);
             actualDocs += scorer.visit(knnCollector);
+            if (knnCollector.earlyTerminated()) {
+                break;
+            }
         }
         if (acceptDocs != null) {
             float unfilteredRatioVisited = (float) expectedDocs / numVectors;
@@ -267,6 +273,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
                 long offset = centroidIterator.nextPostingListOffset();
                 scorer.resetPostingsScorer(offset);
                 actualDocs += scorer.visit(knnCollector);
+                if (knnCollector.earlyTerminated()) {
+                    break;
+                }
             }
         }
     }
@@ -318,6 +327,17 @@ interface CentroidIterator {
         long nextPostingListOffset() throws IOException;
     }
 
+    // TODO : change this to allow (batched) centroid search (with current top)
+    public interface ScoredCentroidIterator {
+        boolean hasNext();
+
+        void scorePostingList(long offset) throws IOException;
+
+        long next();
+
+        long nextPostingListOffset() throws IOException;
+    }
+
     interface PostingVisitor {
         // TODO maybe we can not specifically pass the centroid...
 
@@ -327,4 +347,16 @@ interface PostingVisitor {
         /** returns the number of scored documents */
         int visit(KnnCollector collector) throws IOException;
     }
+
+    public IndexInput getIvfCentroids(FieldInfo fieldInfo) throws IOException {
+        return fields.get(fieldInfo.number).centroidSlice(ivfCentroids);
+    }
+
+    public int getNumCentroids(FieldInfo fieldInfo) {
+        return fields.get(fieldInfo.number).numCentroids;
+    }
+
+    public float[] getGlobalCentroid(FieldInfo fieldInfo) {
+        return fields.get(fieldInfo.number).globalCentroid;
+    }
 }