Select a different reader to use for merges

Author: thecoop

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java

@@ -23,6 +23,7 @@
 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.codecs.lucene99.Lucene99FlatVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
@@ -62,7 +63,9 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new DirectIOLucene99FlatVectorsReader(state, vectorsScorer);
+        return new MergeReaderWrapper(
+            new DirectIOLucene99FlatVectorsReader(state, vectorsScorer),
+            new Lucene99FlatVectorsReader(state, vectorsScorer));
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/MergeReaderWrapper.java

@@ -0,0 +1,77 @@
+/*
+ * 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.es818;
+
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+
+import java.io.IOException;
+import java.util.Collection;
+
+class MergeReaderWrapper extends FlatVectorsReader {
+
+    private final FlatVectorsReader mainReader;
+    private final FlatVectorsReader mergeReader;
+
+    protected MergeReaderWrapper(FlatVectorsReader mainReader, FlatVectorsReader mergeReader) {
+        super(mainReader.getFlatVectorScorer());
+        this.mainReader = mainReader;
+        this.mergeReader = mergeReader;
+    }
+
+    @Override
+    public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException {
+        return mainReader.getRandomVectorScorer(field, target);
+    }
+
+    @Override
+    public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException {
+        return mainReader.getRandomVectorScorer(field, target);
+    }
+
+    @Override
+    public void checkIntegrity() throws IOException {
+        mainReader.checkIntegrity();
+    }
+
+    @Override
+    public FloatVectorValues getFloatVectorValues(String field) throws IOException {
+        return mainReader.getFloatVectorValues(field);
+    }
+
+    @Override
+    public ByteVectorValues getByteVectorValues(String field) throws IOException {
+        return mainReader.getByteVectorValues(field);
+    }
+
+    @Override
+    public FlatVectorsReader getMergeInstance() {
+        return mergeReader;
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        return mainReader.ramBytesUsed();
+    }
+
+    @Override
+    public Collection<Accountable> getChildResources() {
+        return mainReader.getChildResources();
+    }
+
+    @Override
+    public void close() throws IOException {
+        IOUtils.close(mainReader, mergeReader);
+    }
+}