KMeansIntermediate shares assigments

Author: iverase

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

@@ -10,7 +10,6 @@
 package org.elasticsearch.index.codec.vectors.cluster;
 
 import org.apache.lucene.index.FloatVectorValues;
-import org.apache.lucene.util.VectorUtil;
 
 import java.io.IOException;
 
@@ -85,42 +84,16 @@ KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int ta
 
         // TODO: instead of creating a sub-cluster assignments reuse the parent array each time
         int[] assignments = new int[vectors.size()];
-
         KMeansLocal kmeans = new KMeansLocal(m, maxIterations);
         float[][] centroids = KMeansLocal.pickInitialCentroids(vectors, k);
-        KMeansIntermediate kMeansIntermediate = new KMeansIntermediate(centroids);
+        KMeansIntermediate kMeansIntermediate = new KMeansIntermediate(centroids, assignments, vectors::ordToDoc);
         kmeans.cluster(vectors, kMeansIntermediate);
 
         // 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];
-        float[][] nextCentroids = new float[centroids.length][dimension];
-        for (int i = 0; i < vectors.size(); i++) {
-            float smallest = Float.MAX_VALUE;
-            int centroidIdx = -1;
-            float[] vector = vectors.vectorValue(i);
-            for (int j = 0; j < centroids.length; j++) {
-                float[] centroid = centroids[j];
-                float d = VectorUtil.squareDistance(vector, centroid);
-                if (d < smallest) {
-                    smallest = d;
-                    centroidIdx = j;
-                }
-            }
-            centroidVectorCount[centroidIdx]++;
-            for (int j = 0; j < dimension; j++) {
-                nextCentroids[centroidIdx][j] += vector[j];
-            }
-            assignments[i] = centroidIdx;
-        }
-
-        // update centroids based on assignments of all vectors
-        for (int i = 0; i < centroids.length; i++) {
-            if (centroidVectorCount[i] > 0) {
-                for (int j = 0; j < dimension; j++) {
-                    centroids[i][j] = nextCentroids[i][j] / centroidVectorCount[i];
-                }
-            }
+        for (int assigment : assignments) {
+            centroidVectorCount[assigment]++;
         }
 
         int effectiveK = 0;
@@ -130,8 +103,6 @@ KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int ta
             }
         }
 
-        kMeansIntermediate = new KMeansIntermediate(centroids, assignments, vectors::ordToDoc);
-
         if (effectiveK == 1) {
             return kMeansIntermediate;
         }

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

@@ -31,10 +31,6 @@ private KMeansIntermediate(float[][] centroids, int[] assignments, IntToIntFunct
         this(new float[0][0], new int[0], i -> i, new int[0]);
     }
 
-    KMeansIntermediate(float[][] centroids) {
-        this(centroids, new int[0], i -> i, new int[0]);
-    }
-
     KMeansIntermediate(float[][] centroids, int[] assignments) {
         this(centroids, assignments, i -> i, new int[0]);
     }

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

@@ -276,7 +276,8 @@ void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate, L
             return;
         }
 
-        int[] assignments = new int[n];
+        int[] assignments = kMeansIntermediate.assignments();
+        assert assignments.length == n;
         float[][] nextCentroids = new float[centroids.length][vectors.dimension()];
         for (int i = 0; i < maxIterations; i++) {
             if (stepLloyd(vectors, centroids, nextCentroids, assignments, sampleSize, neighborhoods) == false) {
@@ -296,7 +297,7 @@ void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate, L
      * @param maxIterations the max iterations to shift centroids
      */
     public static void cluster(FloatVectorValues vectors, float[][] centroids, int sampleSize, int maxIterations) throws IOException {
-        KMeansIntermediate kMeansIntermediate = new KMeansIntermediate(centroids);
+        KMeansIntermediate kMeansIntermediate = new KMeansIntermediate(centroids, new int[vectors.size()], vectors::ordToDoc);
         KMeansLocal kMeans = new KMeansLocal(sampleSize, maxIterations);
         kMeans.cluster(vectors, kMeansIntermediate);
     }