Fixing centroid merge

apache · Sep 17, 2024 · 3acf852 · 3acf852
1 parent 48a8bd0
commit 3acf852
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/...re/src/java/org/apache/lucene/codecs/lucene912/Lucene912BinaryQuantizedVectorsWriter.java b/...re/src/java/org/apache/lucene/codecs/lucene912/Lucene912BinaryQuantizedVectorsWriter.java
@@ -426,6 +426,9 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE
       final float[][] centroids;
       final float[] mergedCentroid = new float[fieldInfo.getVectorDimension()];
       int vectorCount = mergeAndRecalculateCentroids(mergeState, fieldInfo, mergedCentroid);
+      if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
+        VectorUtil.l2normalize(mergedCentroid);
+      }
       // If we have more vectors than allowed for a single cluster, we will use KMeans to cluster
       // we do an early check here to avoid the overhead of `mergeOneFieldToIndex` which might
       // not be as efficient as mergeOneField
@@ -884,7 +887,7 @@ static int mergeAndRecalculateCentroids(
       return calculateCentroid(mergeState, fieldInfo, mergedCentroid);
     } else {
       for (int j = 0; j < mergedCentroid.length; j++) {
-        mergedCentroid[j] += mergedCentroid[j] / totalVectorCount;
+        mergedCentroid[j] = mergedCentroid[j] / totalVectorCount;
       }
       if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
         VectorUtil.l2normalize(mergedCentroid);