iterated on better mechanism for utilizing parent centroids

Author: john-wagster

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

@@ -126,18 +126,12 @@ public float[] centroid(int centroidOrdinal) throws IOException {
                 return centroid;
             }
 
-            @Override
-            public void bulkScore(NeighborQueue queue) throws IOException {
-                // TODO: bulk score centroids like we do with posting lists
-                centroids.seek(quantizedCentroidsOffset);
-                for (int i = 0; i < numCentroids; i++) {
-                    queue.add(i, score());
-                }
-            }
-
             @Override
             public void bulkScore(NeighborQueue queue, int start, int end) throws IOException {
                 // TODO: bulk score centroids like we do with posting lists
+                assert start > 0;
+                assert end > 0;
+                assert start + end <= numCentroids;
                 centroids.seek(quantizedCentroidsOffset + quantizedVectorByteSize * start);
                 for (int i = start; i < end; i++) {
                     queue.add(i, score());
@@ -222,17 +216,11 @@ public int getChildCount(int centroidOrdinal) throws IOException {
                 return childCount;
             }
 
-            @Override
-            public void bulkScore(NeighborQueue queue) throws IOException {
-                // TODO: bulk score centroids like we do with posting lists
-                centroids.seek(0L);
-                for (int i = 0; i < numParentCentroids; i++) {
-                    queue.add(i, score());
-                }
-            }
-
             @Override
             public void bulkScore(NeighborQueue queue, int start, int end) throws IOException {
+                assert start > 0;
+                assert end > 0;
+                assert start + end <= numParentCentroids;
                 // TODO: bulk score centroids like we do with posting lists
                 centroids.seek(parentNodeByteSize * start);
                 for (int i = start; i < end; i++) {

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

@@ -290,45 +290,35 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             parentCentroidQueue.add(-1, 0.f);
         }
 
-        int bulkParentThreshold = (int) Math.ceil(parentCentroidQueue.size() * 0.5);
-
         while (parentCentroidQueue.size() > 0 && (centroidsVisited < nProbe || knnCollectorImpl.numCollected() < knnCollector.k())) {
             NeighborQueue centroidQueue = new NeighborQueue(centroidQueryScorer.size(), true);
-            int parentsToExplore = 0;
-            while (parentCentroidQueue.size() > 0 && parentsToExplore < bulkParentThreshold) {
-                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++;
-            }
+            updateCentroidQueueWNextParent(parentCentroidQueryScorer, parentCentroidQueue, centroidQueryScorer, centroidQueue);
 
             PostingVisitor scorer = getPostingVisitor(fieldInfo, ivfClusters, target, needsScoring);
             // initially we visit only the "centroids to search"
             // Note, numCollected is doing the bare minimum here.
             // TODO do we need to handle nested doc counts similarly to how we handle
             // filtering? E.g. keep exploring until we hit an expected number of parent documents vs. child vectors?
+            float nextParentDistance = Float.MAX_VALUE;
+            if (parentCentroidQueue.size() > 0) {
+                nextParentDistance = parentCentroidQueue.topScore();
+            }
             while (centroidQueue.size() > 0 && (centroidsVisited < nProbe || knnCollectorImpl.numCollected() < knnCollector.k())) {
                 ++centroidsVisited;
-                // todo do we actually need to know the score???
+                float centroidDistance = centroidQueue.topScore();
+                // the next parent likely contains centroids we need to evaluate prior to evaluating this next centroid
+                while (parentCentroidQueue.size() > 0 && centroidDistance > nextParentDistance) {
+                    updateCentroidQueueWNextParent(parentCentroidQueryScorer, parentCentroidQueue, centroidQueryScorer, centroidQueue);
+                    if (parentCentroidQueue.size() > 0) {
+                        nextParentDistance = parentCentroidQueue.topScore();
+                    } else {
+                        nextParentDistance = Float.MAX_VALUE;
+                    }
+                    centroidDistance = centroidQueue.topScore();
+                }
+
                 int centroidOrdinal = centroidQueue.pop();
-                // todo do we need direct access to the raw centroid???, this is used for quantizing, maybe hydrating and quantizing
-                // is enough?
+                // TODO need direct access to the raw centroid???, this is used for quantizing, maybe hydrating and quantizing is enough?
                 expectedDocs += scorer.resetPostingsScorer(centroidOrdinal, centroidQueryScorer.centroid(centroidOrdinal));
                 actualDocs += scorer.visit(knnCollector);
             }
@@ -346,6 +336,28 @@ public final void search(String field, float[] target, KnnCollector knnCollector
         }
     }
 
+    private static void updateCentroidQueueWNextParent(
+        CentroidWChildrenQueryScorer parentCentroidQueryScorer,
+        NeighborQueue parentCentroidQueue,
+        CentroidQueryScorer centroidQueryScorer,
+        NeighborQueue centroidQueue
+    ) throws IOException {
+        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);
+        }
+        // TODO: add back scorePostingLists? seems like it's not doing anything at this point
+        centroidQueryScorer.bulkScore(centroidQueue, childCentroidOrdinal, childCentroidOrdinal + childCentroidCount);
+    }
+
     @Override
     public final void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
         final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field);
@@ -409,8 +421,6 @@ interface CentroidQueryScorer {
 
         float[] centroid(int centroidOrdinal) throws IOException;
 
-        void bulkScore(NeighborQueue queue) throws IOException;
-
         void bulkScore(NeighborQueue queue, int start, int end) throws IOException;
     }