iter

Author: benwtrent

server/src/main/java/org/elasticsearch/search/vectors/AbstractMaxScoreKnnCollector.java

@@ -16,7 +16,6 @@
 public abstract class AbstractMaxScoreKnnCollector extends AbstractKnnCollector {
     public static final long LEAST_COMPETITIVE = NeighborQueue.encodeRaw(Integer.MAX_VALUE, Float.NEGATIVE_INFINITY);
 
-
     protected AbstractMaxScoreKnnCollector(int k, long visitLimit, KnnSearchStrategy searchStrategy) {
         super(k, visitLimit, searchStrategy);
     }

server/src/main/java/org/elasticsearch/search/vectors/DiversifyingNearestChildrenKnnCollector.java

@@ -26,7 +26,7 @@
 import org.apache.lucene.search.knn.KnnSearchStrategy;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BitSet;
-import org.apache.lucene.util.NumericUtils;
+import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue;
 
 /**
  * This collects the nearest children vectors. Diversifying the results over the provided parent
@@ -36,7 +36,7 @@ class DiversifyingNearestChildrenKnnCollector extends AbstractMaxScoreKnnCollect
 
     private final BitSet parentBitSet;
     private final NodeIdCachingHeap heap;
-    private long minCompetitiveDocScore = Long.MIN_VALUE;
+    private long minCompetitiveDocScore = LEAST_COMPETITIVE;
     private float minCompetitiveScore = Float.NEGATIVE_INFINITY;
 
     /**
@@ -74,10 +74,7 @@ public boolean collect(int docId, float nodeScore) {
 
     @Override
     public float minCompetitiveSimilarity() {
-        if (heap.size >= k() || minCompetitiveDocScore > Long.MIN_VALUE) {
-            return Math.max(minCompetitiveScore, heap.topScore());
-        }
-        return Float.NEGATIVE_INFINITY;
+        return heap.size < k() ? Float.NEGATIVE_INFINITY : Math.max(minCompetitiveScore, heap.topScore());
     }
 
     @Override
@@ -108,17 +105,14 @@ public int numCollected() {
 
     @Override
     public long getMinCompetitiveDocScore() {
-        return  (((long) NumericUtils.floatToSortableInt(heap.topScore())) << 32) | (0xFFFFFFFFL & ~heap.topNode());
+        return heap.size() > 0 ? Math.max(NeighborQueue.encodeRaw(heap.topNode(), heap.topScore()), minCompetitiveDocScore) : minCompetitiveDocScore;
     }
 
     @Override
     void updateMinCompetitiveDocScore(long minCompetitiveDocScore) {
-        if (minCompetitiveDocScore == Long.MIN_VALUE) {
-            // no update
-            return;
-        }
-        this.minCompetitiveDocScore = minCompetitiveDocScore;
-        this.minCompetitiveScore = NumericUtils.sortableIntToFloat((int) (minCompetitiveDocScore >> 32));
+        long queueMinCompetitiveDocScore = heap.size() > 0 ? NeighborQueue.encodeRaw(heap.topNode(), heap.topScore()) : LEAST_COMPETITIVE;
+        this.minCompetitiveDocScore = Math.max(this.minCompetitiveDocScore, Math.max(queueMinCompetitiveDocScore, minCompetitiveDocScore));
+        this.minCompetitiveScore = NeighborQueue.decodeScoreRaw(this.minCompetitiveDocScore);
     }
 
     /**
@@ -154,10 +148,16 @@ private static class NodeIdCachingHeap {
         }
 
         public final int topNode() {
+            if (size == 0) {
+                return Integer.MAX_VALUE;
+            }
             return heapNodes[1].child;
         }
 
         public final float topScore() {
+            if (size == 0) {
+                return Float.NEGATIVE_INFINITY;
+            }
             return heapNodes[1].score;
         }
 

server/src/main/java/org/elasticsearch/search/vectors/IVFKnnSearchStrategy.java

@@ -10,17 +10,13 @@
 
 import org.apache.lucene.search.knn.KnnSearchStrategy;
 import org.apache.lucene.util.SetOnce;
-import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue;
-import org.elasticsearch.logging.LogManager;
-import org.elasticsearch.logging.Logger;
 
 import java.util.Objects;
 import java.util.concurrent.atomic.LongAccumulator;
 
 public class IVFKnnSearchStrategy extends KnnSearchStrategy {
-    private static final Logger logger = LogManager.getLogger(IVFKnnSearchStrategy.class);
     private final float visitRatio;
-    private SetOnce<AbstractMaxScoreKnnCollector> collector = new SetOnce<>();
+    private final SetOnce<AbstractMaxScoreKnnCollector> collector = new SetOnce<>();
     private final LongAccumulator accumulator;
 
     IVFKnnSearchStrategy(float visitRatio, LongAccumulator accumulator) {