Turn the generic format into a proper format

Author: thecoop

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

@@ -464,6 +464,7 @@
             org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat,
             org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat,
+            org.elasticsearch.index.codec.vectors.es93.ES93GenericFlatVectorsFormat,
             org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat;
 

server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93BinaryQuantizedVectorsFormat.java

@@ -25,6 +25,7 @@
 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;
@@ -85,20 +86,23 @@
   *  <li>The sparse vector information, if required, mapping vector ordinal to doc ID
   * </ul>
  */
-public class ES93BinaryQuantizedVectorsFormat extends ES93GenericFlatVectorsFormat {
+public class ES93BinaryQuantizedVectorsFormat extends AbstractFlatVectorsFormat {
 
     public static final String NAME = "ES93BinaryQuantizedVectorsFormat";
 
     private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
+    private final ES93GenericFlatVectorsFormat rawFormat;
+
     public ES93BinaryQuantizedVectorsFormat() {
         this(false, false);
     }
 
     public ES93BinaryQuantizedVectorsFormat(boolean useBFloat16, boolean useDirectIO) {
-        super(NAME, useBFloat16, useDirectIO);
+        super(NAME);
+        rawFormat = new ES93GenericFlatVectorsFormat(useBFloat16, useDirectIO);
     }
 
     @Override
@@ -108,11 +112,16 @@ protected FlatVectorsScorer flatVectorsScorer() {
 
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-        return new ES818BinaryQuantizedVectorsWriter(scorer, super.fieldsWriter(state), state);
+        return new ES818BinaryQuantizedVectorsWriter(scorer, rawFormat.fieldsWriter(state), state);
     }
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new ES818BinaryQuantizedVectorsReader(state, super.fieldsReader(state), scorer);
+        return new ES818BinaryQuantizedVectorsReader(state, rawFormat.fieldsReader(state), scorer);
+    }
+
+    @Override
+    public String toString() {
+        return getName() + "(name=" + getName() + ", rawVectorFormat=" + rawFormat + ", scorer=" + scorer + ")";
     }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93GenericFlatVectorsFormat.java

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 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;
@@ -20,8 +21,9 @@
 import java.io.IOException;
 import java.util.Map;
 
-public abstract class ES93GenericFlatVectorsFormat extends AbstractFlatVectorsFormat {
+public class ES93GenericFlatVectorsFormat extends AbstractFlatVectorsFormat {
 
+    static final String NAME = "ES93GenericFlatVectorsFormat";
     static final String VECTOR_FORMAT_INFO_EXTENSION = "vfi";
     static final String META_CODEC_NAME = "ES93GenericFlatVectorsFormatMeta";
 
@@ -35,12 +37,11 @@ public abstract class ES93GenericFlatVectorsFormat extends AbstractFlatVectorsFo
         VERSION_CURRENT
     );
 
-    private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-    );
-    private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-    );
+    private static final FlatVectorsScorer scorer = FlatVectorScorerUtil.getLucene99FlatVectorsScorer();
+
+    private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(scorer);
+    // TODO: a separate scorer for bfloat16
+    private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(scorer);
 
     private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
         float32VectorFormat.getName(),
@@ -52,12 +53,21 @@ public abstract class ES93GenericFlatVectorsFormat extends AbstractFlatVectorsFo
     private final DirectIOCapableFlatVectorsFormat writeFormat;
     private final boolean useDirectIO;
 
-    public ES93GenericFlatVectorsFormat(String name, boolean useBFloat16, boolean useDirectIO) {
-        super(name);
+    public ES93GenericFlatVectorsFormat() {
+        this(false, false);
+    }
+
+    public ES93GenericFlatVectorsFormat(boolean useBFloat16, boolean useDirectIO) {
+        super(NAME);
         writeFormat = useBFloat16 ? bfloat16VectorFormat : float32VectorFormat;
         this.useDirectIO = useDirectIO;
     }
 
+    @Override
+    protected FlatVectorsScorer flatVectorsScorer() {
+        return scorer;
+    }
+
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new ES93GenericFlatVectorsWriter(META, writeFormat.getName(), useDirectIO, state, writeFormat.fieldsWriter(state));
@@ -74,6 +84,6 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException
 
     @Override
     public String toString() {
-        return getName() + "(name=" + getName() + ", writeFlatVectorFormat=" + writeFormat + ")";
+        return getName() + "(name=" + getName() + ", format=" + writeFormat + ")";
     }
 }

server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat

@@ -9,5 +9,6 @@ org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat
 org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat
 org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat
 org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat
+org.elasticsearch.index.codec.vectors.es93.ES93GenericFlatVectorsFormat
 org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat
 org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat

server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93BinaryQuantizedVectorsFormatTests.java

@@ -62,12 +62,9 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Locale;
 
-import static java.lang.String.format;
 import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
-import static org.hamcrest.Matchers.either;
-import static org.hamcrest.Matchers.startsWith;
+import static org.hamcrest.Matchers.oneOf;
 
 public class ES93BinaryQuantizedVectorsFormatTests extends BaseKnnVectorsFormatTestCase {
 
@@ -196,11 +193,12 @@ public KnnVectorsFormat knnVectorsFormat() {
             }
         };
         String expectedPattern = "ES93BinaryQuantizedVectorsFormat(name=ES93BinaryQuantizedVectorsFormat,"
-            + " writeFlatVectorFormat=Lucene99FlatVectorsFormat(name=Lucene99FlatVectorsFormat,"
-            + " flatVectorScorer=%s())";
-        var defaultScorer = format(Locale.ROOT, expectedPattern, "DefaultFlatVectorScorer");
-        var memSegScorer = format(Locale.ROOT, expectedPattern, "Lucene99MemorySegmentFlatVectorsScorer");
-        assertThat(customCodec.knnVectorsFormat().toString(), either(startsWith(defaultScorer)).or(startsWith(memSegScorer)));
+            + " rawVectorFormat=ES93GenericFlatVectorsFormat(name=ES93GenericFlatVectorsFormat,"
+            + " format=Lucene99FlatVectorsFormat(name=Lucene99FlatVectorsFormat, flatVectorScorer={}())),"
+            + " scorer=ES818BinaryFlatVectorsScorer(nonQuantizedDelegate={}()))";
+        var defaultScorer = expectedPattern.replaceAll("\\{}", "DefaultFlatVectorScorer");
+        var memSegScorer = expectedPattern.replaceAll("\\{}", "Lucene99MemorySegmentFlatVectorsScorer");
+        assertThat(customCodec.knnVectorsFormat().toString(), oneOf(defaultScorer, memSegScorer));
     }
 
     @Override

server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswBinaryQuantizedVectorsFormatTests.java

@@ -50,15 +50,13 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Locale;
 
 import static java.lang.String.format;
 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.index.VectorSimilarityFunction.DOT_PRODUCT;
 import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
-import static org.hamcrest.Matchers.either;
-import static org.hamcrest.Matchers.startsWith;
+import static org.hamcrest.Matchers.oneOf;
 
 public class ES93HnswBinaryQuantizedVectorsFormatTests extends BaseKnnVectorsFormatTestCase {
 
@@ -91,14 +89,16 @@ public KnnVectorsFormat knnVectorsFormat() {
                 return new ES93HnswBinaryQuantizedVectorsFormat(10, 20, false, false, 1, null);
             }
         };
-        String expectedPattern = "ES93HnswBinaryQuantizedVectorsFormat(name=ES93HnswBinaryQuantizedVectorsFormat, maxConn=10, beamWidth=20,"
+        String expectedPattern = "ES93HnswBinaryQuantizedVectorsFormat(name=ES93HnswBinaryQuantizedVectorsFormat,"
+            + " maxConn=10, beamWidth=20,"
             + " flatVectorFormat=ES93BinaryQuantizedVectorsFormat(name=ES93BinaryQuantizedVectorsFormat,"
-            + " writeFlatVectorFormat=Lucene99FlatVectorsFormat(name=Lucene99FlatVectorsFormat,"
-            + " flatVectorScorer=%s())";
+            + " rawVectorFormat=ES93GenericFlatVectorsFormat(name=ES93GenericFlatVectorsFormat,"
+            + " format=Lucene99FlatVectorsFormat(name=Lucene99FlatVectorsFormat, flatVectorScorer={}())),"
+            + " scorer=ES818BinaryFlatVectorsScorer(nonQuantizedDelegate={}())))";
 
-        var defaultScorer = format(Locale.ROOT, expectedPattern, "DefaultFlatVectorScorer");
-        var memSegScorer = format(Locale.ROOT, expectedPattern, "Lucene99MemorySegmentFlatVectorsScorer");
-        assertThat(customCodec.knnVectorsFormat().toString(), either(startsWith(defaultScorer)).or(startsWith(memSegScorer)));
+        var defaultScorer = expectedPattern.replaceAll("\\{}", "DefaultFlatVectorScorer");
+        var memSegScorer = expectedPattern.replaceAll("\\{}", "Lucene99MemorySegmentFlatVectorsScorer");
+        assertThat(customCodec.knnVectorsFormat().toString(), oneOf(defaultScorer, memSegScorer));
     }
 
     public void testSingleVectorCase() throws Exception {