addressing PR comments

Author: benwtrent

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

@@ -179,25 +179,19 @@ private void computeNeighborhoods(float[][] centers, List<NeighborHood> neighbor
             }
         }
 
-        float[] scores = new float[clustersPerNeighborhood];
         for (int i = 0; i < k; i++) {
             NeighborQueue queue = neighborQueues.get(i);
-            int neighborCount = queue.size();
-            int[] neighbors = new int[neighborCount];
-            float maxIntraDistance = queue.consumeNodesWithWorstScore(neighbors, scores);
-            // Sort neighbors by their score
-            for (int j = 0; j < neighborCount; j++) {
-                for (int l = j + 1; l < neighborCount; l++) {
-                    if (scores[j] > scores[l]) {
-                        // swap
-                        int tmp = neighbors[j];
-                        neighbors[j] = neighbors[l];
-                        neighbors[l] = tmp;
-                        float tmpScore = scores[j];
-                        scores[j] = scores[l];
-                        scores[l] = tmpScore;
-                    }
-                }
+            if (queue.size() == 0) {
+                // no neighbors, skip
+                neighborhoods.set(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();
             }
             NeighborHood neighborHood = new NeighborHood(neighbors, maxIntraDistance);
             neighborhoods.set(i, neighborHood);
@@ -283,7 +277,9 @@ void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate) t
         cluster(vectors, kMeansIntermediate, false);
     }
 
-    record NeighborHood(int[] neighbors, float maxIntraDistance) {}
+    record NeighborHood(int[] neighbors, float maxIntraDistance) {
+        static final NeighborHood EMPTY = new NeighborHood(new int[0], Float.POSITIVE_INFINITY);
+    }
 
     /**
      * cluster using a lloyd kmeans algorithm that also considers prior clustered neighborhoods when adjusting centroids

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

@@ -121,23 +121,6 @@ public int pop() {
         return decodeNodeId(heap.pop());
     }
 
-    public float consumeNodesWithWorstScore(int[] dest, float[] scores) {
-        if (dest.length < size()) {
-            throw new IllegalArgumentException("Destination array is too small. Expected at least " + size() + " elements.");
-        }
-        float worstScore = Float.NEGATIVE_INFINITY;
-        for (int i = 0; i < size(); i++) {
-            long heapValue = heap.get(i + 1);
-            float score = decodeScore(heapValue);
-            dest[i] = decodeNodeId(heapValue);
-            scores[i] = score;
-            if (score > worstScore) {
-                worstScore = score;
-            }
-        }
-        return worstScore;
-    }
-
     public void clear() {
         heap.clear();
     }