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

@@ -36,7 +36,7 @@ record CmdLineArgs(
     KnnIndexTester.IndexType indexType,
     int numCandidates,
     int k,
-    int[] nProbes,
+    double[] visitPercentages,
     int ivfClusterSize,
     int overSamplingFactor,
     int hnswM,
@@ -53,7 +53,8 @@ record CmdLineArgs(
     VectorEncoding vectorEncoding,
     int dimensions,
     boolean earlyTermination,
-    KnnIndexTester.MergePolicyType mergePolicy
+    KnnIndexTester.MergePolicyType mergePolicy,
+    float vectorsRatio
 ) implements ToXContentObject {
 
     static final ParseField DOC_VECTORS_FIELD = new ParseField("doc_vectors");
@@ -63,7 +64,8 @@ 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 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");
     static final ParseField HNSW_M_FIELD = new ParseField("hnsw_m");
@@ -81,6 +83,7 @@ record CmdLineArgs(
     static final ParseField FILTER_SELECTIVITY_FIELD = new ParseField("filter_selectivity");
     static final ParseField SEED_FIELD = new ParseField("seed");
     static final ParseField MERGE_POLICY_FIELD = new ParseField("merge_policy");
+    static final ParseField VECTORS_RATIO = new ParseField("vectors_ratio");
 
     static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         Builder builder = PARSER.apply(parser, null);
@@ -97,7 +100,8 @@ 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.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);
         PARSER.declareInt(Builder::setHnswM, HNSW_M_FIELD);
@@ -115,6 +119,7 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareFloat(Builder::setFilterSelectivity, FILTER_SELECTIVITY_FIELD);
         PARSER.declareLong(Builder::setSeed, SEED_FIELD);
         PARSER.declareString(Builder::setMergePolicy, MERGE_POLICY_FIELD);
+        PARSER.declareFloat(Builder::setVectorsRatio, VECTORS_RATIO);
     }
 
     @Override
@@ -132,7 +137,8 @@ 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(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);
         builder.field(HNSW_M_FIELD.getPreferredName(), hnswM);
@@ -149,6 +155,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(EARLY_TERMINATION_FIELD.getPreferredName(), earlyTermination);
         builder.field(FILTER_SELECTIVITY_FIELD.getPreferredName(), filterSelectivity);
         builder.field(SEED_FIELD.getPreferredName(), seed);
+        builder.field(VECTORS_RATIO.getPreferredName(), vectorsRatio);
         return builder.endObject();
     }
 
@@ -165,7 +172,7 @@ static class Builder {
         private KnnIndexTester.IndexType indexType = KnnIndexTester.IndexType.HNSW;
         private int numCandidates = 1000;
         private int k = 10;
-        private int[] nProbes = new int[] { 10 };
+        private double[] visitPercentages = new double[] { 1.0 };
         private int ivfClusterSize = 1000;
         private int overSamplingFactor = 1;
         private int hnswM = 16;
@@ -183,6 +190,7 @@ static class Builder {
         private float filterSelectivity = 1f;
         private long seed = 1751900822751L;
         private KnnIndexTester.MergePolicyType mergePolicy = null;
+        private float vectorsRatio = 1f;
 
         public Builder setDocVectors(List<String> docVectors) {
             if (docVectors == null || docVectors.isEmpty()) {
@@ -223,8 +231,8 @@ public Builder setK(int k) {
             return this;
         }
 
-        public Builder setNProbe(List<Integer> nProbes) {
-            this.nProbes = nProbes.stream().mapToInt(Integer::intValue).toArray();
+        public Builder setVisitPercentages(List<Double> visitPercentages) {
+            this.visitPercentages = visitPercentages.stream().mapToDouble(Double::doubleValue).toArray();
             return this;
         }
 
@@ -313,6 +321,11 @@ public Builder setMergePolicy(String mergePolicy) {
             return this;
         }
 
+        public Builder setVectorsRatio(float vectorsRatio) {
+            this.vectorsRatio = vectorsRatio;
+            return this;
+        }
+
         public CmdLineArgs build() {
             if (docVectors == null) {
                 throw new IllegalArgumentException("Document vectors path must be provided");
@@ -330,7 +343,7 @@ public CmdLineArgs build() {
                 indexType,
                 numCandidates,
                 k,
-                nProbes,
+                visitPercentages,
                 ivfClusterSize,
                 overSamplingFactor,
                 hnswM,
@@ -347,7 +360,8 @@ public CmdLineArgs build() {
                 vectorEncoding,
                 dimensions,
                 earlyTermination,
-                mergePolicy
+                mergePolicy,
+                vectorsRatio
             );
         }
     }

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

@@ -191,18 +191,18 @@ public static void main(String[] args) throws Exception {
         FormattedResults formattedResults = new FormattedResults();
 
         for (CmdLineArgs cmdLineArgs : cmdLineArgsList) {
-            int[] nProbes = cmdLineArgs.indexType().equals(IndexType.IVF) && cmdLineArgs.numQueries() > 0
-                ? cmdLineArgs.nProbes()
-                : new int[] { 0 };
+            double[] visitPercentages = cmdLineArgs.indexType().equals(IndexType.IVF) && cmdLineArgs.numQueries() > 0
+                ? cmdLineArgs.visitPercentages()
+                : new double[] { 0 };
             String indexType = cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT);
             Results indexResults = new Results(
                 cmdLineArgs.docVectors().get(0).getFileName().toString(),
                 indexType,
                 cmdLineArgs.numDocs(),
                 cmdLineArgs.filterSelectivity()
             );
-            Results[] results = new Results[nProbes.length];
-            for (int i = 0; i < nProbes.length; i++) {
+            Results[] results = new Results[visitPercentages.length];
+            for (int i = 0; i < visitPercentages.length; i++) {
                 results[i] = new Results(
                     cmdLineArgs.docVectors().get(0).getFileName().toString(),
                     indexType,
@@ -240,8 +240,7 @@ public static void main(String[] args) throws Exception {
             numSegments(indexPath, indexResults);
             if (cmdLineArgs.queryVectors() != null && cmdLineArgs.numQueries() > 0) {
                 for (int i = 0; i < results.length; i++) {
-                    int nProbe = nProbes[i];
-                    KnnSearcher knnSearcher = new KnnSearcher(indexPath, cmdLineArgs, nProbe);
+                    KnnSearcher knnSearcher = new KnnSearcher(indexPath, cmdLineArgs, visitPercentages[i]);
                     knnSearcher.runSearch(results[i], cmdLineArgs.earlyTermination());
                 }
             }
@@ -293,7 +292,7 @@ public String toString() {
             String[] searchHeaders = {
                 "index_name",
                 "index_type",
-                "n_probe",
+                "visit_percentage(%)",
                 "latency(ms)",
                 "net_cpu_time(ms)",
                 "avg_cpu_count",
@@ -324,7 +323,7 @@ public String toString() {
                 queryResultsArray[i] = new String[] {
                     queryResult.indexName,
                     queryResult.indexType,
-                    Integer.toString(queryResult.nProbe),
+                    String.format(Locale.ROOT, "%.2f", queryResult.visitPercentage),
                     String.format(Locale.ROOT, "%.2f", queryResult.avgLatency),
                     String.format(Locale.ROOT, "%.2f", queryResult.netCpuTimeMS),
                     String.format(Locale.ROOT, "%.2f", queryResult.avgCpuCount),
@@ -400,7 +399,7 @@ static class Results {
         long indexTimeMS;
         long forceMergeTimeMS;
         int numSegments;
-        int nProbe;
+        double visitPercentage;
         double avgLatency;
         double qps;
         double avgRecall;

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

@@ -107,7 +107,7 @@ class KnnSearcher {
     private final float selectivity;
     private final int topK;
     private final int efSearch;
-    private final int nProbe;
+    private final double visitPercentage;
     private final KnnIndexTester.IndexType indexType;
     private int dim;
     private final VectorSimilarityFunction similarityFunction;
@@ -116,7 +116,7 @@ class KnnSearcher {
     private final int searchThreads;
     private final int numSearchers;
 
-    KnnSearcher(Path indexPath, CmdLineArgs cmdLineArgs, int nProbe) {
+    KnnSearcher(Path indexPath, CmdLineArgs cmdLineArgs, double visitPercentage) {
         this.docPath = cmdLineArgs.docVectors();
         this.indexPath = indexPath;
         this.queryPath = cmdLineArgs.queryVectors();
@@ -131,7 +131,7 @@ class KnnSearcher {
             throw new IllegalArgumentException("numQueryVectors must be > 0");
         }
         this.efSearch = cmdLineArgs.numCandidates();
-        this.nProbe = nProbe;
+        this.visitPercentage = visitPercentage;
         this.indexType = cmdLineArgs.indexType();
         this.searchThreads = cmdLineArgs.searchThreads();
         this.numSearchers = cmdLineArgs.numSearchers();
@@ -298,7 +298,7 @@ void runSearch(KnnIndexTester.Results finalResults, boolean earlyTermination) th
         }
         logger.info("checking results");
         int[][] nn = getOrCalculateExactNN(offsetByteSize, filterQuery);
-        finalResults.nProbe = indexType == KnnIndexTester.IndexType.IVF ? nProbe : 0;
+        finalResults.visitPercentage = indexType == KnnIndexTester.IndexType.IVF ? visitPercentage : 0;
         finalResults.avgRecall = checkResults(resultIds, nn, topK);
         finalResults.qps = (1000f * numQueryVectors) / elapsed;
         finalResults.avgLatency = (float) elapsed / numQueryVectors;
@@ -424,7 +424,8 @@ TopDocs doVectorQuery(float[] vector, IndexSearcher searcher, Query filterQuery,
         }
         int efSearch = Math.max(topK, this.efSearch);
         if (indexType == KnnIndexTester.IndexType.IVF) {
-            knnQuery = new IVFKnnFloatVectorQuery(VECTOR_FIELD, vector, topK, efSearch, filterQuery, nProbe);
+            float visitRatio = (float) (visitPercentage / 100);
+            knnQuery = new IVFKnnFloatVectorQuery(VECTOR_FIELD, vector, topK, efSearch, filterQuery, visitRatio);
         } else {
             knnQuery = new ESKnnFloatVectorQuery(
                 VECTOR_FIELD,

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

@@ -86,6 +86,67 @@ 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);
+            final int maxChildrenSize = centroids.readVInt();
+            final NeighborQueue currentParentQueue = new NeighborQueue(maxChildrenSize, true);
+            final int bufferSize = (int) Math.max(numCentroids * CENTROID_SAMPLING_PERCENTAGE, 1);
+            final NeighborQueue neighborQueue = new NeighborQueue(bufferSize, 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/IVFVectorsFormat.java

@@ -60,8 +60,8 @@ public class IVFVectorsFormat extends KnnVectorsFormat {
     );
 
     // This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
-    // useful when searching with 'efSearch' type parameters instead of requiring a specific nprobe.
-    public static final int DYNAMIC_NPROBE = -1;
+    // useful when searching with 'efSearch' type parameters instead of requiring a specific ratio.
+    public static final float DYNAMIC_VISIT_RATIO = 0.0f;
     public static final int DEFAULT_VECTORS_PER_CLUSTER = 384;
     public static final int MIN_VECTORS_PER_CLUSTER = 64;
     public static final int MAX_VECTORS_PER_CLUSTER = 1 << 16; // 65536