iter

Author: john-wagster

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

@@ -47,10 +47,10 @@ public DefaultIVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVect
         super(state, rawVectorsReader);
     }
 
-    @Override
     CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice) throws IOException {
         return new CentroidIterator() {
-            CentroidOffsetAndLength nextOffsetAndLength = null;
+            CentroidOffsetAndLength nextOffsetAndLength =
+                centroidIterator.hasNext() ? centroidIterator.nextPostingListOffsetAndLength() : null;
 
             private void prefetch(CentroidOffsetAndLength offsetAndLength) throws IOException {
                 postingListSlice.prefetch(offsetAndLength.offset(), offsetAndLength.length());
@@ -64,30 +64,21 @@ public boolean hasNext() {
 
             @Override
             public CentroidOffsetAndLength nextPostingListOffsetAndLength() throws IOException {
-                CentroidOffsetAndLength offsetAndLength;
-                // first time we need to fetch two if possible
-                if (nextOffsetAndLength == null) {
-                    offsetAndLength = centroidIterator.nextPostingListOffsetAndLength();
-                    if (centroidIterator.hasNext()) {
-                        nextOffsetAndLength = centroidIterator.nextPostingListOffsetAndLength();
-                        prefetch(nextOffsetAndLength);
-                    }
+                CentroidOffsetAndLength offsetAndLength = nextOffsetAndLength;
+                if (centroidIterator.hasNext()) {
+                    nextOffsetAndLength = centroidIterator.nextPostingListOffsetAndLength();
+                    prefetch(nextOffsetAndLength);
                 } else {
-                    offsetAndLength = nextOffsetAndLength;
-                    if (centroidIterator.hasNext()) {
-                        nextOffsetAndLength = centroidIterator.nextPostingListOffsetAndLength();
-                        prefetch(nextOffsetAndLength);
-                    } else {
-                        nextOffsetAndLength = null;  // indicate we reached the end
-                    }
+                    nextOffsetAndLength = null;  // indicate we reached the end
                 }
                 return offsetAndLength;
             }
         };
     }
 
     @Override
-    CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] targetQuery)
+    CentroidIterator getCentroidIterator(
+        FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] targetQuery, IndexInput postingListSlice)
         throws IOException {
         final FieldEntry fieldEntry = fields.get(fieldInfo.number);
         final float globalCentroidDp = fieldEntry.globalCentroidDp();
@@ -110,8 +101,9 @@ CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, Inde
         final ES92Int7VectorsScorer scorer = ESVectorUtil.getES92Int7VectorsScorer(centroids, fieldInfo.getVectorDimension());
         centroids.seek(0L);
         int numParents = centroids.readVInt();
+        CentroidIterator centroidIterator;
         if (numParents > 0) {
-            return getCentroidIteratorWithParents(
+            centroidIterator = getCentroidIteratorWithParents(
                 fieldInfo,
                 centroids,
                 numParents,
@@ -121,8 +113,10 @@ CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, Inde
                 queryParams,
                 globalCentroidDp
             );
+        } else {
+            centroidIterator = getCentroidIteratorNoParent(fieldInfo, centroids, numCentroids, scorer, quantized, queryParams, globalCentroidDp);
         }
-        return getCentroidIteratorNoParent(fieldInfo, centroids, numCentroids, scorer, quantized, queryParams, globalCentroidDp);
+        return getPostingListPrefetchIterator(centroidIterator, postingListSlice);
     }
 
     private static CentroidIterator getCentroidIteratorNoParent(

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

@@ -85,10 +85,8 @@ protected IVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsR
         }
     }
 
-    abstract CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice)
-        throws IOException;
-
-    abstract CentroidIterator getCentroidIterator(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target)
+    abstract CentroidIterator getCentroidIterator(
+        FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] target, IndexInput postingListSlice)
         throws IOException;
 
     private static IndexInput openDataInput(
@@ -245,10 +243,8 @@ public final void search(String field, float[] target, KnnCollector knnCollector
         // we account for soar vectors here. We can potentially visit a vector twice so we multiply by 2 here.
         long maxVectorVisited = (long) (2.0 * visitRatio * numVectors);
         IndexInput postListSlice = entry.postingListSlice(ivfClusters);
-        CentroidIterator centroidPrefetchingIterator = getPostingListPrefetchIterator(
-            getCentroidIterator(fieldInfo, entry.numCentroids, entry.centroidSlice(ivfCentroids), target),
-            postListSlice
-        );
+        CentroidIterator centroidPrefetchingIterator = getCentroidIterator(
+            fieldInfo, entry.numCentroids, entry.centroidSlice(ivfCentroids), target, postListSlice);
         PostingVisitor scorer = getPostingVisitor(fieldInfo, postListSlice, target, acceptDocs);
         long expectedDocs = 0;
         long actualDocs = 0;