minor tweaks, add knn tester param

Author: tteofili

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

@@ -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");
@@ -81,6 +82,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);
@@ -115,6 +117,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
@@ -149,6 +152,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();
     }
 
@@ -183,6 +187,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()) {
@@ -313,6 +318,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");
@@ -347,7 +357,8 @@ public CmdLineArgs build() {
                 vectorEncoding,
                 dimensions,
                 earlyTermination,
-                mergePolicy
+                mergePolicy,
+                vectorsRatio
             );
         }
     }

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

@@ -115,6 +115,7 @@ class KnnSearcher {
     private final float overSamplingFactor;
     private final int searchThreads;
     private final int numSearchers;
+    private final float vectorsRatio;
 
     KnnSearcher(Path indexPath, CmdLineArgs cmdLineArgs, int nProbe) {
         this.docPath = cmdLineArgs.docVectors();
@@ -137,6 +138,7 @@ class KnnSearcher {
         this.numSearchers = cmdLineArgs.numSearchers();
         this.randomSeed = cmdLineArgs.seed();
         this.selectivity = cmdLineArgs.filterSelectivity();
+        this.vectorsRatio = cmdLineArgs.vectorsRatio();
     }
 
     void runSearch(KnnIndexTester.Results finalResults, boolean earlyTermination) throws IOException {
@@ -424,7 +426,7 @@ 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);
+            knnQuery = new IVFKnnFloatVectorQuery(VECTOR_FIELD, vector, topK, efSearch, filterQuery, nProbe, vectorsRatio);
         } else {
             knnQuery = new ESKnnFloatVectorQuery(
                 VECTOR_FIELD,

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

@@ -257,6 +257,7 @@ public final void search(String field, float[] target, KnnCollector knnCollector
         // TODO do we need to handle nested doc counts similarly to how we handle
         // filtering? E.g. keep exploring until we hit an expected number of parent documents vs. child vectors?
         while (centroidIterator.hasNext() && (centroidsVisited < nProbe || knnCollectorImpl.numCollected() < knnCollector.k())) {
+            // TODO : check previous centroid max score, and exit?
             ++centroidsVisited;
             // todo do we actually need to know the score???
             long offset = centroidIterator.nextPostingListOffset();

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -2748,9 +2748,10 @@ && isNotUnitVector(squaredMagnitude)) {
                         numCands,
                         filter,
                         parentFilter,
-                        bbqIndexOptions.defaultNProbe
+                        bbqIndexOptions.defaultNProbe,
+                        1f
                     )
-                    : new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe);
+                    : new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe, 1);
             } else {
                 knnQuery = parentFilter != null
                     ? new ESDiversifyingChildrenFloatKnnVectorQuery(

server/src/main/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQuery.java

@@ -60,16 +60,15 @@ abstract class AbstractIVFKnnVectorQuery extends Query implements QueryProfilerP
 
     static final TopDocs NO_RESULTS = TopDocsCollector.EMPTY_TOPDOCS;
 
-    private static final double VECTOR_VISITED_PERCENTAGE_BUDGET = 0.05;
-
     protected final String field;
     protected final int nProbe;
     protected final int k;
     protected final int numCands;
     protected final Query filter;
+    private final float visitedRatio;
     protected int vectorOpsCount;
 
-    protected AbstractIVFKnnVectorQuery(String field, int nProbe, int k, int numCands, Query filter) {
+    protected AbstractIVFKnnVectorQuery(String field, int nProbe, int k, int numCands, float visitedRatio, Query filter) {
         if (k < 1) {
             throw new IllegalArgumentException("k must be at least 1, got: " + k);
         }
@@ -84,6 +83,7 @@ protected AbstractIVFKnnVectorQuery(String field, int nProbe, int k, int numCand
         this.k = k;
         this.filter = filter;
         this.numCands = numCands;
+        this.visitedRatio = visitedRatio;
     }
 
     @Override
@@ -132,7 +132,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
         TaskExecutor taskExecutor = indexSearcher.getTaskExecutor();
         List<LeafReaderContext> leafReaderContexts = reader.leaves();
 
-        int totalBudget = (int) (reader.numDocs() * VECTOR_VISITED_PERCENTAGE_BUDGET);
+        int totalBudget = (int) (reader.numDocs() * visitedRatio);
 
         List<Callable<TopDocs>> tasks;
         if (leafReaderContexts.isEmpty() == false) {
@@ -146,12 +146,12 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
                 .mapToDouble(Double::doubleValue)
                 .toArray();
 
-            // max affinity for decreasing nProbe
             double averageAffinity = Arrays.stream(affinityScores).average().orElse(Double.NaN);
+            // max affinity for decreasing nProbe
             double maxAffinity = Arrays.stream(affinityScores).max().orElse(Double.NaN);
             double lowerAffinity = (maxAffinity + averageAffinity) * 0.5;
             double cutoffAffinity = lowerAffinity * 0.5; // minimum affinity score for a segment to be considered
-            double affinityTreshold = (maxAffinity + lowerAffinity) * 0.66; // min affinity for increasing nProbe
+            double affinityThreshold = (maxAffinity + lowerAffinity) * 0.66; // min affinity for increasing nProbe
             int maxAdjustments = (int) (nProbe * 1.5);
 
             if (Double.isNaN(maxAffinity) || Double.isNaN(averageAffinity)) {
@@ -171,10 +171,10 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
                     if (score < cutoffAffinity) {
                         continue;
                     }
-                    int adjustedNProbe = adjustNProbeForSegment(score, affinityTreshold, maxAdjustments);
+                    int adjustedNProbe = adjustNProbeForSegment(score, affinityThreshold, maxAdjustments);
                     LeafReaderContext context = segmentAffinity.context();
 
-                    // budgetᵢ = total_budget × (affinityᵢ × |vectors|ᵢ) / ∑ (affinityⱼ × |vectors|ⱼ)
+                    // distribute the budget according to : budgetᵢ = total_budget × (affinityᵢ × |vectors|ᵢ) / ∑ (affinityⱼ × |vectors|ⱼ)
                     int segmentBudget = (int) (totalBudget * (score * context.reader().numDocs()) / scoreVectorsSum);
                     tasks.add(() -> searchLeaf(context, filterWeight, knnCollectorManager, adjustedNProbe, Math.max(1, segmentBudget)));
                 }
@@ -257,6 +257,7 @@ private List<SegmentAffinity> calculateSegmentAffinities(List<LeafReaderContext>
 
                     // with larger clusters, global centroid might not be a good representative,
                     // so we want to include "some" centroids' scores for higher quality estimate
+                    // TODO: tweak the threshold numCentroids here
                     if (numCentroids > 64) {
                         float[] centroidScores = reader.getCentroidsScores(
                             fieldInfo,

server/src/main/java/org/elasticsearch/search/vectors/DiversifyingChildrenIVFKnnFloatVectorQuery.java

@@ -38,9 +38,10 @@ public DiversifyingChildrenIVFKnnFloatVectorQuery(
         int numCands,
         Query childFilter,
         BitSetProducer parentsFilter,
-        int nProbe
+        int nProbe,
+        float visitedRatio
     ) {
-        super(field, query, k, numCands, childFilter, nProbe);
+        super(field, query, k, numCands, childFilter, nProbe, visitedRatio);
         this.parentsFilter = parentsFilter;
     }
 

server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java

@@ -35,8 +35,8 @@ public class IVFKnnFloatVectorQuery extends AbstractIVFKnnVectorQuery {
      * @param filter the filter to apply to the results
      * @param nProbe the number of probes to use for the IVF search strategy
      */
-    public IVFKnnFloatVectorQuery(String field, float[] query, int k, int numCands, Query filter, int nProbe) {
-        super(field, nProbe, k, numCands, filter);
+    public IVFKnnFloatVectorQuery(String field, float[] query, int k, int numCands, Query filter, int nProbe, float visitedRatio) {
+        super(field, nProbe, k, numCands, visitedRatio, filter);
         this.query = query;
     }
 

server/src/test/java/org/elasticsearch/search/vectors/DiversifyingChildrenIVFKnnFloatVectorQueryTests.java

@@ -18,7 +18,16 @@ public class DiversifyingChildrenIVFKnnFloatVectorQueryTests extends AbstractDiv
 
     @Override
     Query getDiversifyingChildrenKnnQuery(String fieldName, float[] queryVector, Query childFilter, int k, BitSetProducer parentBitSet) {
-        return new DiversifyingChildrenIVFKnnFloatVectorQuery(fieldName, queryVector, k, k, childFilter, parentBitSet, -1);
+        return new DiversifyingChildrenIVFKnnFloatVectorQuery(
+            fieldName,
+            queryVector,
+            k,
+            k,
+            childFilter,
+            parentBitSet,
+            -1,
+            1f
+        );
     }
 
     @Override

server/src/test/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQueryTests.java

@@ -27,7 +27,7 @@ public class IVFKnnFloatVectorQueryTests extends AbstractIVFKnnVectorQueryTestCa
 
     @Override
     IVFKnnFloatVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter, int nProbe) {
-        return new IVFKnnFloatVectorQuery(field, query, k, k, queryFilter, nProbe);
+        return new IVFKnnFloatVectorQuery(field, query, k, k, queryFilter, nProbe, 1);
     }
 
     @Override