address review comments

Author: iverase

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

@@ -210,16 +210,6 @@ private static int getBestCentroid(float[][] centroids, float[] vector, float[]
         return bestCentroidOffset;
     }
 
-    private NeighborHood[] computeNeighborhoods(float[][] centers, int clustersPerNeighborhood) throws IOException {
-        assert centers.length > clustersPerNeighborhood;
-        // experiments shows that below 15k, we better use brute force, otherwise hnsw gives us a nice speed up
-        if (centers.length < 15_000) {
-            return NeighborHood.computeNeighborhoodsBruteForce(centers, clustersPerNeighborhood);
-        } else {
-            return NeighborHood.computeNeighborhoodsGraph(centers, clustersPerNeighborhood);
-        }
-    }
-
     private void assignSpilled(
         FloatVectorValues vectors,
         KMeansIntermediate kmeansIntermediate,
@@ -350,7 +340,7 @@ private void doCluster(FloatVectorValues vectors, KMeansIntermediate kMeansInter
         NeighborHood[] neighborhoods = null;
         // if there are very few centroids, don't bother with neighborhoods or neighbor aware clustering
         if (neighborAware && centroids.length > clustersPerNeighborhood) {
-            neighborhoods = computeNeighborhoods(centroids, clustersPerNeighborhood);
+            neighborhoods = NeighborHood.computeNeighborhoods(centroids, clustersPerNeighborhood);
         }
         cluster(vectors, kMeansIntermediate, neighborhoods);
         if (neighborAware && soarLambda >= 0) {

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

@@ -11,8 +11,8 @@
 
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.search.KnnCollector;
-import org.apache.lucene.search.ScoreDoc;
-import org.apache.lucene.util.Bits;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
 import org.apache.lucene.util.VectorUtil;
 import org.apache.lucene.util.hnsw.HnswGraphBuilder;
 import org.apache.lucene.util.hnsw.HnswGraphSearcher;
@@ -25,8 +25,66 @@
 
 public record NeighborHood(int[] neighbors, float maxIntraDistance) {
 
+    private static final int M = 8;
+    private static final int EF_CONSTRUCTION = 150;
+
     static final NeighborHood EMPTY = new NeighborHood(new int[0], Float.POSITIVE_INFINITY);
 
+    public static NeighborHood[] computeNeighborhoods(float[][] centers, int clustersPerNeighborhood) throws IOException {
+        assert centers.length > clustersPerNeighborhood;
+        // experiments shows that below 15k, we better use brute force, otherwise hnsw gives us a nice speed up
+        if (centers.length < 15_000) {
+            return computeNeighborhoodsBruteForce(centers, clustersPerNeighborhood);
+        } else {
+            return computeNeighborhoodsGraph(centers, clustersPerNeighborhood);
+        }
+    }
+
+    public static NeighborHood[] computeNeighborhoodsBruteForce(float[][] centers, int clustersPerNeighborhood) {
+        int k = centers.length;
+        NeighborQueue[] neighborQueues = new NeighborQueue[k];
+        for (int i = 0; i < k; i++) {
+            neighborQueues[i] = new NeighborQueue(clustersPerNeighborhood, true);
+        }
+        final float[] scores = new float[4];
+        final int limit = k - 3;
+        for (int i = 0; i < k - 1; i++) {
+            float[] center = centers[i];
+            int j = i + 1;
+            for (; j < limit; j += 4) {
+                ESVectorUtil.squareDistanceBulk(center, centers[j], centers[j + 1], centers[j + 2], centers[j + 3], scores);
+                for (int h = 0; h < 4; h++) {
+                    neighborQueues[j + h].insertWithOverflow(i, scores[h]);
+                    neighborQueues[i].insertWithOverflow(j + h, scores[h]);
+                }
+            }
+            for (; j < k; j++) {
+                float dsq = VectorUtil.squareDistance(center, centers[j]);
+                neighborQueues[j].insertWithOverflow(i, dsq);
+                neighborQueues[i].insertWithOverflow(j, dsq);
+            }
+        }
+
+        NeighborHood[] neighborhoods = new NeighborHood[k];
+        for (int i = 0; i < k; i++) {
+            NeighborQueue queue = neighborQueues[i];
+            if (queue.size() == 0) {
+                // no neighbors, skip
+                neighborhoods[i] = NeighborHood.EMPTY;
+                continue;
+            }
+            // consume the queue into the neighbors array and get the maximum intra-cluster distance
+            int[] neighbors = new int[queue.size()];
+            float maxIntraDistance = queue.topScore();
+            int iter = 0;
+            while (queue.size() > 0) {
+                neighbors[neighbors.length - ++iter] = queue.pop();
+            }
+            neighborhoods[i] = new NeighborHood(neighbors, maxIntraDistance);
+        }
+        return neighborhoods;
+    }
+
     public static NeighborHood[] computeNeighborhoodsGraph(float[][] centers, int clustersPerNeighborhood) throws IOException {
         final UpdateableRandomVectorScorer scorer = new UpdateableRandomVectorScorer() {
             int scoringOrdinal;
@@ -57,92 +115,97 @@ public RandomVectorScorerSupplier copy() {
                 return this;
             }
         };
-        final OnHeapHnswGraph graph = HnswGraphBuilder.create(supplier, 16, 100, 42L).build(centers.length);
+        final OnHeapHnswGraph graph = HnswGraphBuilder.create(supplier, M, EF_CONSTRUCTION, 42L).build(centers.length);
         final NeighborHood[] neighborhoods = new NeighborHood[centers.length];
-        final SingleBit singleBit = new SingleBit(centers.length);
+        // oversample the number of neighbors we collect to improve recall
+        final ReusableKnnCollector collector = new ReusableKnnCollector(2 * clustersPerNeighborhood);
         for (int i = 0; i < centers.length; i++) {
+            collector.reset(i);
             scorer.setScoringOrdinal(i);
-            singleBit.indexSet = i;
-            final KnnCollector collector = HnswGraphSearcher.search(scorer, clustersPerNeighborhood, graph, singleBit, Integer.MAX_VALUE);
-            final ScoreDoc[] scoreDocs = collector.topDocs().scoreDocs;
-            if (scoreDocs.length == 0) {
+            HnswGraphSearcher.search(scorer, collector, graph, null);
+            NeighborQueue queue = collector.queue;
+            if (queue.size() == 0) {
                 // no neighbors, skip
                 neighborhoods[i] = NeighborHood.EMPTY;
                 continue;
             }
-            final int[] neighbors = new int[scoreDocs.length];
-            for (int j = 0; j < neighbors.length; j++) {
-                neighbors[j] = scoreDocs[j].doc;
-                assert neighbors[j] != i;
+            while (queue.size() > clustersPerNeighborhood) {
+                queue.pop();
+            }
+            final float minScore = queue.topScore();
+            final int[] neighbors = new int[queue.size()];
+            for (int j = 1; j <= neighbors.length; j++) {
+                neighbors[neighbors.length - j] = queue.pop();
             }
-            final float minCompetitiveSimilarity = (1f / scoreDocs[neighbors.length - 1].score) - 1;
-            neighborhoods[i] = new NeighborHood(neighbors, minCompetitiveSimilarity);
+            neighborhoods[i] = new NeighborHood(neighbors, (1f / minScore) - 1);
         }
         return neighborhoods;
     }
 
-    private static class SingleBit implements Bits {
+    private static class ReusableKnnCollector implements KnnCollector {
 
-        private final int length;
-        private int indexSet;
+        private final NeighborQueue queue;
+        private final int k;
+        int visitedCount;
+        int currenOrd;
 
-        SingleBit(int length) {
-            this.length = length;
+        ReusableKnnCollector(int k) {
+            this.k = k;
+            this.queue = new NeighborQueue(k, false);
+        }
+
+        void reset(int ord) {
+            queue.clear();
+            visitedCount = 0;
+            currenOrd = ord;
         }
 
         @Override
-        public boolean get(int index) {
-            return index != indexSet;
+        public boolean earlyTerminated() {
+            return false;
         }
 
         @Override
-        public int length() {
-            return length;
+        public void incVisitedCount(int count) {
+            visitedCount += count;
         }
-    }
 
-    public static NeighborHood[] computeNeighborhoodsBruteForce(float[][] centers, int clustersPerNeighborhood) {
-        int k = centers.length;
-        NeighborQueue[] neighborQueues = new NeighborQueue[k];
-        for (int i = 0; i < k; i++) {
-            neighborQueues[i] = new NeighborQueue(clustersPerNeighborhood, true);
+        @Override
+        public long visitedCount() {
+            return visitedCount;
         }
-        final float[] scores = new float[4];
-        final int limit = k - 3;
-        for (int i = 0; i < k - 1; i++) {
-            float[] center = centers[i];
-            int j = i + 1;
-            for (; j < limit; j += 4) {
-                ESVectorUtil.squareDistanceBulk(center, centers[j], centers[j + 1], centers[j + 2], centers[j + 3], scores);
-                for (int h = 0; h < 4; h++) {
-                    neighborQueues[j + h].insertWithOverflow(i, scores[h]);
-                    neighborQueues[i].insertWithOverflow(j + h, scores[h]);
-                }
-            }
-            for (; j < k; j++) {
-                float dsq = VectorUtil.squareDistance(center, centers[j]);
-                neighborQueues[j].insertWithOverflow(i, dsq);
-                neighborQueues[i].insertWithOverflow(j, dsq);
-            }
+
+        @Override
+        public long visitLimit() {
+            return Integer.MAX_VALUE;
         }
 
-        NeighborHood[] neighborhoods = new NeighborHood[k];
-        for (int i = 0; i < k; i++) {
-            NeighborQueue queue = neighborQueues[i];
-            if (queue.size() == 0) {
-                // no neighbors, skip
-                neighborhoods[i] = NeighborHood.EMPTY;
-                continue;
-            }
-            // consume the queue into the neighbors array and get the maximum intra-cluster distance
-            int[] neighbors = new int[queue.size()];
-            float maxIntraDistance = queue.topScore();
-            int iter = 0;
-            while (queue.size() > 0) {
-                neighbors[neighbors.length - ++iter] = queue.pop();
+        @Override
+        public int k() {
+            return k;
+        }
+
+        @Override
+        public boolean collect(int docId, float similarity) {
+            if (currenOrd != docId) {
+                return queue.insertWithOverflow(docId, similarity);
             }
-            neighborhoods[i] = new NeighborHood(neighbors, maxIntraDistance);
+            return false;
+        }
+
+        @Override
+        public float minCompetitiveSimilarity() {
+            return queue.size() >= k() ? queue.topScore() : Float.NEGATIVE_INFINITY;
+        }
+
+        @Override
+        public TopDocs topDocs() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public KnnSearchStrategy getSearchStrategy() {
+            return null;
         }
-        return neighborhoods;
     }
 }

server/src/test/java/org/elasticsearch/index/codec/vectors/cluster/KMeansLocalTests.java

@@ -146,23 +146,23 @@ private static FloatVectorValues generateData(int nSamples, int nDims, int nClus
     }
 
     public void testComputeNeighbours() throws IOException {
-        int numCentroids = randomIntBetween(100, 1000);
+        int numCentroids = randomIntBetween(1000, 2000);
         int dims = randomIntBetween(10, 200);
         float[][] vectors = new float[numCentroids][dims];
         for (int i = 0; i < numCentroids; i++) {
             for (int j = 0; j < dims; j++) {
                 vectors[i][j] = randomFloat();
             }
         }
-        int clustersPerNeighbour = randomIntBetween(32, 64);
+        int clustersPerNeighbour = randomIntBetween(32, 128);
         NeighborHood[] neighborHoodsGraph = NeighborHood.computeNeighborhoodsGraph(vectors, clustersPerNeighbour);
         NeighborHood[] neighborHoodsBruteForce = NeighborHood.computeNeighborhoodsBruteForce(vectors, clustersPerNeighbour);
         assertEquals(neighborHoodsGraph.length, neighborHoodsBruteForce.length);
         for (int i = 0; i < neighborHoodsGraph.length; i++) {
             assertEquals(neighborHoodsBruteForce[i].neighbors().length, neighborHoodsGraph[i].neighbors().length);
             int matched = compareNN(i, neighborHoodsBruteForce[i].neighbors(), neighborHoodsGraph[i].neighbors());
             double recall = (double) matched / neighborHoodsGraph[i].neighbors().length;
-            assertThat(recall, greaterThanOrEqualTo(0.7));
+            assertThat(recall, greaterThanOrEqualTo(0.5));
             if (recall == 1.0) {
                 // we cannot assert on array equality as there can be small differences due to numerical errors
                 assertEquals(neighborHoodsBruteForce[i].maxIntraDistance(), neighborHoodsGraph[i].maxIntraDistance(), 1e-5f);