adapt code and comment

Author: jimczi

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

@@ -63,9 +63,16 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new MergeReaderWrapper(
-            new DirectIOLucene99FlatVectorsReader(state, vectorsScorer),
-            new Lucene99FlatVectorsReader(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

@@ -12,9 +12,11 @@
 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.IOUtils;
+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;
@@ -70,6 +72,16 @@ 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);