Handle more edge cases in HierarchicalKMeans (#130622)

This commit handle the following cases:

1) When the number of centroids is one (and better checks when number of
centroids < number of vectors).

2) when number of maxIterations is equal to 0 (no sampling)

fixes #130497

Author: iverase

muted-tests.yml

@@ -578,9 +578,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeIT
   method: test
   issue: https://github.com/elastic/elasticsearch/issues/130067
-- class: org.elasticsearch.index.codec.vectors.cluster.HierarchicalKMeansTests
-  method: testHKmeans
-  issue: https://github.com/elastic/elasticsearch/issues/130497
 - class: org.elasticsearch.backwards.MixedClusterClientYamlTestSuiteIT
   method: test {p0=search.vectors/40_knn_search/Dimensions are dynamically set}
   issue: https://github.com/elastic/elasticsearch/issues/130626

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

@@ -329,8 +329,10 @@ private void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansInterme
         float[][] centroids = kMeansIntermediate.centroids();
         int k = centroids.length;
         int n = vectors.size();
+        int[] assignments = kMeansIntermediate.assignments();
 
-        if (k == 1 || k >= n) {
+        if (k == 1) {
+            Arrays.fill(assignments, 0);
             return;
         }
         IntToIntFunction translateOrd = i -> i;
@@ -339,7 +341,7 @@ private void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansInterme
             sampledVectors = SampleReader.createSampleReader(vectors, sampleSize, 42L);
             translateOrd = sampledVectors::ordToDoc;
         }
-        int[] assignments = kMeansIntermediate.assignments();
+
         assert assignments.length == n;
         float[][] nextCentroids = new float[centroids.length][vectors.dimension()];
         for (int i = 0; i < maxIterations; i++) {
@@ -349,7 +351,7 @@ private void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansInterme
             }
         }
         // If we were sampled, do a once over the full set of vectors to finalize the centroids
-        if (sampleSize < n) {
+        if (sampleSize < n || maxIterations == 0) {
             // No ordinal translation needed here, we are using the full set of vectors
             stepLloyd(vectors, i -> i, centroids, nextCentroids, assignments, neighborhoods);
         }

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

@@ -19,10 +19,10 @@ public class HierarchicalKMeansTests extends ESTestCase {
 
     public void testHKmeans() throws IOException {
         int nClusters = random().nextInt(1, 10);
-        int nVectors = random().nextInt(nClusters * 100, nClusters * 200);
+        int nVectors = random().nextInt(1, nClusters * 200);
         int dims = random().nextInt(2, 20);
-        int sampleSize = random().nextInt(100, nVectors + 1);
-        int maxIterations = random().nextInt(0, 100);
+        int sampleSize = random().nextInt(Math.min(nVectors, 100), nVectors + 1);
+        int maxIterations = random().nextInt(1, 100);
         int clustersPerNeighborhood = random().nextInt(2, 512);
         float soarLambda = random().nextFloat(0.5f, 1.5f);
         FloatVectorValues vectors = generateData(nVectors, dims, nClusters);
@@ -36,14 +36,16 @@ public void testHKmeans() throws IOException {
         int[] assignments = result.assignments();
         int[] soarAssignments = result.soarAssignments();
 
-        assertEquals(nClusters, centroids.length, 6);
+        assertEquals(Math.min(nClusters, nVectors), centroids.length, 8);
         assertEquals(nVectors, assignments.length);
-        if (centroids.length > 1 && clustersPerNeighborhood > 0) {
+        if (centroids.length > 1 && centroids.length < nVectors) {
             assertEquals(nVectors, soarAssignments.length);
             // verify no duplicates exist
             for (int i = 0; i < assignments.length; i++) {
                 assert assignments[i] != soarAssignments[i];
             }
+        } else {
+            assertEquals(0, soarAssignments.length);
         }
     }