include (parent) centroids scores for affinity, with larger segments

Author: tteofili

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

@@ -87,7 +87,7 @@ CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, Inde
     }
 
     @Override
-    public float[] getParentCentroidsScores(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] targetQuery)
+    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();
@@ -109,11 +109,11 @@ public float[] getParentCentroidsScores(FieldInfo fieldInfo, int numCentroids, I
         }
         final ES92Int7VectorsScorer scorer = ESVectorUtil.getES92Int7VectorsScorer(centroids, fieldInfo.getVectorDimension());
         centroids.seek(0L);
-        // score the parents
+        // final scores
         final float[] scores = new float[ES92Int7VectorsScorer.BULK_SIZE];
 
         int numParents = centroids.readVInt();
-        if (numParents > 0) {
+        if (parents && numParents > 0) {
             final NeighborQueue parentsQueue = new NeighborQueue(numParents, true);
             final int maxChildrenSize = centroids.readVInt();
             final NeighborQueue currentParentQueue = new NeighborQueue(maxChildrenSize, true);
@@ -130,6 +130,19 @@ public float[] getParentCentroidsScores(FieldInfo fieldInfo, int numCentroids, I
                 fieldInfo.getVectorSimilarityFunction(),
                 scores
             );
+        } else {
+            final NeighborQueue neighborQueue = new NeighborQueue(numCentroids, true);
+            score(
+                neighborQueue,
+                numCentroids,
+                0,
+                scorer,
+                quantizedQuery,
+                queryParams,
+                globalCentroidDp,
+                fieldInfo.getVectorSimilarityFunction(),
+                scores
+            );
         }
         return scores;
     }

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

@@ -91,7 +91,7 @@ protected IVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsR
     abstract CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target)
         throws IOException;
 
-    public abstract float[] getParentCentroidsScores(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target)
+    public abstract float[] getCentroidsScores(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target, boolean parents)
         throws IOException;
 
     private static IndexInput openDataInput(

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

@@ -132,19 +132,17 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
         List<LeafReaderContext> leafReaderContexts = reader.leaves();
 
         // calculate the affinity of each segment to the query vector
-        // (need information from each segment: no. of clusters, global centroid, density, whatever, ...)
+        // (need information from each segment: no. of clusters, global centroid, density, parent centroids' scores, etc.)
         List<SegmentAffinity> segmentAffinities = calculateSegmentAffinities(leafReaderContexts, getQueryVector());
 
         // TODO: sort segments by affinity score in descending order, and cut the long tail ?
-        
-        // with larger affinity we increase nprobe (and viceversa)
-        // also sort segments by affinity and eventually filter out the long tail
+
         List<LeafReaderContext> selectedSegments = new ArrayList<>();
 
-        // TODO : are these magic numbers ?
-        double cutoff_affinity = 0.01; // minimum affinity score for a segment to be considered
-        double higher_affinity = 0.6; // min affinity for increasing nProbe
-        double lower_affinity = 0.6; // max affinity for decreasing nProbe
+        double cutoff_affinity = 0.3; // minimum affinity score for a segment to be considered
+        double higher_affinity = 0.7; // min affinity for increasing nProbe
+        double lower_affinity = 0.6 ; // max affinity for decreasing nProbe
+
         int max_adjustment = 20;
 
         Map<LeafReaderContext, Integer> segmentNProbeMap = new HashMap<>();
@@ -244,15 +242,18 @@ private List<SegmentAffinity> calculateSegmentAffinities(List<LeafReaderContext>
                     int numCentroids = reader.getNumCentroids(fieldInfo);
                     double centroidDensity = (double) numCentroids / leafReader.numDocs();
 
-                    if (numCentroids > 64) {
-                        float[] parentCentroidsScores = reader.getParentCentroidsScores(
+                    // include some centroids' scores
+                    if (numCentroids > 32) {
+                        float[] centroidScores = reader.getCentroidsScores(
                             fieldInfo,
                             numCentroids,
                             reader.getIvfCentroids(fieldInfo),
-                            queryVector
+                            queryVector,
+                            numCentroids > 64
                         );
-                        Arrays.sort(parentCentroidsScores);
-                        globalCentroidScore = (parentCentroidsScores[0] + parentCentroidsScores[1] + globalCentroidScore) / 3;
+                        Arrays.sort(centroidScores);
+                        globalCentroidScore = (globalCentroidScore + centroidScores[centroidScores.length - 1]
+                            + centroidScores[centroidScores.length - 2]) / 3;
                     }
 
                     double affinityScore = globalCentroidScore * (1 + centroidDensity);