using full hkmeans to gen parents

john-wagster · john-wagster · commit 3a2ba615ba47 · 2025-07-07T15:34:15.000-05:00
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/AssignmentArraySorter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/AssignmentArraySorter.java
@@ -9,34 +9,28 @@
 
 package org.elasticsearch.index.codec.vectors;
 
-import org.apache.lucene.util.TimSorter;
-
-class AssignmentArraySorter extends TimSorter {
+import org.apache.lucene.util.IntroSorter;
 
+class AssignmentArraySorter extends IntroSorter {
+    int pivot = -1;
     private final float[][] centroids;
-    private final float[][] tmpC;
-
     private final int[] centroidsOrds;
-    private final int[] tmpA;
-
     private final int[] sortOrdering;
-    private final int[] tmpSort;
 
     AssignmentArraySorter(float[][] centroids, int[] centroidsOrds, int[] sortOrdering) {
-        super(centroids.length / 64);
         this.centroids = centroids;
         this.centroidsOrds = centroidsOrds;
         this.sortOrdering = sortOrdering;
+    }
 
-        int maxTempSlots = centroids.length / 64;
-        this.tmpC = new float[maxTempSlots][];
-        this.tmpA = new int[maxTempSlots];
-        this.tmpSort = new int[maxTempSlots];
+    @Override
+    protected void setPivot(int i) {
+        pivot = sortOrdering[i];
     }
 
     @Override
-    protected int compare(int i, int j) {
-        return Integer.compare(sortOrdering[i], sortOrdering[j]);
+    protected int comparePivot(int j) {
+        return Integer.compare(pivot, sortOrdering[j]);
     }
 
     @Override
@@ -53,30 +47,4 @@ protected void swap(int i, int j) {
         sortOrdering[i] = sortOrdering[j];
         sortOrdering[j] = tmpSort;
     }
-
-    @Override
-    protected void copy(int src, int dest) {
-        centroids[dest] = centroids[src];
-        centroidsOrds[dest] = centroidsOrds[src];
-        sortOrdering[dest] = sortOrdering[src];
-    }
-
-    @Override
-    protected void save(int start, int len) {
-        System.arraycopy(centroids, start, tmpC, 0, len);
-        System.arraycopy(centroidsOrds, start, tmpA, 0, len);
-        System.arraycopy(sortOrdering, start, tmpSort, 0, len);
-    }
-
-    @Override
-    protected void restore(int src, int dest) {
-        centroids[dest] = tmpC[src];
-        centroidsOrds[dest] = tmpA[src];
-        sortOrdering[dest] = tmpSort[src];
-    }
-
-    @Override
-    protected int compareSaved(int i, int j) {
-        return Integer.compare(tmpSort[i], sortOrdering[j]);
-    }
 }
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java
@@ -167,17 +167,19 @@ ParentCentroidQueryScorer getParentCentroidScorer(
         float[] targetQuery
     ) throws IOException {
         FieldEntry fieldEntry = fields.get(fieldInfo.number);
-        float[] globalCentroid = fieldEntry.globalCentroid();
         float globalCentroidDp = fieldEntry.globalCentroidDp();
         OptimizedScalarQuantizer scalarQuantizer = new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
-        byte[] quantized = new byte[targetQuery.length];
-        float[] targetScratch = ArrayUtil.copyArray(targetQuery);
-        OptimizedScalarQuantizer.QuantizationResult queryParams = scalarQuantizer.scalarQuantize(
-            targetScratch,
-            quantized,
+        final int[] scratch = new int[targetQuery.length];
+        final OptimizedScalarQuantizer.QuantizationResult queryParams = scalarQuantizer.scalarQuantize(
+            ArrayUtil.copyArray(targetQuery),
+            scratch,
             (byte) 4,
-            globalCentroid
+            fieldEntry.globalCentroid()
         );
+        final byte[] quantized = new byte[targetQuery.length];
+        for (int i = 0; i < quantized.length; i++) {
+            quantized[i] = (byte) scratch[i];
+        }
         final ES91Int4VectorsScorer scorer = ESVectorUtil.getES91Int4VectorsScorer(centroids, fieldInfo.getVectorDimension());
         return new ParentCentroidQueryScorer() {
             int currentCentroid = -1;
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
@@ -9,6 +9,9 @@
 
 package org.elasticsearch.index.codec.vectors;
 
+import com.carrotsearch.hppc.IntIntHashMap;
+import com.carrotsearch.hppc.IntIntMap;
+
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
@@ -30,8 +33,6 @@
 import java.util.Arrays;
 import java.util.List;
 
-import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.DEFAULT_VECTORS_PER_CLUSTER;
-
 /**
  * Default implementation of {@link IVFVectorsWriter}. It uses {@link HierarchicalKMeans} algorithm to
  * partition the vector space, and then stores the centroids and posting list in a sequential
@@ -153,7 +154,10 @@ static void writeCentroidsAndPartitions(
                 (byte) 4,
                 globalCentroid
             );
-            writeQuantizedValue(centroidOutput, quantizedScratch, result);
+            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());
         }
@@ -252,40 +256,32 @@ CentroidAssignments calculateAndWriteCentroids(
 
         List<CentroidPartition> centroidPartitions = new ArrayList<>();
 
-        // TODO: make this configurable
-        if (centroids.length > DEFAULT_VECTORS_PER_CLUSTER) {
-            // TODO: sort by global centroids as well
-            // TODO: have this take a function instead of just an int[] for sorting
-            AssignmentArraySorter sorter = new AssignmentArraySorter(centroids, centroidOrds, kMeansResult.parentLayer());
-            sorter.sort(0, centroids.length);
-
-            for (int i = 0; i < kMeansResult.parentLayer().length;) {
-                // for any layer that was not partitioned we treat it duplicatively as a parent and child
-                if (kMeansResult.parentLayer()[i] == -1) {
-                    centroidPartitions.add(new CentroidPartition(centroids[i], i, 1));
-                    i++;
-                } else {
-                    int label = kMeansResult.parentLayer()[i];
-                    int centroidCount = 0;
-                    float[] parentPartitionCentroid = new float[fieldInfo.getVectorDimension()];
-                    int j = i;
-                    for (; j < kMeansResult.parentLayer().length; j++) {
-                        if (kMeansResult.parentLayer()[j] != label) {
-                            break;
-                        }
-                        for (int k = 0; k < parentPartitionCentroid.length; k++) {
-                            parentPartitionCentroid[k] += centroids[i][k];
-                        }
-                        centroidCount++;
-                    }
-                    int childOrdinal = i;
-                    i = j;
-                    for (int d = 0; d < parentPartitionCentroid.length; d++) {
-                        parentPartitionCentroid[d] /= centroidCount;
-                    }
-                    centroidPartitions.add(new CentroidPartition(parentPartitionCentroid, childOrdinal, centroidCount));
+        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));
+        float[][] parentCentroids = result.centroids();
+        int[] parentChildAssignments = result.assignments();
+        // TODO: explore using soar assignments here as well
+        //int[] parentChildSoarAssignments = result.soarAssignments();
+
+        AssignmentArraySorter sorter = new AssignmentArraySorter(centroids, centroidOrds, parentChildAssignments);
+        sorter.sort(0, centroids.length);
+
+        for(int i = 0; i < parentChildAssignments.length; i++) {
+            int label = parentChildAssignments[i];
+            int centroidCount = 0;
+            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));
         }
 
         writeCentroidsAndPartitions(centroidPartitions, centroids, fieldInfo, globalCentroid, centroidOutput);
@@ -298,7 +294,11 @@ CentroidAssignments calculateAndWriteCentroids(
             logger.debug("final centroid count: {}", centroids.length);
         }
 
-        int[][] assignmentsByCluster = mapAssignmentsByCluster(centroids.length, assignments, soarAssignments, centroidOrds);
+        IntIntMap centroidOrdsToIdx = new IntIntHashMap(centroidOrds.length);
+        for(int i = 0; i < centroidOrds.length; i++) {
+            centroidOrdsToIdx.put(centroidOrds[i], i);
+        }
+        int[][] assignmentsByCluster = mapAssignmentsByCluster(centroids.length, assignments, soarAssignments, centroidOrdsToIdx);
 
         if (cacheCentroids) {
             return new CentroidAssignments(centroidPartitions.size(), centroids, assignmentsByCluster);
@@ -307,26 +307,14 @@ CentroidAssignments calculateAndWriteCentroids(
         }
     }
 
-    // FIXME: clean this up
-    static int[][] mapAssignmentsByCluster(int centroidCount, int[] assignments, int[] soarAssignments, int[] centroidOrds) {
+    static int[][] mapAssignmentsByCluster(int centroidCount, int[] assignments, int[] soarAssignments, IntIntMap centroidOrds) {
         int[] centroidVectorCount = new int[centroidCount];
         for (int i = 0; i < assignments.length; i++) {
-            int c = -1;
-            // FIXME: create a reverse mapping prior to this step? .. expensive
-            for (int j = 0; j < centroidOrds.length; j++) {
-                if (assignments[i] == centroidOrds[j]) {
-                    c = j;
-                }
-            }
+            int c = centroidOrds.get(assignments[i]);
             centroidVectorCount[c]++;
             // if soar assignments are present, count them as well
             if (soarAssignments.length > i && soarAssignments[i] != -1) {
-                int s = -1;
-                for (int j = 0; j < centroidOrds.length; j++) {
-                    if (soarAssignments[i] == centroidOrds[j]) {
-                        s = j;
-                    }
-                }
+                int s = centroidOrds.get(soarAssignments[i]);
                 centroidVectorCount[s]++;
             }
         }
@@ -338,21 +326,11 @@ static int[][] mapAssignmentsByCluster(int centroidCount, int[] assignments, int
         Arrays.fill(centroidVectorCount, 0);
 
         for (int i = 0; i < assignments.length; i++) {
-            int c = -1;
-            for (int j = 0; j < centroidOrds.length; j++) {
-                if (assignments[i] == centroidOrds[j]) {
-                    c = j;
-                }
-            }
+            int c = centroidOrds.get(assignments[i]);
             assignmentsByCluster[c][centroidVectorCount[c]++] = i;
             // if soar assignments are present, add them to the cluster as well
             if (soarAssignments.length > i) {
-                int s = -1;
-                for (int j = 0; j < centroidOrds.length; j++) {
-                    if (soarAssignments[i] == centroidOrds[j]) {
-                        s = j;
-                    }
-                }
+                int s = centroidOrds.getOrDefault(soarAssignments[i], -1);
                 if (s != -1) {
                     assignmentsByCluster[s][centroidVectorCount[s]++] = i;
                 }
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java
@@ -221,10 +221,6 @@ public final ByteVectorValues getByteVectorValues(String field) throws IOExcepti
         return rawVectorsReader.getByteVectorValues(field);
     }
 
-    // FIXME: remove the diagnostics
-    int centroidsRead = 0;
-    int ii = 0;
-
     @Override
     public final void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
         final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field);
@@ -296,8 +292,7 @@ public final void search(String field, float[] target, KnnCollector knnCollector
 
         while (parentCentroidQueue.size() > 0 && (centroidsVisited < nProbe || knnCollectorImpl.numCollected() < knnCollector.k())) {
             NeighborQueue centroidQueue = new NeighborQueue(centroidQueryScorer.size(), true);
-            centroidsRead++;
-            centroidsRead += updateCentroidQueueWNextParent(
+            updateCentroidQueueWNextParent(
                 parentCentroidQueryScorer,
                 parentCentroidQueue,
                 centroidQueryScorer,
@@ -317,9 +312,12 @@ public final void search(String field, float[] target, KnnCollector knnCollector
                 ++centroidsVisited;
                 float centroidScore = centroidQueue.topScore();
                 // the next parent likely contains centroids we need to evaluate prior to evaluating this next centroid
-                while (parentCentroidQueue.size() > 0 && centroidScore < nextParentScore) {
-                    centroidsRead++;
-                    centroidsRead += updateCentroidQueueWNextParent(
+                // TODO: for each parent centroid I could store furtherest centroid distance from that parent and then the comparison here
+                // ... 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) ) {
+                    updateCentroidQueueWNextParent(
                         parentCentroidQueryScorer,
                         parentCentroidQueue,
                         centroidQueryScorer,
@@ -350,11 +348,6 @@ public final void search(String field, float[] target, KnnCollector knnCollector
                 }
             }
         }
-
-        if (ii == 1999) {
-            System.out.println("total centroids (parent & child) read:" + (centroidsRead / (ii + 1)));
-        }
-        ii++;
     }
 
     private static int updateCentroidQueueWNextParent(
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriterTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriterTests.java
@@ -9,47 +9,39 @@
 
 package org.elasticsearch.index.codec.vectors;
 
+import com.carrotsearch.hppc.IntIntHashMap;
+import com.carrotsearch.hppc.IntIntMap;
 import org.elasticsearch.test.ESTestCase;
 
-import java.util.Arrays;
-
 import static org.elasticsearch.index.codec.vectors.DefaultIVFVectorsWriter.mapAssignmentsByCluster;
 
 public class DefaultIVFVectorsWriterTests extends ESTestCase {
 
-    // FIXME: clean this up
     public void testAssignmentsByCluster() {
 
         // the assignments represent where vectors (by index) got assigned
         int[] assignments = new int[] { 3, 3, 2, 1, 1, 1, 1, 0, 2, 3, 4, 4, 2, 2, 2 };
         int[] soarAssignments = new int[] { 0, 0, 0, 1, 2, 3, 2, 1, 1, 1, 4, 4, 1, 2, 3 };
 
-        assert assignments.length == soarAssignments.length;
+        assertEquals(assignments.length, soarAssignments.length);
 
         // 0, 1, 2, 3, 4
         // these subsequently get sorted in the order in which the centroids would have been sorted in logic
         int[] centroidOrds = new int[] { 3, 2, 4, 0, 1 };
 
-        int[][] assignmentsByCluster = mapAssignmentsByCluster(centroidOrds.length, assignments, soarAssignments, centroidOrds);
-
-        /*
-         * correct answer
-         * [0, 1, 9, 5, 14]
-         * [2, 8, 12, 13, 14, 4, 6, 13]
-         * [10, 11, 10, 11]
-         * [7, 0, 1, 2]
-         * [3, 4, 5, 6, 3, 7, 8, 9, 12]
-         */
-        for (int i = 0; i < assignmentsByCluster.length; i++) {
-            System.out.println(Arrays.toString(assignmentsByCluster[i]));
+        IntIntMap centroidOrdsToIdx = new IntIntHashMap(centroidOrds.length);
+        for(int i = 0; i < centroidOrds.length; i++) {
+            // idx 0 1 2 3 4 5
+            // ord 3 2 0 4 1 5
+            centroidOrdsToIdx.put(centroidOrds[i], i);
         }
 
-        /*
-         * [0, 1, 5, 9, 14]
-         * [2, 4, 6, 8, 12, 13, 13, 14]
-         * [10, 10, 11, 11]
-         * [0, 1, 2, 7]
-         * [3, 3, 4, 5, 6, 7, 8, 9, 12]
-         */
+        int[][] assignmentsByCluster = mapAssignmentsByCluster(centroidOrds.length, assignments, soarAssignments, centroidOrdsToIdx);
+
+        assertArrayEquals(assignmentsByCluster[0], new int[] {0, 1, 5, 9, 14});
+        assertArrayEquals(assignmentsByCluster[1], new int[] {2, 4, 6, 8, 12, 13, 13, 14});
+        assertArrayEquals(assignmentsByCluster[2], new int[] {10, 10, 11, 11});
+        assertArrayEquals(assignmentsByCluster[3], new int[] {0, 1, 2, 7});
+        assertArrayEquals(assignmentsByCluster[4], new int[] {3, 3, 4, 5, 6, 7, 8, 9, 12});
     }
 }
