Only use direct IO when opening with Default IOContext and using the hybrid directory (#129530)

This commit updates the DIO format so that it only use direct IO when opening with Default IOContext and using the hybrid directory. I also took the opportunity to better align the code with where we ultimately want to get to (what is in the lucene_snapshot branch).

Author: ChrisHegarty

muted-tests.yml

@@ -472,12 +472,6 @@ tests:
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test040JavaUsesTheOsProvidedKeystore
   issue: https://github.com/elastic/elasticsearch/issues/128230
-- class: org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormatTests
-  method: testSimpleOffHeapSizeFSDir
-  issue: https://github.com/elastic/elasticsearch/issues/128799
-- class: org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormatTests
-  method: testSimpleOffHeapSizeFSDir
-  issue: https://github.com/elastic/elasticsearch/issues/128800
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test150MachineDependentHeap
   issue: https://github.com/elastic/elasticsearch/issues/128120

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

@@ -27,6 +27,9 @@
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.store.FilterDirectory;
+import org.apache.lucene.store.IOContext;
+import org.elasticsearch.index.store.FsDirectoryFactory;
 
 import java.io.IOException;
 
@@ -61,9 +64,15 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
         return new Lucene99FlatVectorsWriter(state, vectorsScorer);
     }
 
+    static boolean shouldUseDirectIO(SegmentReadState state) {
+        assert ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO;
+        return FsDirectoryFactory.isHybridFs(state.directory)
+            && FilterDirectory.unwrap(state.directory) instanceof DirectIOIndexInputSupplier;
+    }
+
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        if (DirectIOLucene99FlatVectorsReader.shouldUseDirectIO(state)) {
+        if (shouldUseDirectIO(state) && state.context.context() == IOContext.Context.DEFAULT) {
             // 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(

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

@@ -57,8 +57,6 @@
 @SuppressForbidden(reason = "Copied from lucene")
 public class DirectIOLucene99FlatVectorsReader extends FlatVectorsReader implements OffHeapStats {
 
-    private static final boolean USE_DIRECT_IO = Boolean.parseBoolean(System.getProperty("vector.rescoring.directio", "true"));
-
     private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(DirectIOLucene99FlatVectorsReader.class);
 
     private final IntObjectHashMap<FieldEntry> fields = new IntObjectHashMap<>();
@@ -89,10 +87,6 @@ public DirectIOLucene99FlatVectorsReader(SegmentReadState state, FlatVectorsScor
         }
     }
 
-    public static boolean shouldUseDirectIO(SegmentReadState state) {
-        return USE_DIRECT_IO && FilterDirectory.unwrap(state.directory) instanceof DirectIOIndexInputSupplier;
-    }
-
     private int readMetadata(SegmentReadState state) throws IOException {
         String metaFileName = IndexFileNames.segmentFileName(
             state.segmentInfo.name,
@@ -130,7 +124,8 @@ 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 && FilterDirectory.unwrap(state.directory) instanceof DirectIOIndexInputSupplier did
+        assert ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO;
+        IndexInput in = FilterDirectory.unwrap(state.directory) instanceof DirectIOIndexInputSupplier did
             ? did.openInputDirect(fileName, context)
             : state.directory.openInput(fileName, context);
         boolean success = false;

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

@@ -23,6 +23,7 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
@@ -86,6 +87,8 @@
  */
 public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
 
+    static final boolean USE_DIRECT_IO = Boolean.parseBoolean(System.getProperty("vector.rescoring.directio", "true"));
+
     public static final String BINARIZED_VECTOR_COMPONENT = "BVEC";
     public static final String NAME = "ES818BinaryQuantizedVectorsFormat";
 
@@ -97,9 +100,9 @@ public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
     static final String VECTOR_DATA_EXTENSION = "veb";
     static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
-    private static final DirectIOLucene99FlatVectorsFormat rawVectorFormat = new DirectIOLucene99FlatVectorsFormat(
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-    );
+    private static final FlatVectorsFormat rawVectorFormat = USE_DIRECT_IO
+        ? new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
+        : new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
 
     private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()