[IVF] Simplify how we buildCentroidAssignments

Author: iverase

server/src/main/java/org/elasticsearch/index/codec/vectors/CentroidAssignments.java

@@ -9,36 +9,10 @@
 
 package org.elasticsearch.index.codec.vectors;
 
-final class CentroidAssignments {
-
-    private final int numCentroids;
-    private final float[][] cachedCentroids;
-    private final int[][] assignmentsByCluster;
-
-    private CentroidAssignments(int numCentroids, float[][] cachedCentroids, int[][] assignmentsByCluster) {
-        this.numCentroids = numCentroids;
-        this.cachedCentroids = cachedCentroids;
-        this.assignmentsByCluster = assignmentsByCluster;
-    }
+record CentroidAssignments(int numCentroids, float[][] centroids, int[][] assignmentsByCluster) {
 
     CentroidAssignments(float[][] centroids, int[][] assignmentsByCluster) {
         this(centroids.length, centroids, assignmentsByCluster);
-    }
-
-    CentroidAssignments(int numCentroids, int[][] assignmentsByCluster) {
-        this(numCentroids, null, assignmentsByCluster);
-    }
-
-    // Getters and setters
-    public int numCentroids() {
-        return numCentroids;
-    }
-
-    public float[][] cachedCentroids() {
-        return cachedCentroids;
-    }
-
-    public int[][] assignmentsByCluster() {
-        return assignmentsByCluster;
+        assert centroids.length == assignmentsByCluster.length;
     }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java

@@ -148,6 +148,7 @@ static void writeCentroids(float[][] centroids, FieldInfo fieldInfo, float[] glo
         }
     }
 
+    @Override
     CentroidAssignments calculateAndWriteCentroids(
         FieldInfo fieldInfo,
         FloatVectorValues floatVectorValues,
@@ -156,16 +157,7 @@ CentroidAssignments calculateAndWriteCentroids(
         float[] globalCentroid
     ) throws IOException {
         // TODO: take advantage of prior generated clusters from mergeState in the future
-        return calculateAndWriteCentroids(fieldInfo, floatVectorValues, centroidOutput, globalCentroid, false);
-    }
-
-    CentroidAssignments calculateAndWriteCentroids(
-        FieldInfo fieldInfo,
-        FloatVectorValues floatVectorValues,
-        IndexOutput centroidOutput,
-        float[] globalCentroid
-    ) throws IOException {
-        return calculateAndWriteCentroids(fieldInfo, floatVectorValues, centroidOutput, globalCentroid, true);
+        return calculateAndWriteCentroids(fieldInfo, floatVectorValues, centroidOutput, globalCentroid);
     }
 
     /**
@@ -176,26 +168,22 @@ CentroidAssignments calculateAndWriteCentroids(
      * @param floatVectorValues the float vector values to merge
      * @param centroidOutput the centroid output
      * @param globalCentroid the global centroid, calculated by this method and used to quantize the centroids
-     * @param cacheCentroids whether the centroids are kept or discarded once computed
      * @return the vector assignments, soar assignments, and if asked the centroids themselves that were computed
      * @throws IOException if an I/O error occurs
      */
+    @Override
     CentroidAssignments calculateAndWriteCentroids(
         FieldInfo fieldInfo,
         FloatVectorValues floatVectorValues,
         IndexOutput centroidOutput,
-        float[] globalCentroid,
-        boolean cacheCentroids
-    ) throws IOException {
+        float[] globalCentroid
+    )       throws IOException {
 
         long nanoTime = System.nanoTime();
 
         // TODO: consider hinting / bootstrapping hierarchical kmeans with the prior segments centroids
         KMeansResult kMeansResult = new HierarchicalKMeans(floatVectorValues.dimension()).cluster(floatVectorValues, vectorPerCluster);
         float[][] centroids = kMeansResult.centroids();
-        int[] assignments = kMeansResult.assignments();
-        int[] soarAssignments = kMeansResult.soarAssignments();
-
         // TODO: for flush we are doing this over the vectors and here centroids which seems duplicative
         // preliminary tests suggest recall is good using only centroids but need to do further evaluation
         // TODO: push this logic into vector util?
@@ -215,7 +203,13 @@ CentroidAssignments calculateAndWriteCentroids(
             logger.debug("calculate centroids and assign vectors time ms: {}", (System.nanoTime() - nanoTime) / 1000000.0);
             logger.debug("final centroid count: {}", centroids.length);
         }
+        return buildCentroidAssignments(kMeansResult);
+    }
 
+    static CentroidAssignments buildCentroidAssignments(KMeansResult kMeansResult) {
+        float[][] centroids = kMeansResult.centroids();
+        int[] assignments = kMeansResult.assignments();
+        int[] soarAssignments = kMeansResult.soarAssignments();
         int[] centroidVectorCount = new int[centroids.length];
         for (int i = 0; i < assignments.length; i++) {
             centroidVectorCount[assignments[i]]++;
@@ -242,12 +236,7 @@ CentroidAssignments calculateAndWriteCentroids(
                 }
             }
         }
-
-        if (cacheCentroids) {
-            return new CentroidAssignments(centroids, assignmentsByCluster);
-        } else {
-            return new CentroidAssignments(centroids.length, assignmentsByCluster);
-        }
+        return new CentroidAssignments(centroids, assignmentsByCluster);
     }
 
     static void writeQuantizedValue(IndexOutput indexOutput, byte[] binaryValue, OptimizedScalarQuantizer.QuantizationResult corrections)

server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsWriter.java

@@ -166,7 +166,7 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
                 globalCentroid
             );
 
-            CentroidSupplier centroidSupplier = new OnHeapCentroidSupplier(centroidAssignments.cachedCentroids());
+            CentroidSupplier centroidSupplier = new OnHeapCentroidSupplier(centroidAssignments.centroids());
 
             long centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
             final long[] offsets = buildAndWritePostingsLists(
@@ -280,26 +280,27 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
             IndexInput docs = docsFileName == null ? null : mergeState.segmentInfo.dir.openInput(docsFileName, IOContext.DEFAULT)
         ) {
             final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, docs, vectors, numVectors);
-            success = false;
-            long centroidOffset;
-            long centroidLength;
+
+            final long centroidOffset;
+            final long centroidLength;
+            final int numCentroids;
+            final int[][] assignmentsByCluster;
+            final float[] calculatedGlobalCentroid = new float[fieldInfo.getVectorDimension()];
             String centroidTempName = null;
-            int numCentroids;
             IndexOutput centroidTemp = null;
-            CentroidAssignments centroidAssignments;
-            float[] calculatedGlobalCentroid = new float[fieldInfo.getVectorDimension()];
+            success = false;
             try {
                 centroidTemp = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "civf_", IOContext.DEFAULT);
                 centroidTempName = centroidTemp.getName();
-                centroidAssignments = calculateAndWriteCentroids(
+                CentroidAssignments centroidAssignments = calculateAndWriteCentroids(
                     fieldInfo,
-                    floatVectorValues,
+                    getFloatVectorValues(fieldInfo, docs, vectors, numVectors),
                     centroidTemp,
                     mergeState,
                     calculatedGlobalCentroid
                 );
                 numCentroids = centroidAssignments.numCentroids();
-
+                assignmentsByCluster = centroidAssignments.assignmentsByCluster();
                 success = true;
             } finally {
                 if (success == false && centroidTempName != null) {
@@ -336,7 +337,7 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
                         centroidSupplier,
                         floatVectorValues,
                         ivfClusters,
-                        centroidAssignments.assignmentsByCluster()
+                        assignmentsByCluster
                     );
                     assert offsets.length == centroidSupplier.size();
                     writeMeta(fieldInfo, centroidOffset, centroidLength, offsets, calculatedGlobalCentroid);