diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java index 3999622fdc52e..a38122596a9a8 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java @@ -18,7 +18,7 @@ import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.Bits; import org.apache.lucene.util.VectorUtil; -import org.apache.lucene.util.hnsw.NeighborQueue; +import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue; import org.elasticsearch.index.codec.vectors.reflect.OffHeapStats; import org.elasticsearch.simdvec.ES91OSQVectorsScorer; import org.elasticsearch.simdvec.ES92Int7VectorsScorer; @@ -243,22 +243,19 @@ public boolean hasNext() { @Override public CentroidOffsetAndLength nextPostingListOffsetAndLength() throws IOException { - int centroidOrdinal = neighborQueue.pop(); - updateQueue(); // add one children if available so the queue remains fully populated + int centroidOrdinal = nextCentroid(); centroids.seek(childrenFileOffsets + (long) Long.BYTES * 2 * centroidOrdinal); long postingListOffset = centroids.readLong(); long postingListLength = centroids.readLong(); return new CentroidOffsetAndLength(postingListOffset, postingListLength); } - private void updateQueue() throws IOException { + private int nextCentroid() throws IOException { if (currentParentQueue.size() > 0) { - // add a children from the current parent queue - float score = currentParentQueue.topScore(); - int children = currentParentQueue.pop(); - neighborQueue.add(children, score); + // return next centroid and maybe add a children from the current parent queue + return neighborQueue.popAndAddRaw(currentParentQueue.popRaw()); } else if (parentsQueue.size() > 0) { - // add a new parent from the current parent queue + // current parent queue is empty, populate it again with the next parent int pop = parentsQueue.pop(); populateOneChildrenGroup( currentParentQueue, @@ -273,7 +270,9 @@ private void updateQueue() throws IOException { globalCentroidDp, scores ); - updateQueue(); + return nextCentroid(); + } else { + return neighborQueue.pop(); } } }; diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/NeighborQueue.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/NeighborQueue.java index b0e261bb3aefb..737a7b078024a 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/NeighborQueue.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/NeighborQueue.java @@ -137,6 +137,25 @@ public int pop() { return decodeNodeId(heap.pop()); } + /** Removes the top element and returns it */ + public long popRaw() { + return heap.pop(); + } + + /** + * if the new element is the new top then return its node id. Otherwise, + * removes the current top element, returns its node id and adds the new element + * to the queue. + * */ + public int popAndAddRaw(long raw) { + long top = heap.top(); + if (raw < top) { + return decodeNodeId(raw); + } + heap.updateTop(raw); + return decodeNodeId(top); + } + public void clear() { heap.clear(); }