IVF Hierarchical KMeans Flush & Merge (#128675)

added hierarchical kmeans as a clustering algorithm to better partitionin the space when running ivf on flush and merge

Author: john-wagster

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnSearcher.java

@@ -276,12 +276,11 @@ TopDocs doVectorQuery(byte[] vector, IndexSearcher searcher) throws IOException
     TopDocs doVectorQuery(float[] vector, IndexSearcher searcher) throws IOException {
         Query knnQuery;
         int topK = this.topK;
-        int efSearch = this.efSearch;
         if (overSamplingFactor > 1f) {
             // oversample the topK results to get more candidates for the final result
             topK = (int) Math.ceil(topK * overSamplingFactor);
-            efSearch = Math.max(topK, efSearch);
         }
+        int efSearch = Math.max(topK, this.efSearch);
         if (indexType == KnnIndexTester.IndexType.IVF) {
             knnQuery = new IVFKnnFloatVectorQuery(VECTOR_FIELD, vector, topK, efSearch, null, nProbe);
         } else {

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

@@ -0,0 +1,46 @@
+/*
+ * 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;
+
+import org.apache.lucene.internal.hppc.IntArrayList;
+
+final class CentroidAssignments {
+
+    private final int numCentroids;
+    private final float[][] cachedCentroids;
+    private final IntArrayList[] assignmentsByCluster;
+
+    private CentroidAssignments(int numCentroids, float[][] cachedCentroids, IntArrayList[] assignmentsByCluster) {
+        this.numCentroids = numCentroids;
+        this.cachedCentroids = cachedCentroids;
+        this.assignmentsByCluster = assignmentsByCluster;
+    }
+
+    CentroidAssignments(float[][] centroids, IntArrayList[] assignmentsByCluster) {
+        this(centroids.length, centroids, assignmentsByCluster);
+    }
+
+    CentroidAssignments(int numCentroids, IntArrayList[] assignmentsByCluster) {
+        this(numCentroids, null, assignmentsByCluster);
+    }
+
+    // Getters and setters
+    public int numCentroids() {
+        return numCentroids;
+    }
+
+    public float[][] cachedCentroids() {
+        return cachedCentroids;
+    }
+
+    public IntArrayList[] assignmentsByCluster() {
+        return assignmentsByCluster;
+    }
+}

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

@@ -112,15 +112,6 @@ public float score(int centroidOrdinal) throws IOException {
         };
     }
 
-    @Override
-    protected FloatVectorValues getCentroids(IndexInput indexInput, int numCentroids, FieldInfo info) {
-        FieldEntry entry = fields.get(info.number);
-        if (entry == null) {
-            return null;
-        }
-        return new OffHeapCentroidFloatVectorValues(numCentroids, indexInput, info.getVectorDimension());
-    }
-
     @Override
     NeighborQueue scorePostingLists(FieldInfo fieldInfo, KnnCollector knnCollector, CentroidQueryScorer centroidQueryScorer, int nProbe)
         throws IOException {