Refactor/disk bbq class movements (#135367)

This is a minor refactor. But it makes some the interfaces cleaner

Author: benwtrent

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/CentroidSupplier.java

@@ -0,0 +1,50 @@
+/*
+ * 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.diskbbq;
+
+import java.io.IOException;
+
+/**
+ * An interface for that supply centroids.
+ */
+interface CentroidSupplier {
+    CentroidSupplier EMPTY = new CentroidSupplier() {
+        @Override
+        public int size() {
+            return 0;
+        }
+
+        @Override
+        public float[] centroid(int centroidOrdinal) {
+            throw new IllegalStateException("No centroids");
+        }
+    };
+
+    int size();
+
+    float[] centroid(int centroidOrdinal) throws IOException;
+
+    static CentroidSupplier fromArray(float[][] centroids) {
+        if (centroids.length == 0) {
+            return EMPTY;
+        }
+        return new CentroidSupplier() {
+            @Override
+            public int size() {
+                return centroids.length;
+            }
+
+            @Override
+            public float[] centroid(int centroidOrdinal) {
+                return centroids[centroidOrdinal];
+            }
+        };
+    }
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/DiskBBQBulkWriter.java

@@ -29,8 +29,7 @@ protected DiskBBQBulkWriter(int bulkSize, IndexOutput out) {
         this.out = out;
     }
 
-    abstract void writeVectors(ES920DiskBBQVectorsWriter.QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter)
-        throws IOException;
+    abstract void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException;
 
     static class OneBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
         private final OptimizedScalarQuantizer.QuantizationResult[] corrections;
@@ -41,8 +40,7 @@ static class OneBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
         }
 
         @Override
-        void writeVectors(ES920DiskBBQVectorsWriter.QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter)
-            throws IOException {
+        void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException {
             int limit = qvv.count() - bulkSize + 1;
             int i = 0;
             for (; i < limit; i += bulkSize) {
@@ -104,8 +102,7 @@ static class SevenBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
         }
 
         @Override
-        void writeVectors(ES920DiskBBQVectorsWriter.QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter)
-            throws IOException {
+        void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException {
             int limit = qvv.count() - bulkSize + 1;
             int i = 0;
             for (; i < limit; i += bulkSize) {

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java

@@ -17,7 +17,6 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
-import org.apache.lucene.util.IntroSorter;
 import org.apache.lucene.util.VectorUtil;
 import org.apache.lucene.util.hnsw.IntToIntFunction;
 import org.apache.lucene.util.packed.PackedInts;
@@ -384,7 +383,7 @@ private void writeCentroidsWithParents(
         centroidOutput.writeVInt(centroidGroups.centroids.length);
         centroidOutput.writeVInt(centroidGroups.maxVectorsPerCentroidLength);
         QuantizedCentroids parentQuantizeCentroid = new QuantizedCentroids(
-            new OnHeapCentroidSupplier(centroidGroups.centroids),
+            CentroidSupplier.fromArray(centroidGroups.centroids),
             fieldInfo.getVectorDimension(),
             osq,
             globalCentroid
@@ -577,18 +576,6 @@ public float[] centroid(int centroidOrdinal) throws IOException {
         }
     }
 
-    interface QuantizedVectorValues {
-        int count();
-
-        byte[] next() throws IOException;
-
-        OptimizedScalarQuantizer.QuantizationResult getCorrections() throws IOException;
-    }
-
-    interface IntToBooleanFunction {
-        boolean apply(int ord);
-    }
-
     static class QuantizedCentroids implements QuantizedVectorValues {
         private final CentroidSupplier supplier;
         private final OptimizedScalarQuantizer quantizer;
@@ -760,37 +747,4 @@ public void readQuantizedVector(int ord, boolean isOverspill) throws IOException
             bitSum = quantizedVectorsInput.readShort();
         }
     }
-
-    private static class IntSorter extends IntroSorter {
-        int pivot = -1;
-        private final int[] arr;
-        private final IntToIntFunction func;
-
-        private IntSorter(int[] arr, IntToIntFunction func) {
-            this.arr = arr;
-            this.func = func;
-        }
-
-        @Override
-        protected void setPivot(int i) {
-            pivot = func.apply(arr[i]);
-        }
-
-        @Override
-        protected int comparePivot(int j) {
-            return Integer.compare(pivot, func.apply(arr[j]));
-        }
-
-        @Override
-        protected int compare(int a, int b) {
-            return Integer.compare(func.apply(arr[a]), func.apply(arr[b]));
-        }
-
-        @Override
-        protected void swap(int i, int j) {
-            final int tmp = arr[i];
-            arr[i] = arr[j];
-            arr[j] = tmp;
-        }
-    }
 }

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

@@ -400,8 +400,6 @@ interface CentroidIterator {
     }
 
     interface PostingVisitor {
-        // TODO maybe we can not specifically pass the centroid...
-
         /** returns the number of documents in the posting list */
         int resetPostingsScorer(long offset) throws IOException;
 

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

@@ -183,7 +183,7 @@ public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
             // build centroids
             final CentroidAssignments centroidAssignments = calculateCentroids(fieldWriter.fieldInfo, floatVectorValues, globalCentroid);
             // wrap centroids with a supplier
-            final CentroidSupplier centroidSupplier = new OnHeapCentroidSupplier(centroidAssignments.centroids());
+            final CentroidSupplier centroidSupplier = CentroidSupplier.fromArray(centroidAssignments.centroids());
             // write posting lists
             final long postingListOffset = ivfClusters.alignFilePointer(Float.BYTES);
             final CentroidOffsetAndLength centroidOffsetAndLength = buildAndWritePostingsLists(
@@ -549,40 +549,4 @@ public final long ramBytesUsed() {
 
     private record FieldWriter(FieldInfo fieldInfo, FlatFieldVectorsWriter<float[]> delegate) {}
 
-    interface CentroidSupplier {
-        CentroidSupplier EMPTY = new CentroidSupplier() {
-            @Override
-            public int size() {
-                return 0;
-            }
-
-            @Override
-            public float[] centroid(int centroidOrdinal) {
-                throw new IllegalStateException("No centroids");
-            }
-        };
-
-        int size();
-
-        float[] centroid(int centroidOrdinal) throws IOException;
-    }
-
-    // TODO throw away rawCentroids
-    static class OnHeapCentroidSupplier implements CentroidSupplier {
-        private final float[][] centroids;
-
-        OnHeapCentroidSupplier(float[][] centroids) {
-            this.centroids = centroids;
-        }
-
-        @Override
-        public int size() {
-            return centroids.length;
-        }
-
-        @Override
-        public float[] centroid(int centroidOrdinal) throws IOException {
-            return centroids[centroidOrdinal];
-        }
-    }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IntSorter.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.diskbbq;
+
+import org.apache.lucene.util.IntroSorter;
+import org.apache.lucene.util.hnsw.IntToIntFunction;
+
+class IntSorter extends IntroSorter {
+    int pivot = -1;
+    private final int[] arr;
+    private final IntToIntFunction func;
+
+    IntSorter(int[] arr, IntToIntFunction func) {
+        this.arr = arr;
+        this.func = func;
+    }
+
+    @Override
+    protected void setPivot(int i) {
+        pivot = func.apply(arr[i]);
+    }
+
+    @Override
+    protected int comparePivot(int j) {
+        return Integer.compare(pivot, func.apply(arr[j]));
+    }
+
+    @Override
+    protected int compare(int a, int b) {
+        return Integer.compare(func.apply(arr[a]), func.apply(arr[b]));
+    }
+
+    @Override
+    protected void swap(int i, int j) {
+        final int tmp = arr[i];
+        arr[i] = arr[j];
+        arr[j] = tmp;
+    }
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IntToBooleanFunction.java

@@ -0,0 +1,18 @@
+/*
+ * 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.diskbbq;
+
+/**
+ * Functional interface representing a function that takes an integer input
+ * and produces a boolean output.
+ */
+interface IntToBooleanFunction {
+    boolean apply(int value);
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/QuantizedVectorValues.java

@@ -0,0 +1,27 @@
+/*
+ * 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.diskbbq;
+
+import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
+
+import java.io.IOException;
+
+/**
+ * Interface representing a collection of quantized vector values.
+ * Provides methods to iterate through the vectors and retrieve
+ * associated quantization correction data.
+ */
+interface QuantizedVectorValues {
+    int count();
+
+    byte[] next() throws IOException;
+
+    OptimizedScalarQuantizer.QuantizationResult getCorrections() throws IOException;
+}