Make centroid oversampling depend on the visit ratio

Author: iverase

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

@@ -40,8 +40,8 @@
  */
 public class DefaultIVFVectorsReader extends IVFVectorsReader implements OffHeapStats {
 
-    // The percentage of centroids that are scored to keep recall
-    public static final double CENTROID_SAMPLING_PERCENTAGE = 0.2;
+    // How many extra centroids we need to collect for each visited centroid for hierarchical centroids.
+    public static final float CENTROID_OVERSAMPLING = 5.0f;
 
     public DefaultIVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader) throws IOException {
         super(state, rawVectorsReader);
@@ -89,7 +89,8 @@ CentroidIterator getCentroidIterator(
         int numCentroids,
         IndexInput centroids,
         float[] targetQuery,
-        IndexInput postingListSlice
+        IndexInput postingListSlice,
+        float visitRatio
     ) throws IOException {
         final FieldEntry fieldEntry = fields.get(fieldInfo.number);
         final float globalCentroidDp = fieldEntry.globalCentroidDp();
@@ -122,7 +123,8 @@ CentroidIterator getCentroidIterator(
                 scorer,
                 quantized,
                 queryParams,
-                globalCentroidDp
+                globalCentroidDp,
+                visitRatio
             );
         } else {
             centroidIterator = getCentroidIteratorNoParent(
@@ -185,13 +187,14 @@ private static CentroidIterator getCentroidIteratorWithParents(
         ES92Int7VectorsScorer scorer,
         byte[] quantizeQuery,
         OptimizedScalarQuantizer.QuantizationResult queryParams,
-        float globalCentroidDp
+        float globalCentroidDp,
+        float visitRatio
     ) throws IOException {
         // build the three queues we are going to use
         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 int bufferSize = (int) Math.min(Math.max(visitRatio * numCentroids * CENTROID_OVERSAMPLING, 1), numCentroids);
         final NeighborQueue neighborQueue = new NeighborQueue(bufferSize, true);
         // score the parents
         final float[] scores = new float[ES92Int7VectorsScorer.BULK_SIZE];

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

@@ -90,7 +90,8 @@ abstract CentroidIterator getCentroidIterator(
         int numCentroids,
         IndexInput centroids,
         float[] target,
-        IndexInput postingListSlice
+        IndexInput postingListSlice,
+        float visitRatio
     ) throws IOException;
 
     private static IndexInput openDataInput(
@@ -252,7 +253,8 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             entry.numCentroids,
             entry.centroidSlice(ivfCentroids),
             target,
-            postListSlice
+            postListSlice,
+            visitRatio
         );
         PostingVisitor scorer = getPostingVisitor(fieldInfo, postListSlice, target, acceptDocs);
         long expectedDocs = 0;