Adding a new development format for diskbbq

server/src/main/java/module-info.java

@@ -7,7 +7,6 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat;
 import org.elasticsearch.plugins.internal.RestExtension;
 import org.elasticsearch.reservedstate.ReservedStateHandlerProvider;
 
@@ -463,7 +462,8 @@
             org.elasticsearch.index.codec.vectors.es816.ES816HnswBinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
-            ES920DiskBBQVectorsFormat;
+            org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat,
+            org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat;
 
     provides org.apache.lucene.codecs.Codec
         with

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

@@ -9,9 +9,9 @@
 
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
-record CentroidAssignments(int numCentroids, float[][] centroids, int[] assignments, int[] overspillAssignments) {
+public record CentroidAssignments(int numCentroids, float[][] centroids, int[] assignments, int[] overspillAssignments) {
 
-    CentroidAssignments(float[][] centroids, int[] assignments, int[] overspillAssignments) {
+    public CentroidAssignments(float[][] centroids, int[] assignments, int[] overspillAssignments) {
         this(centroids.length, centroids, assignments, overspillAssignments);
         assert assignments.length == overspillAssignments.length || overspillAssignments.length == 0
             : "assignments and overspillAssignments must have the same length";

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

@@ -14,7 +14,7 @@
 /**
  * An interface for that supply centroids.
  */
-interface CentroidSupplier {
+public interface CentroidSupplier {
     CentroidSupplier EMPTY = new CentroidSupplier() {
         @Override
         public int size() {

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

@@ -20,7 +20,7 @@
  * This class provides the structure for writing vectors in bulk, with specific
  * implementations for different bit sizes strategies.
  */
-abstract class DiskBBQBulkWriter {
+public abstract class DiskBBQBulkWriter {
     protected final int bulkSize;
     protected final IndexOutput out;
 
@@ -29,18 +29,18 @@ protected DiskBBQBulkWriter(int bulkSize, IndexOutput out) {
         this.out = out;
     }
 
-    abstract void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException;
+    public abstract void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException;
 
-    static class OneBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
+    public static class OneBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
         private final OptimizedScalarQuantizer.QuantizationResult[] corrections;
 
-        OneBitDiskBBQBulkWriter(int bulkSize, IndexOutput out) {
+        public OneBitDiskBBQBulkWriter(int bulkSize, IndexOutput out) {
             super(bulkSize, out);
             this.corrections = new OptimizedScalarQuantizer.QuantizationResult[bulkSize];
         }
 
         @Override
-        void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException {
+        public void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException {
             int limit = qvv.count() - bulkSize + 1;
             int i = 0;
             for (; i < limit; i += bulkSize) {
@@ -93,16 +93,16 @@ private void writeCorrection(OptimizedScalarQuantizer.QuantizationResult correct
         }
     }
 
-    static class SevenBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
+    public static class SevenBitDiskBBQBulkWriter extends DiskBBQBulkWriter {
         private final OptimizedScalarQuantizer.QuantizationResult[] corrections;
 
-        SevenBitDiskBBQBulkWriter(int bulkSize, IndexOutput out) {
+        public SevenBitDiskBBQBulkWriter(int bulkSize, IndexOutput out) {
             super(bulkSize, out);
             this.corrections = new OptimizedScalarQuantizer.QuantizationResult[bulkSize];
         }
 
         @Override
-        void writeVectors(QuantizedVectorValues qvv, CheckedIntConsumer<IOException> docsWriter) throws IOException {
+        public 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/DocIdsWriter.java

@@ -30,7 +30,7 @@
  *
  * <p>It is copied from the BKD implementation.
  */
-final class DocIdsWriter {
+public final class DocIdsWriter {
 
     private static final byte CONTINUOUS_IDS = (byte) -2;
     private static final byte DELTA_BPV_16 = (byte) 16;
@@ -40,7 +40,7 @@ final class DocIdsWriter {
 
     private int[] scratch = new int[0];
 
-    DocIdsWriter() {}
+    public DocIdsWriter() {}
 
     /**
      * Calculate the best encoding that will be used to write blocks of doc ids of blockSize.
@@ -51,7 +51,7 @@ final class DocIdsWriter {
      * @param blockSize the block size
      * @return the byte encoding to use for the blocks
      */
-    byte calculateBlockEncoding(IntToIntFunction docIds, int count, int blockSize) {
+    public byte calculateBlockEncoding(IntToIntFunction docIds, int count, int blockSize) {
         if (count == 0) {
             return CONTINUOUS_IDS;
         }
@@ -90,7 +90,7 @@ byte calculateBlockEncoding(IntToIntFunction docIds, int count, int blockSize) {
         }
     }
 
-    void writeDocIds(IntToIntFunction docIds, int count, byte encoding, DataOutput out) throws IOException {
+    public void writeDocIds(IntToIntFunction docIds, int count, byte encoding, DataOutput out) throws IOException {
         if (count == 0) {
             return;
         }
@@ -206,7 +206,7 @@ private static int[] sortedAndMaxAndMin2Max(IntToIntFunction docIds, int count)
         return new int[] { (strictlySorted && min2max == count) ? 1 : 0, max, min2max };
     }
 
-    void writeDocIds(IntToIntFunction docIds, int count, DataOutput out) throws IOException {
+    public void writeDocIds(IntToIntFunction docIds, int count, DataOutput out) throws IOException {
         if (count == 0) {
             return;
         }
@@ -253,7 +253,7 @@ void writeDocIds(IntToIntFunction docIds, int count, DataOutput out) throws IOEx
         }
     }
 
-    void readInts(IndexInput in, int count, byte encoding, int[] docIDs) throws IOException {
+    public void readInts(IndexInput in, int count, byte encoding, int[] docIDs) throws IOException {
         if (count == 0) {
             return;
         }
@@ -271,7 +271,7 @@ void readInts(IndexInput in, int count, byte encoding, int[] docIDs) throws IOEx
     }
 
     /** Read {@code count} integers into {@code docIDs}. */
-    void readInts(IndexInput in, int count, int[] docIDs) throws IOException {
+    public void readInts(IndexInput in, int count, int[] docIDs) throws IOException {
         if (count == 0) {
             return;
         }

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

@@ -43,7 +43,8 @@ public ES920DiskBBQVectorsReader(SegmentReadState state, Map<String, FlatVectors
         super(state, rawVectorsReader);
     }
 
-    CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice) throws IOException {
+    public CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice)
+        throws IOException {
         return new CentroidIterator() {
             CentroidOffsetAndLength nextOffsetAndLength = centroidIterator.hasNext()
                 ? centroidIterator.nextPostingListOffsetAndLength()
@@ -80,7 +81,7 @@ public CentroidOffsetAndLength nextPostingListOffsetAndLength() throws IOExcepti
     }
 
     @Override
-    CentroidIterator getCentroidIterator(
+    public CentroidIterator getCentroidIterator(
         FieldInfo fieldInfo,
         int numCentroids,
         IndexInput centroids,
@@ -348,7 +349,8 @@ private static void score(
     }
 
     @Override
-    PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput indexInput, float[] target, Bits acceptDocs) throws IOException {
+    public PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput indexInput, float[] target, Bits acceptDocs)
+        throws IOException {
         FieldEntry entry = fields.get(fieldInfo.number);
         final int maxPostingListSize = indexInput.readVInt();
         return new MemorySegmentPostingsVisitor(target, indexInput, entry, fieldInfo, maxPostingListSize, acceptDocs);

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

@@ -64,7 +64,7 @@ public ES920DiskBBQVectorsWriter(
     }
 
     @Override
-    CentroidOffsetAndLength buildAndWritePostingsLists(
+    public CentroidOffsetAndLength buildAndWritePostingsLists(
         FieldInfo fieldInfo,
         CentroidSupplier centroidSupplier,
         FloatVectorValues floatVectorValues,
@@ -160,7 +160,7 @@ CentroidOffsetAndLength buildAndWritePostingsLists(
 
     @Override
     @SuppressForbidden(reason = "require usage of Lucene's IOUtils#deleteFilesIgnoringExceptions(...)")
-    CentroidOffsetAndLength buildAndWritePostingsLists(
+    public CentroidOffsetAndLength buildAndWritePostingsLists(
         FieldInfo fieldInfo,
         CentroidSupplier centroidSupplier,
         FloatVectorValues floatVectorValues,
@@ -347,12 +347,17 @@ private static void printClusterQualityStatistics(int[][] clusters) {
     }
 
     @Override
-    CentroidSupplier createCentroidSupplier(IndexInput centroidsInput, int numCentroids, FieldInfo fieldInfo, float[] globalCentroid) {
+    public CentroidSupplier createCentroidSupplier(
+        IndexInput centroidsInput,
+        int numCentroids,
+        FieldInfo fieldInfo,
+        float[] globalCentroid
+    ) {
         return new OffHeapCentroidSupplier(centroidsInput, numCentroids, fieldInfo);
     }
 
     @Override
-    void writeCentroids(
+    public void writeCentroids(
         FieldInfo fieldInfo,
         CentroidSupplier centroidSupplier,
         float[] globalCentroid,
@@ -502,7 +507,7 @@ public int size() {
      * @throws IOException if an I/O error occurs
      */
     @Override
-    CentroidAssignments calculateCentroids(FieldInfo fieldInfo, FloatVectorValues floatVectorValues, float[] globalCentroid)
+    public CentroidAssignments calculateCentroids(FieldInfo fieldInfo, FloatVectorValues floatVectorValues, float[] globalCentroid)
         throws IOException {
 
         long nanoTime = System.nanoTime();

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

@@ -106,7 +106,7 @@ protected IVFVectorsReader(SegmentReadState state, Map<String, FlatVectorsReader
         }
     }
 
-    abstract CentroidIterator getCentroidIterator(
+    public abstract CentroidIterator getCentroidIterator(
         FieldInfo fieldInfo,
         int numCentroids,
         IndexInput centroids,
@@ -388,18 +388,18 @@ IndexInput postingListSlice(IndexInput postingListFile) throws IOException {
         }
     }
 
-    abstract PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput postingsLists, float[] target, Bits needsScoring)
+    public abstract PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput postingsLists, float[] target, Bits needsScoring)
         throws IOException;
 
-    record CentroidOffsetAndLength(long offset, long length) {}
+    public record CentroidOffsetAndLength(long offset, long length) {}
 
-    interface CentroidIterator {
+    public interface CentroidIterator {
         boolean hasNext();
 
         CentroidOffsetAndLength nextPostingListOffsetAndLength() throws IOException;
     }
 
-    interface PostingVisitor {
+    public interface PostingVisitor {
         /** 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

@@ -127,20 +127,20 @@ public final KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOExc
         return rawVectorDelegate;
     }
 
-    abstract CentroidAssignments calculateCentroids(FieldInfo fieldInfo, FloatVectorValues floatVectorValues, float[] globalCentroid)
+    public abstract CentroidAssignments calculateCentroids(FieldInfo fieldInfo, FloatVectorValues floatVectorValues, float[] globalCentroid)
         throws IOException;
 
-    record CentroidOffsetAndLength(LongValues offsets, LongValues lengths) {}
+    public record CentroidOffsetAndLength(LongValues offsets, LongValues lengths) {}
 
-    abstract void writeCentroids(
+    public abstract void writeCentroids(
         FieldInfo fieldInfo,
         CentroidSupplier centroidSupplier,
         float[] globalCentroid,
         CentroidOffsetAndLength centroidOffsetAndLength,
         IndexOutput centroidOutput
     ) throws IOException;
 
-    abstract CentroidOffsetAndLength buildAndWritePostingsLists(
+    public abstract CentroidOffsetAndLength buildAndWritePostingsLists(
         FieldInfo fieldInfo,
         CentroidSupplier centroidSupplier,
         FloatVectorValues floatVectorValues,
@@ -150,7 +150,7 @@ abstract CentroidOffsetAndLength buildAndWritePostingsLists(
         int[] overspillAssignments
     ) throws IOException;
 
-    abstract CentroidOffsetAndLength buildAndWritePostingsLists(
+    public abstract CentroidOffsetAndLength buildAndWritePostingsLists(
         FieldInfo fieldInfo,
         CentroidSupplier centroidSupplier,
         FloatVectorValues floatVectorValues,
@@ -161,7 +161,7 @@ abstract CentroidOffsetAndLength buildAndWritePostingsLists(
         int[] overspillAssignments
     ) throws IOException;
 
-    abstract CentroidSupplier createCentroidSupplier(
+    public abstract CentroidSupplier createCentroidSupplier(
         IndexInput centroidsInput,
         int numCentroids,
         FieldInfo fieldInfo,

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

@@ -12,12 +12,12 @@
 import org.apache.lucene.util.IntroSorter;
 import org.apache.lucene.util.hnsw.IntToIntFunction;
 
-class IntSorter extends IntroSorter {
+public class IntSorter extends IntroSorter {
     int pivot = -1;
     private final int[] arr;
     private final IntToIntFunction func;
 
-    IntSorter(int[] arr, IntToIntFunction func) {
+    public IntSorter(int[] arr, IntToIntFunction func) {
         this.arr = arr;
         this.func = func;
     }

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

@@ -13,6 +13,6 @@
  * Functional interface representing a function that takes an integer input
  * and produces a boolean output.
  */
-interface IntToBooleanFunction {
+public interface IntToBooleanFunction {
     boolean apply(int value);
 }

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

@@ -18,7 +18,7 @@
  * Provides methods to iterate through the vectors and retrieve
  * associated quantization correction data.
  */
-interface QuantizedVectorValues {
+public interface QuantizedVectorValues {
     int count();
 
     byte[] next() throws IOException;