iter

Author: john-wagster

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

@@ -13,70 +13,70 @@
 
 class AssignmentArraySorter extends TimSorter {
 
-  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 int compare(int i, int j) {
-    return Integer.compare(sortOrdering[i], sortOrdering[j]);
-  }
-
-  @Override
-  protected void swap(int i, int j) {
-    final float[] tmpC = centroids[i];
-    centroids[i] = centroids[j];
-    centroids[j] = tmpC;
-
-    final int tmpA = centroidsOrds[i];
-    centroidsOrds[i] = centroidsOrds[j];
-    centroidsOrds[j] = tmpA;
-
-    final int tmpSort = sortOrdering[i];
-    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]);
-  }
+    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 int compare(int i, int j) {
+        return Integer.compare(sortOrdering[i], sortOrdering[j]);
+    }
+
+    @Override
+    protected void swap(int i, int j) {
+        final float[] tmpC = centroids[i];
+        centroids[i] = centroids[j];
+        centroids[j] = tmpC;
+
+        final int tmpA = centroidsOrds[i];
+        centroidsOrds[i] = centroidsOrds[j];
+        centroidsOrds[j] = tmpA;
+
+        final int tmpSort = sortOrdering[i];
+        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]);
+    }
 }

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

@@ -47,9 +47,13 @@ public DefaultIVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVect
     }
 
     @Override
-    CentroidQueryScorer getCentroidScorer(FieldInfo fieldInfo, int numParentCentroids,
-                                          int numCentroids, IndexInput centroids, float[] targetQuery)
-        throws IOException {
+    CentroidQueryScorer getCentroidScorer(
+        FieldInfo fieldInfo,
+        int numParentCentroids,
+        int numCentroids,
+        IndexInput centroids,
+        float[] targetQuery
+    ) throws IOException {
         final FieldEntry fieldEntry = fields.get(fieldInfo.number);
         final float globalCentroidDp = fieldEntry.globalCentroidDp();
         final OptimizedScalarQuantizer scalarQuantizer = new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
@@ -153,12 +157,8 @@ private float int4QuantizedScore(
 
     // FIXME: clean up duplicative code between the scorers
     @Override
-    CentroidQueryScorerWChildren getCentroidScorerWChildren(
-        FieldInfo fieldInfo,
-        int numCentroids,
-        IndexInput centroids,
-        float[] targetQuery
-    ) throws IOException {
+    ParentCentroidQueryScorer getParentCentroidScorer(FieldInfo fieldInfo, int numCentroids, IndexInput centroids, float[] targetQuery)
+        throws IOException {
         FieldEntry fieldEntry = fields.get(fieldInfo.number);
         float[] globalCentroid = fieldEntry.globalCentroid();
         float globalCentroidDp = fieldEntry.globalCentroidDp();
@@ -172,7 +172,7 @@ CentroidQueryScorerWChildren getCentroidScorerWChildren(
             globalCentroid
         );
         final ES91Int4VectorsScorer scorer = ESVectorUtil.getES91Int4VectorsScorer(centroids, fieldInfo.getVectorDimension());
-        return new CentroidQueryScorerWChildren() {
+        return new ParentCentroidQueryScorer() {
             int currentCentroid = -1;
             private final float[] centroidCorrectiveValues = new float[3];
             private final long quantizedVectorByteSize = fieldInfo.getVectorDimension() + 3 * Float.BYTES + Short.BYTES;

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

@@ -28,12 +28,12 @@
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 
 import static org.apache.lucene.codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat.INDEX_BITS;
 import static org.elasticsearch.index.codec.vectors.BQVectorUtils.discretize;
 import static org.elasticsearch.index.codec.vectors.BQVectorUtils.packAsBinary;
+import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.DEFAULT_VECTORS_PER_CLUSTER;
 
 /**
  * Default implementation of {@link IVFVectorsWriter}. It uses {@link HierarchicalKMeans} algorithm to
@@ -167,14 +167,23 @@ private void writePostingList(IntArrayList cluster, IndexOutput postingsOutput,
     }
 
     @Override
-    CentroidSupplier createCentroidSupplier(IndexInput centroidsInput, int numParentCentroids,
-                                            int numCentroids, FieldInfo fieldInfo, float[] globalCentroid) {
+    CentroidSupplier createCentroidSupplier(
+        IndexInput centroidsInput,
+        int numParentCentroids,
+        int numCentroids,
+        FieldInfo fieldInfo,
+        float[] globalCentroid
+    ) {
         return new OffHeapCentroidSupplier(centroidsInput, numParentCentroids, numCentroids, fieldInfo);
     }
 
-    static void writeCentroidsAndPartitions(List<CentroidPartition> centroidPartitions, float[][] centroids,
-                                     FieldInfo fieldInfo, float[] globalCentroid, IndexOutput centroidOutput)
-        throws IOException {
+    static void writeCentroidsAndPartitions(
+        List<CentroidPartition> centroidPartitions,
+        float[][] centroids,
+        FieldInfo fieldInfo,
+        float[] globalCentroid,
+        IndexOutput centroidOutput
+    ) throws IOException {
         final OptimizedScalarQuantizer osq = new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
         byte[] quantizedScratch = new byte[fieldInfo.getVectorDimension()];
         float[] centroidScratch = new float[fieldInfo.getVectorDimension()];
@@ -208,7 +217,7 @@ static void writeCentroidsAndPartitions(List<CentroidPartition> centroidPartitio
             writeQuantizedValue(centroidOutput, quantizedScratch, result);
         }
 
-        //write the raw float vectors so we can quantize the query vector relative to the centroid on read
+        // write the raw float vectors so we can quantize the query vector relative to the centroid on read
         final ByteBuffer buffer = ByteBuffer.allocate(fieldInfo.getVectorDimension() * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN);
         for (float[] centroid : centroids) {
             buffer.asFloatBuffer().put(centroid);
@@ -281,41 +290,44 @@ CentroidAssignments calculateAndWriteCentroids(
         // TODO: sort while constructing the hkmeans structure
         // we do this so we don't have to sort the assignments which is much more expensive
         int[] centroidOrds = new int[centroids.length];
-        for(int i = 0; i < centroidOrds.length; i++) {
+        for (int i = 0; i < centroidOrds.length; i++) {
             centroidOrds[i] = i;
         }
 
-        // 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);
-
         List<CentroidPartition> centroidPartitions = new ArrayList<>();
-        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 totalCentroids = 0;
-                float[] parentPartitionCentroid = new float[fieldInfo.getVectorDimension()];
-                int j = i;
-                for (; j < kMeansResult.parentLayer().length; j++) {
-                    if(kMeansResult.parentLayer()[j] != label) {
-                        break;
+
+        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 totalCentroids = 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];
+                        }
+                        totalCentroids++;
                     }
-                    for (int k = 0; k < parentPartitionCentroid.length; k++) {
-                        parentPartitionCentroid[k] += centroids[i][k];
+                    int childOrdinal = i;
+                    i = j;
+                    for (int d = 0; d < parentPartitionCentroid.length; d++) {
+                        parentPartitionCentroid[d] /= totalCentroids;
                     }
-                    totalCentroids++;
-                }
-                int childOrdinal = i;
-                i = j;
-                for (int d = 0; d < parentPartitionCentroid.length; d++) {
-                    parentPartitionCentroid[d] /= totalCentroids;
+                    centroidPartitions.add(new CentroidPartition(parentPartitionCentroid, childOrdinal, totalCentroids));
                 }
-                centroidPartitions.add(new CentroidPartition(parentPartitionCentroid, childOrdinal, totalCentroids));
             }
         }
 
@@ -330,7 +342,7 @@ CentroidAssignments calculateAndWriteCentroids(
         for (int c = 0; c < assignmentsByCluster.length; c++) {
             IntArrayList cluster = new IntArrayList(vectorPerCluster);
             for (int j = 0; j < assignments.length; j++) {
-                if(assignments[j] == -1) {
+                if (assignments[j] == -1) {
                     continue;
                 }
                 if (assignments[j] == centroidOrds[c]) {
@@ -339,7 +351,7 @@ CentroidAssignments calculateAndWriteCentroids(
             }
 
             for (int j = 0; j < soarAssignments.length; j++) {
-                if(soarAssignments[j] == -1) {
+                if (soarAssignments[j] == -1) {
                     continue;
                 }
                 if (soarAssignments[j] == centroidOrds[c]) {