make centroid oversampling dependent on centroidsPerParentcluster

Author: iverase

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

@@ -41,10 +41,12 @@
 public class DefaultIVFVectorsReader extends IVFVectorsReader implements OffHeapStats {
 
     // How many extra centroids we need to collect for each visited centroid for hierarchical centroids.
-    public static final float CENTROID_OVERSAMPLING = 8.0f;
+    public final float centroidOversampling;
 
-    public DefaultIVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader) throws IOException {
+    public DefaultIVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader, int centroidsPerParentCluster)
+        throws IOException {
         super(state, rawVectorsReader);
+        centroidOversampling = (float) centroidsPerParentCluster / 2;
     }
 
     CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice) throws IOException {
@@ -124,7 +126,7 @@ CentroidIterator getCentroidIterator(
                 quantized,
                 queryParams,
                 globalCentroidDp,
-                visitRatio
+                visitRatio * centroidOversampling
             );
         } else {
             centroidIterator = getCentroidIteratorNoParent(
@@ -188,13 +190,13 @@ private static CentroidIterator getCentroidIteratorWithParents(
         byte[] quantizeQuery,
         OptimizedScalarQuantizer.QuantizationResult queryParams,
         float globalCentroidDp,
-        float visitRatio
+        float centroidRatio
     ) 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.min(Math.max(visitRatio * numCentroids * CENTROID_OVERSAMPLING, 1), numCentroids);
+        final int bufferSize = (int) Math.min(Math.max(centroidRatio * numCentroids, 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/IVFVectorsFormat.java

@@ -110,7 +110,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
 
     @Override
     public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new DefaultIVFVectorsReader(state, rawVectorFormat.fieldsReader(state));
+        return new DefaultIVFVectorsReader(state, rawVectorFormat.fieldsReader(state), centroidsPerParentCluster);
     }
 
     @Override