remove explicit budget as it conflicts with visitedRatio, normalized affinity as visited ratio modifier

Author: tteofili

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

@@ -138,63 +138,6 @@ CentroidIterator getCentroidIterator(
         return getPostingListPrefetchIterator(centroidIterator, postingListSlice);
     }
 
-    @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

@@ -93,9 +93,6 @@ abstract CentroidIterator getCentroidIterator(
         IndexInput postingListSlice
     ) 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,
@@ -272,9 +269,6 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             // is enough?
             expectedDocs += scorer.resetPostingsScorer(offsetAndLength.offset());
             actualDocs += scorer.visit(knnCollector);
-            if (knnCollector.earlyTerminated()) {
-                break;
-            }
         }
         if (acceptDocs != null) {
             float unfilteredRatioVisited = (float) expectedDocs / numVectors;
@@ -284,9 +278,6 @@ public final void search(String field, float[] target, KnnCollector knnCollector
                 CentroidOffsetAndLength offsetAndLength = centroidPrefetchingIterator.nextPostingListOffsetAndLength();
                 scorer.resetPostingsScorer(offsetAndLength.offset());
                 actualDocs += scorer.visit(knnCollector);
-                if (knnCollector.earlyTerminated()) {
-                    break;
-                }
             }
         }
     }

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

@@ -62,6 +62,9 @@ abstract class AbstractIVFKnnVectorQuery extends Query implements QueryProfilerP
 
     static final TopDocs NO_RESULTS = TopDocsCollector.EMPTY_TOPDOCS;
 
+    private static final float MAX_VISIT_INCREASE_RATIO = 0.1f;
+    private static final double MAX_VISIT_DECREASE_RATIO = 0.01f;
+
     protected final String field;
     protected final float providedVisitRatio;
     protected final int k;
@@ -166,79 +169,48 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
             visitRatio = providedVisitRatio;
         }
 
-        // FIXME: pick a reasonable min budget and make sure visitRatio is used appropriately throughout and not pushing the budget to zero
-        int totalBudget = Math.max(100, (int) (totalDocsWVectors * visitRatio));
-
         List<Callable<TopDocs>> tasks;
         if (leafReaderContexts.isEmpty() == false) {
             // calculate the affinity of each segment to the query vector
-            // (need information from each segment: no. of clusters, global centroid, density, parent centroids' scores, etc.)
             List<SegmentAffinity> segmentAffinities = calculateSegmentAffinities(leafReaderContexts, getQueryVector(), costs);
+            segmentAffinities.sort((a, b) -> Double.compare(b.affinityScore(), a.affinityScore()));
 
-            // TODO: sort segments by affinity score in descending order, and cut the long tail ?
             double[] affinityScores = segmentAffinities.stream()
                 .map(SegmentAffinity::affinityScore)
                 .mapToDouble(Double::doubleValue)
-                .toArray();
-
-            double[] filteredAffinityScores = Arrays.stream(affinityScores)
                 .filter(x -> Double.isNaN(x) == false && Double.isInfinite(x) == false)
                 .toArray();
 
-            final double averageAffinity = Arrays.stream(filteredAffinityScores).average().orElse(0.0);
-            double variance = Arrays.stream(filteredAffinityScores).map(x -> (x - averageAffinity) * (x - averageAffinity)).sum()
-                / filteredAffinityScores.length;
-            double stdDev = Math.sqrt(variance);
-
-            double maxAffinity;
-            double cutoffAffinity;
-            double affinityThreshold;
-
-            if (stdDev > averageAffinity) {
-                // adjust calculation when distribution is very skewed
-                double minAffinity = Arrays.stream(affinityScores).min().orElse(Double.NaN);
-                maxAffinity = Arrays.stream(affinityScores).max().orElse(Double.NaN);
-                double lowerAffinity = (minAffinity + averageAffinity) * 0.5;
-                cutoffAffinity = lowerAffinity * 0.1;
-                affinityThreshold = (minAffinity + lowerAffinity) * 0.66;
-            } else {
-                maxAffinity = averageAffinity + 50 * stdDev;
-                cutoffAffinity = averageAffinity - 50 * stdDev;
-                affinityThreshold = averageAffinity + stdDev;
-            }
+            double minAffinity = Arrays.stream(affinityScores).min().orElse(Double.NaN);
+            double maxAffinity = Arrays.stream(affinityScores).max().orElse(Double.NaN);
 
-            float maxAdjustments = visitRatio * 1.5f;
+            double[] normalizedAffinityScores = Arrays.stream(affinityScores)
+                .map(d -> (d - minAffinity) / (maxAffinity - minAffinity))
+                .toArray();
 
-            if (Double.isNaN(maxAffinity) || Double.isNaN(averageAffinity)) {
+            if (normalizedAffinityScores.length != segmentAffinities.size()) {
                 tasks = new ArrayList<>(leafReaderContexts.size());
                 for (LeafReaderContext context : leafReaderContexts) {
-                    tasks.add(() -> searchLeaf(context, filterWeight, knnCollectorManager, visitRatio, Integer.MAX_VALUE));
+                    tasks.add(() -> searchLeaf(context, filterWeight, knnCollectorManager, visitRatio));
                 }
             } else {
                 tasks = new ArrayList<>(segmentAffinities.size());
-                double scoreVectorsSum = segmentAffinities.stream().map(segmentAffinity -> {
-                    double affinity = Double.isNaN(segmentAffinity.affinityScore) ? maxAffinity : segmentAffinity.affinityScore;
-                    affinity = Math.clamp(affinity, 0.0, maxAffinity);
-                    return affinity * segmentAffinity.numVectors;
-                }).mapToDouble(Double::doubleValue).sum();
-
+                int j = 0;
                 for (SegmentAffinity segmentAffinity : segmentAffinities) {
-                    double score = segmentAffinity.affinityScore();
-                    if (score < cutoffAffinity) {
-                        continue;
-                    }
-                    float adjustedVisitRatio = adjustVisitRatioForSegment(score, affinityThreshold, maxAdjustments, visitRatio);
-                    LeafReaderContext context = segmentAffinity.context();
+                    double normalizedAffinityScore = normalizedAffinityScores[j];
 
-                    // distribute the budget according to : budgetᵢ = total_budget × (affinityᵢ × |vectors|ᵢ) / ∑ (affinityⱼ × |vectors|ⱼ)
-                    int segmentBudget = (int) (totalBudget * (score * segmentAffinity.numVectors) / scoreVectorsSum);
+                    float adjustedVisitRatio = (float) Math.clamp(
+                        (visitRatio * (MAX_VISIT_INCREASE_RATIO + normalizedAffinityScore)),
+                        visitRatio * MAX_VISIT_DECREASE_RATIO,
+                        visitRatio * (1 + MAX_VISIT_INCREASE_RATIO)
+                    );
 
-                    // TODO : should we always grant a min budget for each affine-enough segment
-                    if (segmentBudget > 0) {
-                        tasks.add(
-                            () -> searchLeaf(context, filterWeight, knnCollectorManager, adjustedVisitRatio, Math.max(1, segmentBudget))
-                        );
+                    LeafReaderContext context = segmentAffinity.context();
+
+                    if (adjustedVisitRatio > 0) {
+                        tasks.add(() -> searchLeaf(context, filterWeight, knnCollectorManager, adjustedVisitRatio));
                     }
+                    j++;
                 }
             }
         } else {
@@ -257,21 +229,6 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
 
     abstract VectorScorer createVectorScorer(LeafReaderContext context, FieldInfo fi) throws IOException;
 
-    private float adjustVisitRatioForSegment(double affinityScore, double affinityThreshold, float maxAdjustment, float visitRatio) {
-        // for high affinity scores, increase visited ratio
-        if (affinityScore > affinityThreshold) {
-            int adjustment = (int) Math.ceil((affinityScore - affinityThreshold) * maxAdjustment);
-            return Math.min(visitRatio * adjustment, visitRatio + maxAdjustment);
-        }
-
-        // for low affinity scores, decrease visited ratio
-        if (affinityScore <= affinityThreshold) {
-            return Math.max(visitRatio * 0.5f, 0.01f);
-        }
-
-        return visitRatio;
-    }
-
     abstract float[] getQueryVector() throws IOException;
 
     private IVFVectorsReader unwrapReader(KnnVectorsReader knnVectorsReader) {
@@ -287,8 +244,7 @@ private IVFVectorsReader unwrapReader(KnnVectorsReader knnVectorsReader) {
         return result;
     }
 
-    private List<SegmentAffinity> calculateSegmentAffinities(List<LeafReaderContext> leafReaderContexts, float[] queryVector, int[] costs)
-        throws IOException {
+    private List<SegmentAffinity> calculateSegmentAffinities(List<LeafReaderContext> leafReaderContexts, float[] queryVector, int[] costs) {
         List<SegmentAffinity> segmentAffinities = new ArrayList<>(leafReaderContexts.size());
 
         int i = 0;
@@ -325,25 +281,8 @@ private List<SegmentAffinity> calculateSegmentAffinities(List<LeafReaderContext>
                     int numCentroids = reader.getNumCentroids(fieldInfo);
                     double centroidDensity = (double) numCentroids / numVectors;
 
-                    // 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,
-                            numCentroids,
-                            reader.getIvfCentroids(fieldInfo),
-                            queryVector,
-                            numCentroids > 128
-                        );
-                        Arrays.sort(centroidScores);
-                        float first = centroidScores[centroidScores.length - 1];
-                        float second = centroidScores[centroidScores.length - 2];
-                        centroidsScore = (centroidsScore + first + second) / 3;
-                    }
-
+                    // TODO : we may want to include some actual centroids' scores for higher quality estimate
                     double affinityScore = centroidsScore * (1 + centroidDensity);
-
                     segmentAffinities.add(new SegmentAffinity(context, affinityScore, numVectors));
                 } else {
                     segmentAffinities.add(new SegmentAffinity(context, Float.NaN, 0));
@@ -357,14 +296,9 @@ private List<SegmentAffinity> calculateSegmentAffinities(List<LeafReaderContext>
 
     private record SegmentAffinity(LeafReaderContext context, double affinityScore, int numVectors) {}
 
-    private TopDocs searchLeaf(
-        LeafReaderContext ctx,
-        Weight filterWeight,
-        KnnCollectorManager knnCollectorManager,
-        float visitRatio,
-        int visitingBudget
-    ) throws IOException {
-        TopDocs results = getLeafResults(ctx, filterWeight, knnCollectorManager, visitRatio, visitingBudget);
+    private TopDocs searchLeaf(LeafReaderContext ctx, Weight filterWeight, KnnCollectorManager knnCollectorManager, float visitRatio)
+        throws IOException {
+        TopDocs results = getLeafResults(ctx, filterWeight, knnCollectorManager, visitRatio);
         IntHashSet dedup = new IntHashSet(results.scoreDocs.length * 4 / 3);
         int deduplicateCount = 0;
         for (ScoreDoc scoreDoc : results.scoreDocs) {
@@ -384,19 +318,13 @@ private TopDocs searchLeaf(
         return new TopDocs(results.totalHits, deduplicatedScoreDocs);
     }
 
-    TopDocs getLeafResults(
-        LeafReaderContext ctx,
-        Weight filterWeight,
-        KnnCollectorManager knnCollectorManager,
-        float visitRatio,
-        int visitingBudget
-    ) throws IOException {
+    TopDocs getLeafResults(LeafReaderContext ctx, Weight filterWeight, KnnCollectorManager knnCollectorManager, float visitRatio)
+        throws IOException {
         final LeafReader reader = ctx.reader();
         final Bits liveDocs = reader.getLiveDocs();
 
-        KnnSearchStrategy searchStrategy = new IVFKnnSearchStrategy(visitRatio);
         if (filterWeight == null) {
-            return approximateSearch(ctx, liveDocs, visitingBudget, knnCollectorManager, searchStrategy);
+            return approximateSearch(ctx, liveDocs, Integer.MAX_VALUE, knnCollectorManager, visitRatio);
         }
 
         Scorer scorer = filterWeight.scorer(ctx);
@@ -406,15 +334,15 @@ TopDocs getLeafResults(
 
         BitSet acceptDocs = createBitSet(scorer.iterator(), liveDocs, reader.maxDoc());
         final int cost = acceptDocs.cardinality();
-        return approximateSearch(ctx, acceptDocs, Math.min(visitingBudget, cost + 1), knnCollectorManager, searchStrategy);
+        return approximateSearch(ctx, acceptDocs, cost + 1, knnCollectorManager, visitRatio);
     }
 
     abstract TopDocs approximateSearch(
         LeafReaderContext context,
         Bits acceptDocs,
         int visitedLimit,
         KnnCollectorManager knnCollectorManager,
-        KnnSearchStrategy searchStrategy
+        float visitRatio
     ) throws IOException;
 
     protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {

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

@@ -92,7 +92,7 @@ protected TopDocs approximateSearch(
         Bits acceptDocs,
         int visitedLimit,
         KnnCollectorManager knnCollectorManager,
-        KnnSearchStrategy searchStrategy
+        float visitRatio
     ) throws IOException {
         LeafReader reader = context.reader();
         FloatVectorValues floatVectorValues = reader.getFloatVectorValues(field);
@@ -103,7 +103,8 @@ protected TopDocs approximateSearch(
         if (floatVectorValues.size() == 0) {
             return NO_RESULTS;
         }
-        KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, searchStrategy, context);
+        KnnSearchStrategy strategy = new IVFKnnSearchStrategy(visitRatio);
+        KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, strategy, context);
         if (knnCollector == null) {
             return NO_RESULTS;
         }