IVF Hierarchical KMeans Flush & Merge

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.codec.vectors;
+
+record CentroidAssignments(int numCentroids, float[][] cachedCentroids, short[] assignments, short[] soarAssignments) {
+
+    CentroidAssignments(float[][] centroids, short[] assignments, short[] soarAssignments) {
+        this(centroids.length, centroids, assignments, soarAssignments);
+    }
+
+    CentroidAssignments(int numCentroids, short[] assignments, short[] soarAssignments) {
+        this(numCentroids, null, assignments, soarAssignments);
+    }
+}

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

@@ -32,7 +32,6 @@
 import org.apache.lucene.util.InfoStream;
 import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.core.IOUtils;
-import org.elasticsearch.core.SuppressForbidden;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -122,38 +121,34 @@ public final KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOExc
         return rawVectorDelegate;
     }
 
-    protected abstract int calculateAndWriteCentroids(
+    abstract CentroidAssignments calculateAndWriteCentroids(
         FieldInfo fieldInfo,
         FloatVectorValues floatVectorValues,
-        IndexOutput temporaryCentroidOutput,
+        IndexOutput centroidOutput,
         MergeState mergeState,
         float[] globalCentroid
     ) throws IOException;
 
-    abstract long[] buildAndWritePostingsLists(
-        FieldInfo fieldInfo,
-        CentroidAssignmentScorer scorer,
-        FloatVectorValues floatVectorValues,
-        IndexOutput postingsOutput,
-        MergeState mergeState
-    ) throws IOException;
-
-    abstract CentroidAssignmentScorer calculateAndWriteCentroids(
+    abstract CentroidAssignments calculateAndWriteCentroids(
         FieldInfo fieldInfo,
         FloatVectorValues floatVectorValues,
         IndexOutput centroidOutput,
+        InfoStream infoStream,
         float[] globalCentroid
     ) throws IOException;
 
     abstract long[] buildAndWritePostingsLists(
         FieldInfo fieldInfo,
-        InfoStream infoStream,
-        CentroidAssignmentScorer scorer,
+        CentroidSupplier centroidSupplier,
         FloatVectorValues floatVectorValues,
-        IndexOutput postingsOutput
+        IndexOutput postingsOutput,
+        InfoStream infoStream,
+        CentroidAssignments centroidAssignments
     ) throws IOException;
 
-    abstract CentroidAssignmentScorer createCentroidScorer(
+    abstract CentroidSupplier createCentroidSupplier(float[][] cachedCentroids) throws IOException;
+
+    abstract CentroidSupplier createCentroidSupplier(
         IndexInput centroidsInput,
         int numCentroids,
         FieldInfo fieldInfo,
@@ -165,33 +160,31 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
         rawVectorDelegate.flush(maxDoc, sortMap);
         for (FieldWriter fieldWriter : fieldWriters) {
             float[] globalCentroid = new float[fieldWriter.fieldInfo.getVectorDimension()];
-            // calculate global centroid
-            for (var vector : fieldWriter.delegate.getVectors()) {
-                for (int i = 0; i < globalCentroid.length; i++) {
-                    globalCentroid[i] += vector[i];
-                }
-            }
-            for (int i = 0; i < globalCentroid.length; i++) {
-                globalCentroid[i] /= fieldWriter.delegate.getVectors().size();
-            }
             // build a float vector values with random access
             final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldWriter.fieldInfo, fieldWriter.delegate, maxDoc);
             // build centroids
             long centroidOffset = ivfCentroids.alignFilePointer(Float.BYTES);
-            final CentroidAssignmentScorer centroidAssignmentScorer = calculateAndWriteCentroids(
+
+            final CentroidAssignments centroidAssignments = calculateAndWriteCentroids(
                 fieldWriter.fieldInfo,
                 floatVectorValues,
                 ivfCentroids,
+                segmentWriteState.infoStream,
                 globalCentroid
             );
+
+            CentroidSupplier centroidSupplier = createCentroidSupplier(centroidAssignments.cachedCentroids());
+
             long centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
             final long[] offsets = buildAndWritePostingsLists(
                 fieldWriter.fieldInfo,
-                segmentWriteState.infoStream,
-                centroidAssignmentScorer,
+                centroidSupplier,
                 floatVectorValues,
-                ivfClusters
+                ivfClusters,
+                segmentWriteState.infoStream,
+                centroidAssignments
             );
+            // write posting lists
             writeMeta(fieldWriter.fieldInfo, centroidOffset, centroidLength, offsets, globalCentroid);
         }
     }
@@ -250,7 +243,6 @@ static IVFVectorsReader getIVFReader(KnnVectorsReader vectorsReader, String fiel
     }
 
     @Override
-    @SuppressForbidden(reason = "require usage of Lucene's IOUtils#deleteFilesIgnoringExceptions(...)")
     public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
         rawVectorDelegate.mergeOneField(fieldInfo, mergeState);
         if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
@@ -276,26 +268,29 @@ public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) thro
                 float[] calculatedGlobalCentroid = new float[fieldInfo.getVectorDimension()];
                 final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, in, numVectors);
                 success = false;
-                CentroidAssignmentScorer centroidAssignmentScorer;
                 long centroidOffset;
                 long centroidLength;
                 String centroidTempName = null;
                 int numCentroids;
                 IndexOutput centroidTemp = null;
+                CentroidAssignments centroidAssignments;
                 try {
                     centroidTemp = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "civf_", IOContext.DEFAULT);
                     centroidTempName = centroidTemp.getName();
-                    numCentroids = calculateAndWriteCentroids(
+
+                    centroidAssignments = calculateAndWriteCentroids(
                         fieldInfo,
                         floatVectorValues,
                         centroidTemp,
                         mergeState,
                         calculatedGlobalCentroid
                     );
+                    numCentroids = centroidAssignments.numCentroids();
+
                     success = true;
                 } finally {
                     if (success == false && centroidTempName != null) {
-                        IOUtils.closeWhileHandlingException(centroidTemp);
+                        org.apache.lucene.util.IOUtils.closeWhileHandlingException(centroidTemp);
                         org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, centroidTempName);
                     }
                 }
@@ -304,27 +299,34 @@ public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) thro
                         centroidOffset = ivfCentroids.getFilePointer();
                         writeMeta(fieldInfo, centroidOffset, 0, new long[0], null);
                         CodecUtil.writeFooter(centroidTemp);
-                        IOUtils.close(centroidTemp);
+                        org.apache.lucene.util.IOUtils.close(centroidTemp);
                         return;
                     }
                     CodecUtil.writeFooter(centroidTemp);
-                    IOUtils.close(centroidTemp);
+                    org.apache.lucene.util.IOUtils.close(centroidTemp);
                     centroidOffset = ivfCentroids.alignFilePointer(Float.BYTES);
-                    try (IndexInput centroidInput = mergeState.segmentInfo.dir.openInput(centroidTempName, IOContext.DEFAULT)) {
-                        ivfCentroids.copyBytes(centroidInput, centroidInput.length() - CodecUtil.footerLength());
+                    try (IndexInput centroidsInput = mergeState.segmentInfo.dir.openInput(centroidTempName, IOContext.DEFAULT)) {
+                        ivfCentroids.copyBytes(centroidsInput, centroidsInput.length() - CodecUtil.footerLength());
                         centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
-                        centroidAssignmentScorer = createCentroidScorer(centroidInput, numCentroids, fieldInfo, calculatedGlobalCentroid);
-                        assert centroidAssignmentScorer.size() == numCentroids;
+
+                        CentroidSupplier centroidSupplier = createCentroidSupplier(
+                            centroidsInput,
+                            numCentroids,
+                            fieldInfo,
+                            calculatedGlobalCentroid
+                        );
+
                         // build a float vector values with random access
                         // build centroids
                         final long[] offsets = buildAndWritePostingsLists(
                             fieldInfo,
-                            centroidAssignmentScorer,
+                            centroidSupplier,
                             floatVectorValues,
                             ivfClusters,
-                            mergeState
+                            mergeState.infoStream,
+                            centroidAssignments
                         );
-                        assert offsets.length == centroidAssignmentScorer.size();
+                        assert offsets.length == centroidSupplier.size();
                         writeMeta(fieldInfo, centroidOffset, centroidLength, offsets, calculatedGlobalCentroid);
                     }
                 } finally {
@@ -452,8 +454,8 @@ public final long ramBytesUsed() {
 
     private record FieldWriter(FieldInfo fieldInfo, FlatFieldVectorsWriter<float[]> delegate) {}
 
-    interface CentroidAssignmentScorer {
-        CentroidAssignmentScorer EMPTY = new CentroidAssignmentScorer() {
+    interface CentroidSupplier {
+        CentroidSupplier EMPTY = new CentroidSupplier() {
             @Override
             public int size() {
                 return 0;
@@ -463,24 +465,10 @@ public int size() {
             public float[] centroid(int centroidOrdinal) {
                 throw new IllegalStateException("No centroids");
             }
-
-            @Override
-            public float score(int centroidOrdinal) {
-                throw new IllegalStateException("No centroids");
-            }
-
-            @Override
-            public void setScoringVector(float[] vector) {
-                throw new IllegalStateException("No centroids");
-            }
         };
 
         int size();
 
         float[] centroid(int centroidOrdinal) throws IOException;
-
-        void setScoringVector(float[] vector);
-
-        float score(int centroidOrdinal) throws IOException;
     }
 }