Avoid using direct I/O during vector merges (#127406)

The vector file isn't reopened during merge operations, so relying on `IOContext` to disable direct I/O during merges is ineffective.
This change updates the strategy to explicitly open the vector file twice when direct I/O is enabled: once for default reads and once for direct I/O.
We then switch to the appropriate index input in `getMergeInstance`, following the same approach used by other formats.

---------

Co-authored-by: Simon Cooper <[email protected]>

Author: jimczi

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,16 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new DirectIOLucene99FlatVectorsReader(state, vectorsScorer);
+        if (DirectIOLucene99FlatVectorsReader.shouldUseDirectIO(state)) {
+            // Use mmap for merges and direct I/O for searches.
+            // TODO: Open the mmap file with sequential access instead of random (current behavior).
+            return new MergeReaderWrapper(
+                new DirectIOLucene99FlatVectorsReader(state, vectorsScorer),
+                new Lucene99FlatVectorsReader(state, vectorsScorer)
+            );
+        } else {
+            return new Lucene99FlatVectorsReader(state, vectorsScorer);
+        }
     }
 
     @Override

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

@@ -85,6 +85,10 @@ public DirectIOLucene99FlatVectorsReader(SegmentReadState state, FlatVectorsScor
         }
     }
 
+    public static boolean shouldUseDirectIO(SegmentReadState state) {
+        return USE_DIRECT_IO && state.directory instanceof DirectIOIndexInputSupplier;
+    }
+
     private int readMetadata(SegmentReadState state) throws IOException {
         String metaFileName = IndexFileNames.segmentFileName(
             state.segmentInfo.name,
@@ -122,11 +126,9 @@ private static IndexInput openDataInput(
     ) throws IOException {
         String fileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, fileExtension);
         // use direct IO for accessing raw vector data for searches
-        IndexInput in = USE_DIRECT_IO
-            && context.context() == IOContext.Context.DEFAULT
-            && state.directory instanceof DirectIOIndexInputSupplier did
-                ? did.openInputDirect(fileName, context)
-                : state.directory.openInput(fileName, context);
+        IndexInput in = USE_DIRECT_IO && state.directory instanceof DirectIOIndexInputSupplier did
+            ? did.openInputDirect(fileName, context)
+            : state.directory.openInput(fileName, context);
         boolean success = false;
         try {
             int versionVectorData = CodecUtil.checkIndexHeader(

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

@@ -0,0 +1,89 @@
+/*
+ * 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.search.KnnCollector;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.elasticsearch.core.IOUtils;
+
+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 search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
+        mainReader.search(field, target, knnCollector, acceptDocs);
+    }
+
+    @Override
+    public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
+        mainReader.search(field, target, knnCollector, acceptDocs);
+    }
+
+    @Override
+    public void close() throws IOException {
+        IOUtils.close(mainReader, mergeReader);
+    }
+}