Revert "Avoid using direct I/O during vector merges"

This reverts commit e842eeb.

Author: jimczi

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

@@ -45,6 +45,7 @@
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
@@ -57,69 +58,31 @@ public class DirectIOLucene99FlatVectorsReader extends FlatVectorsReader {
 
     private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(DirectIOLucene99FlatVectorsReader.class);
 
-    private final IntObjectHashMap<FieldEntry> fields;
+    private final IntObjectHashMap<FieldEntry> fields = new IntObjectHashMap<>();
     private final IndexInput vectorData;
-    private final IndexInput vectorDataDirect;
-    private final IndexInput vectorDataMerge;
     private final FieldInfos fieldInfos;
-    private final boolean isClone;
 
     public DirectIOLucene99FlatVectorsReader(SegmentReadState state, FlatVectorsScorer scorer) throws IOException {
         super(scorer);
-        this.fields = new IntObjectHashMap<>();
         int versionMeta = readMetadata(state);
         this.fieldInfos = state.fieldInfos;
         boolean success = false;
         try {
-            vectorDataDirect = openDataInput(
+            vectorData = openDataInput(
                 state,
                 versionMeta,
                 DirectIOLucene99FlatVectorsFormat.VECTOR_DATA_EXTENSION,
                 DirectIOLucene99FlatVectorsFormat.VECTOR_DATA_CODEC_NAME,
-                state.context,
-                true
+                // Flat formats are used to randomly access vectors from their node ID that is stored
+                // in the HNSW graph.
+                state.context.withReadAdvice(ReadAdvice.RANDOM)
             );
             success = true;
         } finally {
             if (success == false) {
                 IOUtils.closeWhileHandlingException(this);
             }
         }
-
-        success = false;
-        try {
-            vectorDataMerge = openDataInput(
-                state,
-                versionMeta,
-                DirectIOLucene99FlatVectorsFormat.VECTOR_DATA_EXTENSION,
-                DirectIOLucene99FlatVectorsFormat.VECTOR_DATA_CODEC_NAME,
-                // We use sequential access since this input is only used for merging
-                USE_DIRECT_IO ? state.context.withReadAdvice(ReadAdvice.SEQUENTIAL) : state.context,
-                false
-            );
-            success = true;
-        } finally {
-            if (success == false) {
-                IOUtils.closeWhileHandlingException(this);
-            }
-        }
-        this.vectorData = vectorDataDirect;
-        this.isClone = false;
-    }
-
-    /**
-     * Returns a {@link DirectIOLucene99FlatVectorsReader} that switch the raw vector data to use
-     * the provided {@link IndexInput}.
-     * This is useful for merges since we want to switch from directIO to sequential reads.
-     */
-    private DirectIOLucene99FlatVectorsReader(DirectIOLucene99FlatVectorsReader clone, IndexInput vectorData) {
-        super(clone.vectorScorer);
-        this.fields = clone.fields;
-        this.vectorData = vectorData;
-        this.vectorDataDirect = clone.vectorDataDirect;
-        this.vectorDataMerge = clone.vectorDataMerge;
-        this.fieldInfos = clone.fieldInfos;
-        this.isClone = true;
     }
 
     private int readMetadata(SegmentReadState state) throws IOException {
@@ -155,13 +118,11 @@ private static IndexInput openDataInput(
         int versionMeta,
         String fileExtension,
         String codecName,
-        IOContext context,
-        boolean useDirectIO
+        IOContext context
     ) throws IOException {
         String fileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, fileExtension);
         // use direct IO for accessing raw vector data for searches
-        IndexInput in = useDirectIO
-            && USE_DIRECT_IO
+        IndexInput in = USE_DIRECT_IO
             && context.context() == IOContext.Context.DEFAULT
             && state.directory instanceof DirectIOIndexInputSupplier did
                 ? did.openInputDirect(fileName, context)
@@ -192,11 +153,6 @@ private static IndexInput openDataInput(
         }
     }
 
-    // only for tests
-    IndexInput getRawVectorsInput() {
-        return vectorData;
-    }
-
     private void readFields(ChecksumIndexInput meta, FieldInfos infos) throws IOException {
         for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) {
             FieldInfo info = infos.fieldInfo(fieldNumber);
@@ -220,8 +176,13 @@ public void checkIntegrity() throws IOException {
 
     @Override
     public FlatVectorsReader getMergeInstance() {
-        // for merges we use sequential access instead of direct IO
-        return new DirectIOLucene99FlatVectorsReader(this, vectorDataMerge);
+        try {
+            // Update the read advice since vectors are guaranteed to be accessed sequentially for merge
+            this.vectorData.updateReadAdvice(ReadAdvice.SEQUENTIAL);
+            return this;
+        } catch (IOException exception) {
+            throw new UncheckedIOException(exception);
+        }
     }
 
     private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) {
@@ -306,11 +267,16 @@ public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) thr
         );
     }
 
+    @Override
+    public void finishMerge() throws IOException {
+        // This makes sure that the access pattern hint is reverted back since HNSW implementation
+        // needs it
+        this.vectorData.updateReadAdvice(ReadAdvice.RANDOM);
+    }
+
     @Override
     public void close() throws IOException {
-        if (isClone == false) {
-            IOUtils.close(vectorDataMerge, vectorDataDirect);
-        }
+        IOUtils.close(vectorData);
     }
 
     private record FieldEntry(

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

@@ -119,11 +119,6 @@ public class ES818BinaryQuantizedVectorsReader extends FlatVectorsReader impleme
         }
     }
 
-    // only for tests
-    FlatVectorsReader getRawVectorsReader() {
-        return rawVectorsReader;
-    }
-
     private void readFields(ChecksumIndexInput meta, FieldInfos infos) throws IOException {
         for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) {
             FieldInfo info = infos.fieldInfo(fieldNumber);