diff --git a/muted-tests.yml b/muted-tests.yml index 96dde0c3f30ce..e3a63cb7d87e3 100644 --- a/muted-tests.yml +++ b/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 diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java index 02de4d5450d35..a03cd3fc9fe8e 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java +++ b/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( diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsReader.java index 30829d513c5e7..f304fe3aa4e88 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsReader.java +++ b/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 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; diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java index c3e12a765eb5f..5df81c35bbd40 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java +++ b/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()