fixed 1 off error and other cleanup

Author: john-wagster

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

@@ -29,9 +29,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.List;
 
 /**
  * Default implementation of {@link IVFVectorsWriter}. It uses {@link HierarchicalKMeans} algorithm to
@@ -131,7 +129,7 @@ CentroidSupplier createCentroidSupplier(
     }
 
     static void writeCentroidsAndPartitions(
-        List<CentroidPartition> centroidPartitions,
+        CentroidPartition[] centroidPartitions,
         float[][] centroids,
         FieldInfo fieldInfo,
         float[] globalCentroid,
@@ -144,22 +142,24 @@ static void writeCentroidsAndPartitions(
         // TODO do we want to store these distances as well for future use?
         // TODO: sort centroids by global centroid (was doing so previously here)
 
-        // write the top level partition parent nodes and their pointers to the centroids within the partition
-        // a size of 1 indicates a leaf node that did not have a parent node (orphans)
-        for (CentroidPartition centroidPartition : centroidPartitions) {
-            System.arraycopy(centroidPartition.centroid(), 0, centroidScratch, 0, centroidPartition.centroid().length);
-            OptimizedScalarQuantizer.QuantizationResult result = osq.scalarQuantize(
-                centroidScratch,
-                quantizedScratch,
-                (byte) 4,
-                globalCentroid
-            );
-            for (int i = 0; i < quantizedScratch.length; i++) {
-                quantized[i] = (byte) quantizedScratch[i];
+        if (centroidPartitions != null) {
+            // write the top level partition parent nodes and their pointers to the centroids within the partition
+            // a size of 1 indicates a leaf node that did not have a parent node (orphans)
+            for (CentroidPartition centroidPartition : centroidPartitions) {
+                System.arraycopy(centroidPartition.centroid(), 0, centroidScratch, 0, centroidPartition.centroid().length);
+                OptimizedScalarQuantizer.QuantizationResult result = osq.scalarQuantize(
+                    centroidScratch,
+                    quantizedScratch,
+                    (byte) 4,
+                    globalCentroid
+                );
+                for (int i = 0; i < quantizedScratch.length; i++) {
+                    quantized[i] = (byte) quantizedScratch[i];
+                }
+                writeQuantizedValue(centroidOutput, quantized, result);
+                centroidOutput.writeInt(centroidPartition.childOrdinal());
+                centroidOutput.writeInt(centroidPartition.size());
             }
-            writeQuantizedValue(centroidOutput, quantized, result);
-            centroidOutput.writeInt(centroidPartition.childOrdinal());
-            centroidOutput.writeInt(centroidPartition.size());
         }
 
         // write the quantized centroids which will be duplicate for orphans
@@ -242,43 +242,41 @@ CentroidAssignments calculateAndWriteCentroids(
             centroidOrds[i] = i;
         }
 
-        List<CentroidPartition> centroidPartitions = new ArrayList<>();
+        CentroidPartition[] centroidPartitions = null;
+        int partitionsCount = 0;
 
         if (centroids.length > IVFVectorsFormat.DEFAULT_VECTORS_PER_CLUSTER) {
-            List<float[]> centroidsList = Arrays.stream(centroids).toList();
-            FloatVectorValues centroidsAsFVV = FloatVectorValues.fromFloats(centroidsList, fieldInfo.getVectorDimension());
-
-            HierarchicalKMeans hierarchicalKMeans = new HierarchicalKMeans(fieldInfo.getVectorDimension());
-            KMeansResult result = hierarchicalKMeans.cluster(centroidsAsFVV, centroids.length / (int) Math.sqrt(centroids.length));
+            KMeansResult result = clusterParentCentroids(fieldInfo, centroids);
             float[][] parentCentroids = result.centroids();
             int[] parentChildAssignments = result.assignments();
-            // TODO: explore using soar assignments here as well
-            // int[] parentChildSoarAssignments = result.soarAssignments();
+            // TODO: explore soar assignments here as well
+
+            centroidPartitions = new CentroidPartition[parentCentroids.length];
 
             AssignmentArraySorter sorter = new AssignmentArraySorter(centroids, centroidOrds, parentChildAssignments);
             sorter.sort(0, centroids.length);
 
-            for (int i = 0; i < parentChildAssignments.length; i++) {
+            for (int i = 0; i < parentChildAssignments.length;) {
                 int label = parentChildAssignments[i];
                 int centroidCount = 0;
+                int childOffset = i;
                 int j = i;
                 for (; j < parentChildAssignments.length; j++) {
                     if (parentChildAssignments[j] != label) {
                         break;
                     }
                     centroidCount++;
                 }
-                int childOrdinal = i;
                 i = j;
-                centroidPartitions.add(new CentroidPartition(parentCentroids[label], childOrdinal, centroidCount));
+                centroidPartitions[partitionsCount++] = new CentroidPartition(parentCentroids[label], childOffset, centroidCount);
             }
         }
 
         writeCentroidsAndPartitions(centroidPartitions, centroids, fieldInfo, globalCentroid, centroidOutput);
 
         if (logger.isDebugEnabled()) {
             logger.debug("calculate centroids and assign vectors time ms: {}", (System.nanoTime() - nanoTime) / 1000000.0);
-            logger.debug("final parent centroid count {}: ", centroidPartitions.size());
+            logger.debug("final parent centroid count {}: ", partitionsCount);
             logger.debug("final centroid count: {}", centroids.length);
         }
 
@@ -291,7 +289,40 @@ CentroidAssignments calculateAndWriteCentroids(
         int[] soarAssignments = kMeansResult.soarAssignments();
 
         int[][] assignmentsByCluster = buildCentroidAssignments(centroids.length, assignments, soarAssignments, centroidOrdsToIdx);
-        return new CentroidAssignments(centroidPartitions.size(), centroids, assignmentsByCluster);
+        return new CentroidAssignments(partitionsCount, centroids, assignmentsByCluster);
+    }
+
+    private KMeansResult clusterParentCentroids(FieldInfo fieldInfo, float[][] centroids) throws IOException {
+        FloatVectorValues centroidsAsFVV = new FloatVectorValues() {
+            @Override
+            public int size() {
+                return centroids.length;
+            }
+
+            @Override
+            public int dimension() {
+                return fieldInfo.getVectorDimension();
+            }
+
+            @Override
+            public float[] vectorValue(int targetOrd) {
+                return centroids[targetOrd];
+            }
+
+            @Override
+            public FloatVectorValues copy() {
+                return this;
+            }
+
+            @Override
+            public DocIndexIterator iterator() {
+                return createDenseIterator();
+            }
+        };
+
+        HierarchicalKMeans hierarchicalKMeans = new HierarchicalKMeans(fieldInfo.getVectorDimension());
+        KMeansResult result = hierarchicalKMeans.cluster(centroidsAsFVV, centroids.length / (int) Math.sqrt(centroids.length));
+        return result;
     }
 
     static int[][] buildCentroidAssignments(int centroidCount, int[] assignments, int[] soarAssignments, IntIntMap centroidOrds) {

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

@@ -46,6 +46,8 @@
  */
 public abstract class IVFVectorsReader extends KnnVectorsReader {
 
+    private static final float PARENT_SCORE_BUFFER = 0.05f;
+
     private final IndexInput ivfCentroids, ivfClusters;
     private final SegmentReadState state;
     private final FieldInfos fieldInfos;
@@ -311,7 +313,7 @@ public final void search(String field, float[] target, KnnCollector knnCollector
                 // ... would be centroidScore < (nextParentScore + furthestCentroidScore) which is better than just a buffer
                 // TODO: try a ParentNProbe here that's for instance the sqrt(nProbe) that forces a fixed
                 // ... number of parents to be explored at each step
-                while (parentCentroidQueue.size() > 0 && centroidScore < (nextParentScore + nextParentScore * 0.05)) {
+                while (parentCentroidQueue.size() > 0 && centroidScore < (nextParentScore + nextParentScore * PARENT_SCORE_BUFFER)) {
                     updateCentroidQueueWNextParent(parentCentroidQueryScorer, parentCentroidQueue, centroidQueryScorer, centroidQueue);
                     if (parentCentroidQueue.size() > 0) {
                         nextParentScore = parentCentroidQueue.topScore();
@@ -358,7 +360,7 @@ private static int updateCentroidQueueWNextParent(
             childCentroidOrdinal = parentCentroidQueryScorer.getChildCentroidStart(parentCentroidOrdinal);
             childCentroidCount = parentCentroidQueryScorer.getChildCount(parentCentroidOrdinal);
         }
-        // TODO: add back scorePostingLists? seems like it's not doing anything at this point
+        // TODO: add back scorePostingLists? or make this function abstract? or break it apart? not sure
         centroidQueryScorer.bulkScore(centroidQueue, childCentroidOrdinal, childCentroidOrdinal + childCentroidCount);
 
         return childCentroidCount;

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

@@ -66,7 +66,7 @@ public KMeansResult cluster(FloatVectorValues vectors, int targetSize) throws IO
         }
 
         // partition the space
-        KMeansIntermediate kMeansIntermediate = clusterAndSplit(vectors, targetSize, 0);
+        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);
@@ -76,7 +76,7 @@ public KMeansResult cluster(FloatVectorValues vectors, int targetSize) throws IO
         return kMeansIntermediate;
     }
 
-    KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int targetSize, final int depth) throws IOException {
+    KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int targetSize) throws IOException {
         if (vectors.size() <= targetSize) {
             return new KMeansIntermediate();
         }
@@ -118,7 +118,7 @@ KMeansIntermediate clusterAndSplit(final FloatVectorValues vectors, final int ta
                 // TODO: consider iterative here instead of recursive
                 // recursive call to build out the sub partitions around this centroid c
                 // subsequently reconcile and flatten the space of all centroids and assignments into one structure we can return
-                updateAssignmentsWithRecursiveSplit(kMeansIntermediate, c, clusterAndSplit(sample, targetSize, depth + 1), depth);
+                updateAssignmentsWithRecursiveSplit(kMeansIntermediate, c, clusterAndSplit(sample, targetSize));
             }
         }
 
@@ -138,7 +138,7 @@ static FloatVectorValues createClusterSlice(int clusterSize, int cluster, FloatV
         return new FloatVectorValuesSlice(vectors, slice);
     }
 
-    void updateAssignmentsWithRecursiveSplit(KMeansIntermediate current, int cluster, KMeansIntermediate subPartitions, int depth) {
+    void updateAssignmentsWithRecursiveSplit(KMeansIntermediate current, int cluster, KMeansIntermediate subPartitions) {
         int orgCentroidsSize = current.centroids().length;
         int newCentroidsSize = current.centroids().length + subPartitions.centroids().length - 1;