[DiskBBQ] Add concurrency to KMeansLocal (#139239)

Hierarchical k-means algorithm can be run using several threads providing a significant speed up.

Author: iverase

docs/changelog/139239.yaml

@@ -0,0 +1,5 @@
+pr: 139239
+summary: "[DiskBBQ] Add concurrency on KMeansLocal"
+area: Vector Search
+type: enhancement
+issues: []

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

@@ -10,6 +10,7 @@
 package org.elasticsearch.index.codec.vectors.cluster;
 
 import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.search.TaskExecutor;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -25,27 +26,84 @@ public class HierarchicalKMeans {
     public static final int SAMPLES_PER_CLUSTER_DEFAULT = 64;
     public static final float DEFAULT_SOAR_LAMBDA = 1.0f;
     public static final int NO_SOAR_ASSIGNMENT = -1;
+    private static final int MIN_VECTORS_PRE_THREAD = 64;
 
     final int dimension;
     final int maxIterations;
     final int samplesPerCluster;
     final int clustersPerNeighborhood;
     final float soarLambda;
 
-    public HierarchicalKMeans(int dimension) {
-        this(dimension, MAX_ITERATIONS_DEFAULT, SAMPLES_PER_CLUSTER_DEFAULT, MAXK, DEFAULT_SOAR_LAMBDA);
-    }
+    private final TaskExecutor executor;
+    private final int numWorkers;
 
-    public HierarchicalKMeans(int dimension, int maxIterations, int samplesPerCluster, int clustersPerNeighborhood, float soarLambda) {
+    private HierarchicalKMeans(
+        int dimension,
+        TaskExecutor executor,
+        int numWorkers,
+        int maxIterations,
+        int samplesPerCluster,
+        int clustersPerNeighborhood,
+        float soarLambda
+    ) {
         this.dimension = dimension;
+        this.executor = executor;
+        this.numWorkers = numWorkers;
         this.maxIterations = maxIterations;
         this.samplesPerCluster = samplesPerCluster;
         this.clustersPerNeighborhood = clustersPerNeighborhood;
         this.soarLambda = soarLambda;
     }
 
+    public static HierarchicalKMeans ofSerial(int dimension) {
+        return ofSerial(dimension, MAX_ITERATIONS_DEFAULT, SAMPLES_PER_CLUSTER_DEFAULT, MAXK, DEFAULT_SOAR_LAMBDA);
+    }
+
+    public static HierarchicalKMeans ofSerial(
+        int dimension,
+        int maxIterations,
+        int samplesPerCluster,
+        int clustersPerNeighborhood,
+        float soarLambda
+    ) {
+        return new HierarchicalKMeans(dimension, null, 1, maxIterations, samplesPerCluster, clustersPerNeighborhood, soarLambda);
+    }
+
+    public static HierarchicalKMeans ofConcurrent(int dimension, TaskExecutor executor, int numWorkers) {
+        return ofConcurrent(
+            dimension,
+            executor,
+            numWorkers,
+            MAX_ITERATIONS_DEFAULT,
+            SAMPLES_PER_CLUSTER_DEFAULT,
+            MAXK,
+            DEFAULT_SOAR_LAMBDA
+        );
+    }
+
+    public static HierarchicalKMeans ofConcurrent(
+        int dimension,
+        TaskExecutor executor,
+        int numWorkers,
+        int maxIterations,
+        int samplesPerCluster,
+        int clustersPerNeighborhood,
+        float soarLambda
+    ) {
+        return new HierarchicalKMeans(
+            dimension,
+            executor,
+            numWorkers,
+            maxIterations,
+            samplesPerCluster,
+            clustersPerNeighborhood,
+            soarLambda
+        );
+
+    }
+
     /**
-     * clusters or moreso partitions the set of vectors by starting with a rough number of partitions and then recursively refining those
+     * clusters the set of vectors by starting with a rough number of partitions and then recursively refining those
      * lastly a pass is made to adjust nearby neighborhoods and add an extra assignment per vector to nearby neighborhoods
      *
      * @param vectors the vectors to cluster
@@ -54,7 +112,6 @@ public HierarchicalKMeans(int dimension, int maxIterations, int samplesPerCluste
      * @throws IOException is thrown if vectors is inaccessible
      */
     public KMeansResult cluster(FloatVectorValues vectors, int targetSize) throws IOException {
-
         if (vectors.size() == 0) {
             return new KMeansIntermediate();
         }
@@ -80,14 +137,13 @@ public KMeansResult cluster(FloatVectorValues vectors, int targetSize) throws IO
         KMeansIntermediate kMeansIntermediate = clusterAndSplit(vectors, targetSize);
         if (kMeansIntermediate.centroids().length > 1 && kMeansIntermediate.centroids().length < vectors.size()) {
             int localSampleSize = Math.min(kMeansIntermediate.centroids().length * samplesPerCluster / 2, vectors.size());
-            KMeansLocal kMeansLocal = new KMeansLocal(localSampleSize, maxIterations);
+            KMeansLocal kMeansLocal = buildKmeansLocal(vectors.size(), localSampleSize);
             kMeansLocal.cluster(vectors, kMeansIntermediate, clustersPerNeighborhood, soarLambda);
         }
-
         return kMeansIntermediate;
     }
 
-    KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int targetSize) throws IOException {
+    private KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int targetSize) throws IOException {
         if (vectors.size() <= targetSize) {
             return new KMeansIntermediate();
         }
@@ -99,10 +155,10 @@ KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int ta
         int[] assignments = new int[vectors.size()];
         // ensure we don't over assign to cluster 0 without adjusting it
         Arrays.fill(assignments, -1);
-        KMeansLocal kmeans = new KMeansLocal(m, maxIterations);
         float[][] centroids = KMeansLocal.pickInitialCentroids(vectors, k);
         KMeansIntermediate kMeansIntermediate = new KMeansIntermediate(centroids, assignments, vectors::ordToDoc);
-        kmeans.cluster(vectors, kMeansIntermediate);
+        KMeansLocal kMeansLocal = buildKmeansLocal(vectors.size(), m);
+        kMeansLocal.cluster(vectors, kMeansIntermediate);
 
         // TODO: consider adding cluster size counts to the kmeans algo
         // handle assignment here so we can track distance and cluster size
@@ -164,6 +220,14 @@ KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int ta
         return kMeansIntermediate;
     }
 
+    private KMeansLocal buildKmeansLocal(int numVectors, int localSampleSize) {
+        int numWorkers = Math.min(this.numWorkers, numVectors / MIN_VECTORS_PRE_THREAD);
+        // if there is no executor or there is no enough vectors for more than one thread, use the serial version
+        return executor == null || numWorkers <= 1
+            ? new KMeansLocalSerial(localSampleSize, maxIterations)
+            : new KMeansLocalConcurrent(executor, numWorkers, localSampleSize, maxIterations);
+    }
+
     static FloatVectorValues createClusterSlice(int clusterSize, int cluster, FloatVectorValues vectors, int[] assignments) {
         int[] slice = new int[clusterSize];
         int idx = 0;

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

@@ -26,22 +26,43 @@
  * with finalizing nearby pre-established clusters and generate
  * <a href="https://research.google/blog/soar-new-algorithms-for-even-faster-vector-search-with-scann/">SOAR</a> assignments
  */
-class KMeansLocal {
+abstract class KMeansLocal {
 
     // the minimum distance that is considered to be "far enough" to a centroid in order to compute the soar distance.
     // For vectors that are closer than this distance to the centroid don't get spilled because they are well represented
     // by the centroid itself. In many cases, it indicates a degenerated distribution, e.g the cluster is composed of the
     // many equal vectors.
     private static final float SOAR_MIN_DISTANCE = 1e-16f;
 
-    final int sampleSize;
-    final int maxIterations;
+    private final int sampleSize;
+    private final int maxIterations;
 
     KMeansLocal(int sampleSize, int maxIterations) {
         this.sampleSize = sampleSize;
         this.maxIterations = maxIterations;
     }
 
+    /** Number of workers to use for parallelism **/
+    protected abstract int numWorkers();
+
+    /** assign to each vector the closest centroid **/
+    protected abstract boolean stepLloyd(
+        FloatVectorValues vectors,
+        IntToIntFunction translateOrd,
+        float[][] centroids,
+        FixedBitSet[] centroidChangedSlices,
+        int[] assignments,
+        NeighborHood[] neighborHoods
+    ) throws IOException;
+
+    /** assign to each vector the soar assignment **/
+    protected abstract void assignSpilled(
+        FloatVectorValues vectors,
+        KMeansIntermediate kmeansIntermediate,
+        NeighborHood[] neighborhoods,
+        float soarLambda
+    ) throws IOException;
+
     /**
      * uses a Reservoir Sampling approach to picking the initial centroids which are subsequently expected
      * to be used by a clustering algorithm
@@ -69,20 +90,21 @@ static float[][] pickInitialCentroids(FloatVectorValues vectors, int centroidCou
         return centroids;
     }
 
-    private static boolean stepLloyd(
+    /** Assign vectors from {@code start} to {@code end} to the closest centroid. */
+    protected static boolean stepLloydSlice(
         FloatVectorValues vectors,
         IntToIntFunction translateOrd,
         float[][] centroids,
         FixedBitSet centroidChanged,
-        int[] centroidCounts,
         int[] assignments,
-        NeighborHood[] neighborhoods
+        NeighborHood[] neighborhoods,
+        int start,
+        int end
     ) throws IOException {
         boolean changed = false;
-        int dim = vectors.dimension();
         centroidChanged.clear();
         final float[] distances = new float[4];
-        for (int idx = 0; idx < vectors.size(); idx++) {
+        for (int idx = start; idx < end; idx++) {
             float[] vector = vectors.vectorValue(idx);
             int vectorOrd = translateOrd.apply(idx);
             final int assignment = assignments[vectorOrd];
@@ -101,36 +123,49 @@ private static boolean stepLloyd(
                 changed = true;
             }
         }
-        if (changed) {
-            Arrays.fill(centroidCounts, 0);
-            for (int idx = 0; idx < vectors.size(); idx++) {
-                final int assignment = assignments[translateOrd.apply(idx)];
-                if (centroidChanged.get(assignment)) {
-                    float[] centroid = centroids[assignment];
-                    if (centroidCounts[assignment]++ == 0) {
-                        Arrays.fill(centroid, 0.0f);
-                    }
-                    float[] vector = vectors.vectorValue(idx);
+        return changed;
+    }
+
+    private static void updateCentroids(
+        FloatVectorValues vectors,
+        IntToIntFunction translateOrd,
+        float[][] centroids,
+        FixedBitSet[] centroidChangedSlices,
+        int[] centroidCounts,
+        int[] assignments
+    ) throws IOException {
+        Arrays.fill(centroidCounts, 0);
+        FixedBitSet centroidChanged = centroidChangedSlices[0];
+        for (int j = 1; j < centroidChangedSlices.length; j++) {
+            centroidChanged.or(centroidChangedSlices[j]);
+        }
+        int dim = vectors.dimension();
+        for (int idx = 0; idx < vectors.size(); idx++) {
+            final int assignment = assignments[translateOrd.apply(idx)];
+            if (centroidChanged.get(assignment)) {
+                float[] centroid = centroids[assignment];
+                float[] vector = vectors.vectorValue(idx);
+                if (centroidCounts[assignment]++ == 0) {
+                    System.arraycopy(vector, 0, centroid, 0, dim);
+                } else {
                     for (int d = 0; d < dim; d++) {
                         centroid[d] += vector[d];
                     }
                 }
             }
+        }
 
-            for (int clusterIdx = 0; clusterIdx < centroids.length; clusterIdx++) {
-                if (centroidChanged.get(clusterIdx)) {
-                    float count = (float) centroidCounts[clusterIdx];
-                    if (count > 0) {
-                        float[] centroid = centroids[clusterIdx];
-                        for (int d = 0; d < dim; d++) {
-                            centroid[d] /= count;
-                        }
+        for (int clusterIdx = 0; clusterIdx < centroids.length; clusterIdx++) {
+            if (centroidChanged.get(clusterIdx)) {
+                float count = (float) centroidCounts[clusterIdx];
+                if (count > 0) {
+                    float[] centroid = centroids[clusterIdx];
+                    for (int d = 0; d < dim; d++) {
+                        centroid[d] /= count;
                     }
                 }
             }
         }
-
-        return changed;
     }
 
     private static int getBestCentroidFromNeighbours(
@@ -211,11 +246,14 @@ private static int getBestCentroid(float[][] centroids, float[] vector, float[]
         return bestCentroidOffset;
     }
 
-    private void assignSpilled(
+    /** Assign vectors from {@code start} to {@code end} to the SOAR centroid. */
+    protected static void assignSpilledSlice(
         FloatVectorValues vectors,
         KMeansIntermediate kmeansIntermediate,
         NeighborHood[] neighborhoods,
-        float soarLambda
+        float soarLambda,
+        int start,
+        int end
     ) throws IOException {
         // SOAR uses an adjusted distance for assigning spilled documents which is
         // given by:
@@ -235,7 +273,7 @@ private void assignSpilled(
 
         float[] diffs = new float[vectors.dimension()];
         final float[] distances = new float[4];
-        for (int i = 0; i < vectors.size(); i++) {
+        for (int i = start; i < end; i++) {
             float[] vector = vectors.vectorValue(i);
             int currAssignment = assignments[i];
             float[] currentCentroid = centroids[currAssignment];
@@ -308,7 +346,7 @@ private void assignSpilled(
      *                     passing in a valid output object with a centroids array that is the size of centroids expected
      * @throws IOException is thrown if vectors is inaccessible
      */
-    void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate) throws IOException {
+    final void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate) throws IOException {
         doCluster(vectors, kMeansIntermediate, -1, -1);
     }
 
@@ -326,7 +364,7 @@ void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate) t
      *
      * @throws IOException is thrown if vectors is inaccessible or if the clustersPerNeighborhood is less than 2
      */
-    void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate, int clustersPerNeighborhood, float soarLambda)
+    final void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansIntermediate, int clustersPerNeighborhood, float soarLambda)
         throws IOException {
         if (clustersPerNeighborhood < 2) {
             throw new IllegalArgumentException("clustersPerNeighborhood must be at least 2, got [" + clustersPerNeighborhood + "]");
@@ -370,18 +408,25 @@ private void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansInterme
         }
 
         assert assignments.length == n;
-        FixedBitSet centroidChanged = new FixedBitSet(centroids.length);
+        FixedBitSet[] centroidChangedSlices = new FixedBitSet[numWorkers()];
+        for (int i = 0; i < numWorkers(); i++) {
+            centroidChangedSlices[i] = new FixedBitSet(centroids.length);
+        }
         int[] centroidCounts = new int[centroids.length];
         for (int i = 0; i < maxIterations; i++) {
             // This is potentially sampled, so we need to translate ordinals
-            if (stepLloyd(sampledVectors, translateOrd, centroids, centroidChanged, centroidCounts, assignments, neighborhoods) == false) {
+            if (stepLloyd(sampledVectors, translateOrd, centroids, centroidChangedSlices, assignments, neighborhoods)) {
+                updateCentroids(sampledVectors, translateOrd, centroids, centroidChangedSlices, centroidCounts, assignments);
+            } else {
                 break;
             }
         }
         // If we were sampled, do a once over the full set of vectors to finalize the centroids
         if (sampleSize < n || maxIterations == 0) {
             // No ordinal translation needed here, we are using the full set of vectors
-            stepLloyd(vectors, i -> i, centroids, centroidChanged, centroidCounts, assignments, neighborhoods);
+            if (stepLloyd(vectors, i -> i, centroids, centroidChangedSlices, assignments, neighborhoods)) {
+                updateCentroids(sampledVectors, translateOrd, centroids, centroidChangedSlices, centroidCounts, assignments);
+            }
         }
     }
 
@@ -396,8 +441,7 @@ private void cluster(FloatVectorValues vectors, KMeansIntermediate kMeansInterme
      */
     public static void cluster(FloatVectorValues vectors, float[][] centroids, int sampleSize, int maxIterations) throws IOException {
         KMeansIntermediate kMeansIntermediate = new KMeansIntermediate(centroids, new int[vectors.size()], vectors::ordToDoc);
-        KMeansLocal kMeans = new KMeansLocal(sampleSize, maxIterations);
+        KMeansLocal kMeans = new KMeansLocalSerial(sampleSize, maxIterations);
         kMeans.cluster(vectors, kMeansIntermediate);
     }
-
 }