Use abstract classes

Author: thecoop

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/ES92BFloat16FlatVectorsFormat.java

@@ -19,7 +19,6 @@
  */
 package org.elasticsearch.index.codec.vectors.es92;
 
-import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
@@ -29,15 +28,15 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.MergeInfo;
 import org.elasticsearch.common.util.set.Sets;
+import org.elasticsearch.index.codec.vectors.AbstractFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.MergeReaderWrapper;
 import org.elasticsearch.index.codec.vectors.es818.DirectIOHint;
-import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.store.FsDirectoryFactory;
 
 import java.io.IOException;
 import java.util.Set;
 
-public final class ES92BFloat16FlatVectorsFormat extends FlatVectorsFormat {
+public final class ES92BFloat16FlatVectorsFormat extends AbstractFlatVectorsFormat {
 
     static final String NAME = "ES92BFloat16FlatVectorsFormat";
     static final String META_CODEC_NAME = "ES92BFloat16FlatVectorsFormatMeta";
@@ -56,13 +55,18 @@ public ES92BFloat16FlatVectorsFormat(FlatVectorsScorer vectorsScorer) {
         this.vectorsScorer = vectorsScorer;
     }
 
+    @Override
+    protected FlatVectorsScorer flatVectorsScorer() {
+        return vectorsScorer;
+    }
+
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new ES92BFloat16FlatVectorsWriter(state, vectorsScorer);
     }
 
     static boolean shouldUseDirectIO(SegmentReadState state) {
-        return ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO && FsDirectoryFactory.isHybridFs(state.directory);
+        return USE_DIRECT_IO && FsDirectoryFactory.isHybridFs(state.directory);
     }
 
     @Override
@@ -87,11 +91,6 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException
         }
     }
 
-    @Override
-    public String toString() {
-        return "ES92BFloat16FlatVectorsFormat(" + "vectorsScorer=" + vectorsScorer + ')';
-    }
-
     static class DirectIOContext implements IOContext {
 
         final Set<FileOpenHint> hints;

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/ES92BinaryQuantizedBFloat16VectorsFormat.java

@@ -22,9 +22,11 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.elasticsearch.index.codec.vectors.AbstractFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryFlatVectorsScorer;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsReader;
@@ -87,19 +89,10 @@
   *  <li>The sparse vector information, if required, mapping vector ordinal to doc ID
   * </ul>
  */
-public class ES92BinaryQuantizedBFloat16VectorsFormat extends FlatVectorsFormat {
+public class ES92BinaryQuantizedBFloat16VectorsFormat extends AbstractFlatVectorsFormat {
 
-    public static final String BINARIZED_VECTOR_COMPONENT = "BVEC";
     public static final String NAME = "ES92BinaryQuantizedBFloat16VectorsFormat";
 
-    static final int VERSION_START = 0;
-    static final int VERSION_CURRENT = VERSION_START;
-    static final String META_CODEC_NAME = "ES92BinaryQuantizedBFloat16VectorsFormatMeta";
-    static final String VECTOR_DATA_CODEC_NAME = "ES92BinaryQuantizedBFloat16VectorsFormatData";
-    static final String META_EXTENSION = "vemb";
-    static final String VECTOR_DATA_EXTENSION = "veb";
-    static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
-
     private static final FlatVectorsFormat rawVectorFormat = new ES92BFloat16FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
@@ -113,6 +106,11 @@ public ES92BinaryQuantizedBFloat16VectorsFormat() {
         super(NAME);
     }
 
+    @Override
+    protected FlatVectorsScorer flatVectorsScorer() {
+        return scorer;
+    }
+
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new ES818BinaryQuantizedVectorsWriter(scorer, rawVectorFormat.fieldsWriter(state), state);
@@ -122,14 +120,4 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
         return new ES818BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), scorer);
     }
-
-    @Override
-    public int getMaxDimensions(String fieldName) {
-        return MAX_DIMS_COUNT;
-    }
-
-    @Override
-    public String toString() {
-        return "ES92BinaryQuantizedBFloat16VectorsFormat(name=" + NAME + ", flatVectorScorer=" + scorer + ")";
-    }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/ES92HnswBinaryQuantizedBFloat16VectorsFormat.java

@@ -19,54 +19,32 @@
  */
 package org.elasticsearch.index.codec.vectors.es92;
 
-import org.apache.lucene.codecs.KnnVectorsFormat;
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.KnnVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
-import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
 import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
-import org.apache.lucene.search.TaskExecutor;
-import org.apache.lucene.util.hnsw.HnswGraph;
+import org.elasticsearch.index.codec.vectors.AbstractHnswVectorsFormat;
 
 import java.io.IOException;
 import java.util.concurrent.ExecutorService;
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER;
-import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.MAXIMUM_BEAM_WIDTH;
-import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.MAXIMUM_MAX_CONN;
-import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT;
 
 /**
  * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
  */
-public class ES92HnswBinaryQuantizedBFloat16VectorsFormat extends KnnVectorsFormat {
+public class ES92HnswBinaryQuantizedBFloat16VectorsFormat extends AbstractHnswVectorsFormat {
 
     public static final String NAME = "ES92HnswBinaryQuantizedBFloat16VectorsFormat";
 
-    /**
-     * Controls how many of the nearest neighbor candidates are connected to the new node. Defaults to
-     * {@link Lucene99HnswVectorsFormat#DEFAULT_MAX_CONN}. See {@link HnswGraph} for more details.
-     */
-    private final int maxConn;
-
-    /**
-     * The number of candidate neighbors to track while searching the graph for each newly inserted
-     * node. Defaults to {@link Lucene99HnswVectorsFormat#DEFAULT_BEAM_WIDTH}. See {@link HnswGraph}
-     * for details.
-     */
-    private final int beamWidth;
-
     /** The format for storing, reading, merging vectors on disk */
     private static final FlatVectorsFormat flatVectorsFormat = new ES92BinaryQuantizedBFloat16VectorsFormat();
 
-    private final int numMergeWorkers;
-    private final TaskExecutor mergeExec;
-
     /** Constructs a format using default graph construction parameters */
     public ES92HnswBinaryQuantizedBFloat16VectorsFormat() {
         this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null);
@@ -93,28 +71,12 @@ public ES92HnswBinaryQuantizedBFloat16VectorsFormat(int maxConn, int beamWidth)
      *     generated by this format to do the merge
      */
     public ES92HnswBinaryQuantizedBFloat16VectorsFormat(int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
-        super(NAME);
-        if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
-            throw new IllegalArgumentException(
-                "maxConn must be positive and less than or equal to " + MAXIMUM_MAX_CONN + "; maxConn=" + maxConn
-            );
-        }
-        if (beamWidth <= 0 || beamWidth > MAXIMUM_BEAM_WIDTH) {
-            throw new IllegalArgumentException(
-                "beamWidth must be positive and less than or equal to " + MAXIMUM_BEAM_WIDTH + "; beamWidth=" + beamWidth
-            );
-        }
-        this.maxConn = maxConn;
-        this.beamWidth = beamWidth;
-        if (numMergeWorkers == 1 && mergeExec != null) {
-            throw new IllegalArgumentException("No executor service is needed as we'll use single thread to merge");
-        }
-        this.numMergeWorkers = numMergeWorkers;
-        if (mergeExec != null) {
-            this.mergeExec = new TaskExecutor(mergeExec);
-        } else {
-            this.mergeExec = null;
-        }
+        super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec);
+    }
+
+    @Override
+    protected FlatVectorsFormat flatVectorsFormat() {
+        return flatVectorsFormat;
     }
 
     @Override
@@ -126,20 +88,4 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
     public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
         return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state));
     }
-
-    @Override
-    public int getMaxDimensions(String fieldName) {
-        return MAX_DIMS_COUNT;
-    }
-
-    @Override
-    public String toString() {
-        return "ES92HnswBinaryQuantizedBFloat16VectorsFormat(name=ES92HnswBinaryQuantizedBFloat16VectorsFormat, maxConn="
-            + maxConn
-            + ", beamWidth="
-            + beamWidth
-            + ", flatVectorFormat="
-            + flatVectorsFormat
-            + ")";
-    }
 }