Remove empty clusters in Hierarchical k-means (#132569)

This commit makes sure the result of the hierarchical k-means cannot produce empty clusters.

Author: iverase

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

@@ -106,29 +106,57 @@ KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int ta
         // TODO: consider adding cluster size counts to the kmeans algo
         // handle assignment here so we can track distance and cluster size
         int[] centroidVectorCount = new int[centroids.length];
+        int effectiveCluster = -1;
         int effectiveK = 0;
         for (int assigment : assignments) {
             centroidVectorCount[assigment]++;
             // this cluster has received an assignment, its now effective, but only count it once
             if (centroidVectorCount[assigment] == 1) {
                 effectiveK++;
+                effectiveCluster = assigment;
             }
         }
 
         if (effectiveK == 1) {
+            final float[][] singleClusterCentroid = new float[1][];
+            singleClusterCentroid[0] = centroids[effectiveCluster];
+            kMeansIntermediate.setCentroids(singleClusterCentroid);
+            Arrays.fill(kMeansIntermediate.assignments(), 0);
             return kMeansIntermediate;
         }
 
+        int removedElements = 0;
         for (int c = 0; c < centroidVectorCount.length; c++) {
             // Recurse for each cluster which is larger than targetSize
             // Give ourselves 30% margin for the target size
-            if (100 * centroidVectorCount[c] > 134 * targetSize) {
-                FloatVectorValues sample = createClusterSlice(centroidVectorCount[c], c, vectors, assignments);
-
+            final int count = centroidVectorCount[c];
+            final int adjustedCentroid = c - removedElements;
+            if (100 * count > 134 * targetSize) {
+                final FloatVectorValues sample = createClusterSlice(count, adjustedCentroid, vectors, assignments);
                 // TODO: consider iterative here instead of recursive
                 // recursive call to build out the sub partitions around this centroid c
                 // subsequently reconcile and flatten the space of all centroids and assignments into one structure we can return
-                updateAssignmentsWithRecursiveSplit(kMeansIntermediate, c, clusterAndSplit(sample, targetSize));
+                updateAssignmentsWithRecursiveSplit(kMeansIntermediate, adjustedCentroid, clusterAndSplit(sample, targetSize));
+            } else if (count == 0) {
+                // remove empty clusters
+                final int newSize = kMeansIntermediate.centroids().length - 1;
+                final float[][] newCentroids = new float[newSize][];
+                System.arraycopy(kMeansIntermediate.centroids(), 0, newCentroids, 0, adjustedCentroid);
+                System.arraycopy(
+                    kMeansIntermediate.centroids(),
+                    adjustedCentroid + 1,
+                    newCentroids,
+                    adjustedCentroid,
+                    newSize - adjustedCentroid
+                );
+                // we need to update the assignments to reflect the new centroid ordinals
+                for (int i = 0; i < kMeansIntermediate.assignments().length; i++) {
+                    if (kMeansIntermediate.assignments()[i] > adjustedCentroid) {
+                        kMeansIntermediate.assignments()[i]--;
+                    }
+                }
+                kMeansIntermediate.setCentroids(newCentroids);
+                removedElements++;
             }
         }
 

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

@@ -74,4 +74,63 @@ private static FloatVectorValues generateData(int nSamples, int nDims, int nClus
         }
         return FloatVectorValues.fromFloats(vectors, nDims);
     }
+
+    public void testFewDifferentValues() throws IOException {
+        int nVectors = random().nextInt(100, 1000);
+        int targetSize = random().nextInt(4, 64);
+        int dims = random().nextInt(2, 20);
+        int diffValues = randomIntBetween(1, 5);
+        float[][] values = new float[diffValues][dims];
+        for (int i = 0; i < diffValues; i++) {
+            for (int j = 0; j < dims; j++) {
+                values[i][j] = random().nextFloat();
+            }
+        }
+        List<float[]> vectorList = new ArrayList<>(nVectors);
+        for (int i = 0; i < nVectors; i++) {
+            vectorList.add(values[random().nextInt(diffValues)]);
+        }
+        FloatVectorValues vectors = FloatVectorValues.fromFloats(vectorList, dims);
+
+        HierarchicalKMeans hkmeans = new HierarchicalKMeans(
+            dims,
+            random().nextInt(1, 100),
+            random().nextInt(Math.min(nVectors, 100), nVectors + 1),
+            random().nextInt(2, 512),
+            random().nextFloat(0.5f, 1.5f)
+        );
+
+        KMeansResult result = hkmeans.cluster(vectors, targetSize);
+
+        float[][] centroids = result.centroids();
+        int[] assignments = result.assignments();
+        int[] soarAssignments = result.soarAssignments();
+
+        int[] counts = new int[centroids.length];
+        for (int i = 0; i < assignments.length; i++) {
+            counts[assignments[i]]++;
+        }
+        int totalCount = 0;
+        for (int count : counts) {
+            totalCount += count;
+            assertTrue(count > 0);
+        }
+        assertEquals(nVectors, totalCount);
+
+        assertEquals(nVectors, assignments.length);
+
+        for (int assignment : assignments) {
+            assertTrue(assignment >= 0 && assignment < centroids.length);
+        }
+        if (centroids.length > 1 && centroids.length < nVectors) {
+            assertEquals(nVectors, soarAssignments.length);
+            // verify no duplicates exist
+            for (int i = 0; i < assignments.length; i++) {
+                assertTrue(soarAssignments[i] >= 0 && soarAssignments[i] < centroids.length);
+                assertNotEquals(assignments[i], soarAssignments[i]);
+            }
+        } else {
+            assertEquals(0, soarAssignments.length);
+        }
+    }
 }