fixing test failure, addressing pr comments

benwtrent · benwtrent · commit a4f2c28b7690 · 2025-05-05T09:04:38.000-04:00
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java
@@ -223,44 +223,17 @@ static void writeCentroids(float[][] centroids, FieldInfo fieldInfo, float[] glo
         }
     }
 
-    record SegmentCentroid(int segment, int centroid, int centroidSize) {}
-
-    @Override
-    protected int calculateAndWriteCentroids(
+    static float[][] gatherInitCentroids(
+        List<FloatVectorValues> centroidList,
+        List<SegmentCentroid> segmentCentroids,
+        int desiredClusters,
         FieldInfo fieldInfo,
-        FloatVectorValues floatVectorValues,
-        IndexOutput temporaryCentroidOutput,
-        MergeState mergeState,
-        float[] globalCentroid
+        MergeState mergeState
     ) throws IOException {
-        if (floatVectorValues.size() == 0) {
-            return 0;
+        if (centroidList.size() == 0) {
+            return null;
         }
-        int desiredClusters = ((floatVectorValues.size() - 1) / vectorPerCluster) + 1;
-        // init centroids from merge state
-        List<FloatVectorValues> centroidList = new ArrayList<>();
-        List<SegmentCentroid> segmentCentroids = new ArrayList<>(desiredClusters);
-
-        int segmentIdx = 0;
         long startTime = System.nanoTime();
-        for (var reader : mergeState.knnVectorsReaders) {
-            IVFVectorsReader ivfVectorsReader = IVFVectorsFormat.getIVFReader(reader, fieldInfo.name);
-            if (ivfVectorsReader == null) {
-                continue;
-            }
-
-            FloatVectorValues centroid = ivfVectorsReader.getCentroids(fieldInfo);
-            if (centroid == null) {
-                continue;
-            }
-            centroidList.add(centroid);
-            for (int i = 0; i < centroid.size(); i++) {
-                int size = ivfVectorsReader.centroidSize(fieldInfo.name, i);
-                segmentCentroids.add(new SegmentCentroid(segmentIdx, i, size));
-            }
-            segmentIdx++;
-        }
-
         // sort centroid list by floatvector size
         FloatVectorValues baseSegment = centroidList.get(0);
         for (var l : centroidList) {
@@ -334,6 +307,9 @@ protected int calculateAndWriteCentroids(
             sum[label - 1] += segmentCentroid.centroidSize;
         }
         for (int i = 0; i < initCentroids.length; i++) {
+            if (sum[i] == 0 || sum[i] == 1) {
+                continue;
+            }
             for (int j = 0; j < initCentroids[i].length; j++) {
                 initCentroids[i][j] /= sum[i];
             }
@@ -348,6 +324,67 @@ protected int calculateAndWriteCentroids(
                 "Gathered initCentroids:" + initCentroids.length + " for desired: " + desiredClusters
             );
         }
+        return initCentroids;
+    }
+
+    record SegmentCentroid(int segment, int centroid, int centroidSize) {}
+
+    /**
+     * Calculate the centroids for the given field and write them to the given
+     * temporary centroid output.
+     * When merging, we first bootstrap the KMeans algorithm with the centroids contained in the merging segments.
+     * To prevent centroids that are too similar from having an outsized impact, all centroids that are closer than
+     * the largest segments intra-cluster distance are merged into a single centroid.
+     * The resulting centroids are then used to initialize the KMeans algorithm.
+     *
+     * @param fieldInfo merging field info
+     * @param floatVectorValues the float vector values to merge
+     * @param temporaryCentroidOutput the temporary centroid output
+     * @param mergeState the merge state
+     * @param globalCentroid the global centroid, calculated by this method and used to quantize the centroids
+     * @return the number of centroids written
+     * @throws IOException if an I/O error occurs
+     */
+    @Override
+    protected int calculateAndWriteCentroids(
+        FieldInfo fieldInfo,
+        FloatVectorValues floatVectorValues,
+        IndexOutput temporaryCentroidOutput,
+        MergeState mergeState,
+        float[] globalCentroid
+    ) throws IOException {
+        if (floatVectorValues.size() == 0) {
+            return 0;
+        }
+        int maxNumClusters = ((floatVectorValues.size() - 1) / vectorPerCluster) + 1;
+        int desiredClusters = (int) Math.max(Math.sqrt(floatVectorValues.size()), maxNumClusters);
+        // init centroids from merge state
+        List<FloatVectorValues> centroidList = new ArrayList<>();
+        List<SegmentCentroid> segmentCentroids = new ArrayList<>(desiredClusters);
+
+        int segmentIdx = 0;
+        for (var reader : mergeState.knnVectorsReaders) {
+            IVFVectorsReader ivfVectorsReader = IVFVectorsFormat.getIVFReader(reader, fieldInfo.name);
+            if (ivfVectorsReader == null) {
+                continue;
+            }
+
+            FloatVectorValues centroid = ivfVectorsReader.getCentroids(fieldInfo);
+            if (centroid == null) {
+                continue;
+            }
+            centroidList.add(centroid);
+            for (int i = 0; i < centroid.size(); i++) {
+                int size = ivfVectorsReader.centroidSize(fieldInfo.name, i);
+                if (size == 0) {
+                    continue;
+                }
+                segmentCentroids.add(new SegmentCentroid(segmentIdx, i, size));
+            }
+            segmentIdx++;
+        }
+
+        float[][] initCentroids = gatherInitCentroids(centroidList, segmentCentroids, desiredClusters, fieldInfo, mergeState);
 
         // FIXME: run a custom version of KMeans that is just better...
         long nanoTime = System.nanoTime();
@@ -369,6 +406,15 @@ protected int calculateAndWriteCentroids(
         float[][] centroids = kMeans.centroids();
 
         // write them
+        // calculate the global centroid from all the centroids:
+        for (float[] centroid : centroids) {
+            for (int j = 0; j < centroid.length; j++) {
+                globalCentroid[j] += centroid[j];
+            }
+        }
+        for (int j = 0; j < globalCentroid.length; j++) {
+            globalCentroid[j] /= centroids.length;
+        }
         writeCentroids(centroids, fieldInfo, globalCentroid, temporaryCentroidOutput);
         return centroids.length;
     }
@@ -477,14 +523,11 @@ static void assignCentroids(CentroidAssignmentScorer scorer, FloatVectorValues v
                 // pop the best
                 int sz = neighborsToCheck.size();
                 int best = neighborsToCheck.consumeNodesAndScoresMin(ordScoreIterator.ords, ordScoreIterator.scores);
-                // reset the ordScoreIterator as it has consumed the ords and scores
-                ordScoreIterator.idx = sz;
+                // Set the size to the number of neighbors we actually found
+                ordScoreIterator.setSize(sz);
                 bestScore = ordScoreIterator.getScore(best);
                 bestCentroid = ordScoreIterator.getOrd(best);
             }
-            if (clusters[bestCentroid] == null) {
-                clusters[bestCentroid] = new IntArrayList(16);
-            }
             clusters[bestCentroid].add(docID);
             if (soarClusterCheckCount > 0) {
                 assignCentroidSOAR(
@@ -495,7 +538,7 @@ static void assignCentroids(CentroidAssignmentScorer scorer, FloatVectorValues v
                     bestScore,
                     scratch,
                     scorer,
-                    vectors,
+                    vector,
                     clusters
                 );
             }
@@ -511,10 +554,9 @@ static void assignCentroidSOAR(
         float bestScore,
         float[] scratch,
         CentroidAssignmentScorer scorer,
-        FloatVectorValues vectors,
+        float[] vector,
         IntArrayList[] clusters
     ) throws IOException {
-        float[] vector = vectors.vectorValue(vecOrd);
         ESVectorUtil.subtract(vector, bestCentroid, scratch);
         int bestSecondaryCentroid = -1;
         float minDist = Float.MAX_VALUE;
@@ -546,6 +588,14 @@ static class OrdScoreIterator {
             this.scores = new float[size];
         }
 
+        int setSize(int size) {
+            if (size > ords.length) {
+                throw new IllegalArgumentException("size must be <= " + ords.length);
+            }
+            this.idx = size;
+            return size;
+        }
+
         int getOrd(int idx) {
             return ords[idx];
         }
@@ -606,15 +656,15 @@ static class OffHeapCentroidAssignmentScorer implements CentroidAssignmentScorer
         private final int dimension;
         private final float[] scratch;
         private float[] q;
-        private final long centroidByteSize;
+        private final long rawCentroidOffset;
         private int currOrd = -1;
 
         OffHeapCentroidAssignmentScorer(IndexInput centroidsInput, int numCentroids, FieldInfo info) {
             this.centroidsInput = centroidsInput;
             this.numCentroids = numCentroids;
             this.dimension = info.getVectorDimension();
             this.scratch = new float[dimension];
-            this.centroidByteSize = dimension + 3 * Float.BYTES + Short.BYTES;
+            this.rawCentroidOffset = (dimension + 3 * Float.BYTES + Short.BYTES) * numCentroids;
         }
 
         @Override
@@ -627,7 +677,7 @@ public float[] centroid(int centroidOrdinal) throws IOException {
             if (centroidOrdinal == currOrd) {
                 return scratch;
             }
-            centroidsInput.seek(numCentroids * centroidByteSize + (long) centroidOrdinal * dimension * Float.BYTES);
+            centroidsInput.seek(rawCentroidOffset + (long) centroidOrdinal * dimension * Float.BYTES);
             centroidsInput.readFloats(scratch, 0, dimension);
             this.currOrd = centroidOrdinal;
             return scratch;
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsWriter.java
@@ -192,7 +192,6 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
                 floatVectorValues,
                 ivfClusters
             );
-            // write posting lists
             writeMeta(fieldWriter.fieldInfo, centroidOffset, centroidLength, offsets, globalCentroid);
         }
     }
@@ -256,54 +255,25 @@ public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) thro
         rawVectorDelegate.mergeOneField(fieldInfo, mergeState);
         if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
             final int numVectors;
-            String name = null;
+            String tempRawVectorsFileName = null;
             boolean success = false;
             // build a float vector values with random access. In order to do that we dump the vectors to
             // a temporary file
             // and write the docID follow by the vector
             try (IndexOutput out = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "ivf_", IOContext.DEFAULT)) {
-                name = out.getName();
+                tempRawVectorsFileName = out.getName();
                 // TODO do this better, we shouldn't have to write to a temp file, we should be able to
-                // to just from the merged vector values.
+                // to just from the merged vector values, the tricky part is the random access.
                 numVectors = writeFloatVectorValues(fieldInfo, out, MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState));
+                CodecUtil.writeFooter(out);
                 success = true;
             } finally {
-                if (success == false && name != null) {
-                    org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, name);
+                if (success == false && tempRawVectorsFileName != null) {
+                    org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, tempRawVectorsFileName);
                 }
             }
-            float[] globalCentroid = new float[fieldInfo.getVectorDimension()];
-            int vectorCount = 0;
-            for (int idx = 0; idx < mergeState.knnVectorsReaders.length; idx++) {
-                if (mergeState.fieldInfos[idx] == null
-                    || mergeState.fieldInfos[idx].hasVectorValues() == false
-                    || mergeState.fieldInfos[idx].fieldInfo(fieldInfo.name) == null
-                    || mergeState.fieldInfos[idx].fieldInfo(fieldInfo.name).hasVectorValues() == false) {
-                    continue;
-                }
-                KnnVectorsReader knnReaders = mergeState.knnVectorsReaders[idx];
-                IVFVectorsReader ivfReader = getIVFReader(knnReaders, fieldInfo.name);
-                if (ivfReader != null) {
-                    FloatVectorValues floatVectorValues = knnReaders.getFloatVectorValues(fieldInfo.name);
-                    if (floatVectorValues == null) {
-                        continue;
-                    }
-                    int numVecs = floatVectorValues.size();
-                    float[] readerGlobalCentroid = ivfReader.getGlobalCentroid(fieldInfo);
-                    if (readerGlobalCentroid != null) {
-                        vectorCount += numVecs;
-                        for (int i = 0; i < globalCentroid.length; i++) {
-                            globalCentroid[i] += readerGlobalCentroid[i] * numVecs;
-                        }
-                    }
-                }
-            }
-            if (vectorCount > 0) {
-                for (int i = 0; i < globalCentroid.length; i++) {
-                    globalCentroid[i] /= vectorCount;
-                }
-            }
-            try (IndexInput in = mergeState.segmentInfo.dir.openInput(name, IOContext.DEFAULT)) {
+            try (IndexInput in = mergeState.segmentInfo.dir.openInput(tempRawVectorsFileName, IOContext.DEFAULT)) {
+                float[] calculatedGlobalCentroid = new float[fieldInfo.getVectorDimension()];
                 final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, in, numVectors);
                 success = false;
                 CentroidAssignmentScorer centroidAssignmentScorer;
@@ -315,7 +285,13 @@ public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) thro
                 try {
                     centroidTemp = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "civf_", IOContext.DEFAULT);
                     centroidTempName = centroidTemp.getName();
-                    numCentroids = calculateAndWriteCentroids(fieldInfo, floatVectorValues, centroidTemp, mergeState, globalCentroid);
+                    numCentroids = calculateAndWriteCentroids(
+                        fieldInfo,
+                        floatVectorValues,
+                        centroidTemp,
+                        mergeState,
+                        calculatedGlobalCentroid
+                    );
                     success = true;
                 } finally {
                     if (success == false && centroidTempName != null) {
@@ -337,7 +313,7 @@ public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) thro
                     try (IndexInput centroidInput = mergeState.segmentInfo.dir.openInput(centroidTempName, IOContext.DEFAULT)) {
                         ivfCentroids.copyBytes(centroidInput, centroidInput.length() - CodecUtil.footerLength());
                         centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
-                        centroidAssignmentScorer = createCentroidScorer(centroidInput, numCentroids, fieldInfo, globalCentroid);
+                        centroidAssignmentScorer = createCentroidScorer(centroidInput, numCentroids, fieldInfo, calculatedGlobalCentroid);
                         assert centroidAssignmentScorer.size() == numCentroids;
                         // build a float vector values with random access
                         // build centroids
@@ -348,20 +324,18 @@ public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) thro
                             ivfClusters,
                             mergeState
                         );
-                        // write posting lists
-
-                        // TODO handle this correctly by creating new centroid
-                        if (vectorCount == 0 && offsets.length > 0) {
-                            throw new IllegalStateException("No global centroid found for field: " + fieldInfo.name);
-                        }
                         assert offsets.length == centroidAssignmentScorer.size();
-                        writeMeta(fieldInfo, centroidOffset, centroidLength, offsets, globalCentroid);
+                        writeMeta(fieldInfo, centroidOffset, centroidLength, offsets, calculatedGlobalCentroid);
                     }
                 } finally {
-                    org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, name, centroidTempName);
+                    org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(
+                        mergeState.segmentInfo.dir,
+                        tempRawVectorsFileName,
+                        centroidTempName
+                    );
                 }
             } finally {
-                org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, name);
+                org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, tempRawVectorsFileName);
             }
         }
     }
