added exploration for a percentage of parents

Author: john-wagster

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

@@ -310,7 +310,7 @@ CentroidAssignments calculateAndWriteCentroids(
                     i++;
                 } else {
                     int label = kMeansResult.parentLayer()[i];
-                    int totalCentroids = 0;
+                    int centroidCount = 0;
                     float[] parentPartitionCentroid = new float[fieldInfo.getVectorDimension()];
                     int j = i;
                     for (; j < kMeansResult.parentLayer().length; j++) {
@@ -320,14 +320,14 @@ CentroidAssignments calculateAndWriteCentroids(
                         for (int k = 0; k < parentPartitionCentroid.length; k++) {
                             parentPartitionCentroid[k] += centroids[i][k];
                         }
-                        totalCentroids++;
+                        centroidCount++;
                     }
                     int childOrdinal = i;
                     i = j;
                     for (int d = 0; d < parentPartitionCentroid.length; d++) {
-                        parentPartitionCentroid[d] /= totalCentroids;
+                        parentPartitionCentroid[d] /= centroidCount;
                     }
-                    centroidPartitions.add(new CentroidPartition(parentPartitionCentroid, childOrdinal, totalCentroids));
+                    centroidPartitions.add(new CentroidPartition(parentPartitionCentroid, childOrdinal, centroidCount));
                 }
             }
         }

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

@@ -291,27 +291,31 @@ public final void search(String field, float[] target, KnnCollector knnCollector
         }
 
         while (parentCentroidQueue.size() > 0 && (centroidsVisited < nProbe || knnCollectorImpl.numCollected() < knnCollector.k())) {
-            int parentCentroidOrdinal = parentCentroidQueue.pop();
-
-            int childCentroidOrdinal;
-            int childCentroidCount;
-            if (parentCentroidOrdinal == -1) {
-                // score all centroids
-                childCentroidOrdinal = 0;
-                childCentroidCount = centroidQueryScorer.size();
-            } else {
-                childCentroidOrdinal = parentCentroidQueryScorer.getChildCentroidStart(parentCentroidOrdinal);
-                childCentroidCount = parentCentroidQueryScorer.getChildCount(parentCentroidOrdinal);
-            }
 
-            NeighborQueue centroidQueue = scorePostingLists(
-                fieldInfo,
-                knnCollector,
-                centroidQueryScorer,
-                nProbe,
-                childCentroidOrdinal,
-                childCentroidCount
-            );
+            NeighborQueue centroidQueue = new NeighborQueue(centroidQueryScorer.size(), true);;
+            int parentsToExplore = 0;
+            while(parentCentroidQueue.size() > 0 && parentsToExplore < parentCentroidQueryScorer.size() * 0.5) {
+                int parentCentroidOrdinal = parentCentroidQueue.pop();
+
+                int childCentroidOrdinal;
+                int childCentroidCount;
+                if (parentCentroidOrdinal == -1) {
+                    // score all centroids
+                    childCentroidOrdinal = 0;
+                    childCentroidCount = centroidQueryScorer.size();
+                } else {
+                    childCentroidOrdinal = parentCentroidQueryScorer.getChildCentroidStart(parentCentroidOrdinal);
+                    childCentroidCount = parentCentroidQueryScorer.getChildCount(parentCentroidOrdinal);
+                }
+                // FIXME: modify scorePostingLists to take a queue instead of creating one
+                centroidQueryScorer.bulkScore(centroidQueue, childCentroidOrdinal, childCentroidOrdinal + childCentroidCount);
+
+                if(parentCentroidOrdinal == -1) {
+                    break;
+                }
+
+                parentsToExplore++;
+            }
 
             PostingVisitor scorer = getPostingVisitor(fieldInfo, ivfClusters, target, needsScoring);
             // initially we visit only the "centroids to search"

server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/KMeansLocal.java

@@ -135,6 +135,7 @@ int getBestCentroidOffset(float[][] centroids, float[] vector, int centroidIdx,
                 }
             }
         }
+
         return bestCentroidOffset;
     }