minor test fixes and edge cases

Author: john-wagster

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

@@ -58,7 +58,7 @@ public KMeansResult cluster(FloatVectorValues vectors, int targetSize) throws IO
         }
 
         // if we have a small number of vectors pick one and output that as the centroid
-        if (vectors.size() < targetSize) {
+        if (vectors.size() <= targetSize) {
             float[] centroid = new float[dimension];
             System.arraycopy(vectors.vectorValue(0), 0, centroid, 0, dimension);
             return new KMeansIntermediate(new float[][] { centroid }, new int[vectors.size()]);

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

@@ -196,6 +196,7 @@ private int[] assignSpilled(FloatVectorValues vectors, List<int[]> neighborhoods
 
             int bestAssignment = -1;
             float minSoar = Float.MAX_VALUE;
+            assert neighborhoods.get(currAssignment) != null;
             for (int neighbor : neighborhoods.get(currAssignment)) {
                 if (neighbor == currAssignment) {
                     continue;
@@ -260,7 +261,7 @@ void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate, b
             computeNeighborhoods(centroids, neighborhoods, clustersPerNeighborhood);
         }
         cluster(vectors, kMeansIntermediate, neighborhoods);
-        if (neighborAware) {
+        if (neighborAware && clustersPerNeighborhood > 0) {
             int[] assignments = kMeansIntermediate.assignments();
             assert assignments != null;
             assert assignments.length == vectors.size();

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

@@ -21,7 +21,7 @@ public void testHKmeans() throws IOException {
         int nClusters = random().nextInt(1, 10);
         int nVectors = random().nextInt(nClusters * 100, nClusters * 200);
         int dims = random().nextInt(2, 20);
-        int sampleSize = random().nextInt(100, nVectors);
+        int sampleSize = random().nextInt(100, nVectors + 1);
         int maxIterations = random().nextInt(0, 100);
         int clustersPerNeighborhood = random().nextInt(0, 512);
         float soarLambda = random().nextFloat(0.5f, 1.5f);
@@ -36,13 +36,14 @@ public void testHKmeans() throws IOException {
         int[] assignments = result.assignments();
         int[] soarAssignments = result.soarAssignments();
 
-        assertEquals(nClusters, centroids.length, 5);
+        assertEquals(nClusters, centroids.length, 6);
         assertEquals(nVectors, assignments.length);
-        assertEquals(nVectors, soarAssignments.length);
-
-        // verify no duplicates exist
-        for (int i = 0; i < assignments.length; i++) {
-            assert assignments[i] != soarAssignments[i];
+        if (centroids.length > 1 && clustersPerNeighborhood > 0) {
+            assertEquals(nVectors, soarAssignments.length);
+            // verify no duplicates exist
+            for (int i = 0; i < assignments.length; i++) {
+                assert assignments[i] != soarAssignments[i];
+            }
         }
     }
 

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

@@ -23,7 +23,7 @@ public void testKMeansNeighbors() throws IOException {
         int nClusters = random().nextInt(1, 10);
         int nVectors = random().nextInt(nClusters * 100, nClusters * 200);
         int dims = random().nextInt(2, 20);
-        int sampleSize = random().nextInt(100, nVectors);
+        int sampleSize = random().nextInt(100, nVectors + 1);
         int maxIterations = random().nextInt(0, 100);
         int clustersPerNeighborhood = random().nextInt(0, 512);
         float soarLambda = random().nextFloat(0.5f, 1.5f);