Use direct I/O for merge-time reads of raw vectors

docs/changelog/155919.yaml

@@ -0,0 +1,6 @@
+pr: 155919
+summary: Use direct I/O for merge-time reads of raw vectors
+area: Vector Search
+type: enhancement
+issues:
+ - 155021

docs/reference/elasticsearch/mapping-reference/dense-vector.md

@@ -602,7 +602,7 @@
 ::::
 
 `on_disk_rescore` {applies_to}`stack: preview 9.3` {applies_to}`serverless: unavailable`
-:   (Optional, boolean) Only applicable to quantized HNSW and `bbq_disk` index types. When `true`, vector rescoring will read the raw vector data directly from disk, and will not copy it in memory. This can improve performance when vector data is larger than the amount of available RAM. This setting only applies to newly-indexed vectors; after changing this setting, the vectors must be reindexed or force-merged to apply the new setting to the whole index. Defaults to `false`.
+:   (Optional, boolean) Only applicable to quantized HNSW and `bbq_disk` index types. When `true`, vector rescoring and merging will read the raw vector data directly from disk, and will not copy it in memory. This can improve performance when vector data is larger than the amount of available RAM. This setting only applies to newly-indexed vectors; after changing this setting, the vectors must be reindexed or force-merged to apply the new setting to the whole index. Defaults to `false`.
 
 `auto_calibrate` {applies_to}`stack: ga 9.5`
 :   (Optional, boolean) Only applicable to `bbq_disk`. When `true`, {{es}} automatically selects the optimal quantization encoding, oversampling factor, and preconditioning for each merged segment based on the actual recall characteristics of the merged corpus. Segments containing fewer than 10,000 vectors after merging are not calibrated and, when not otherwise specified in the mappings, use the default oversampling factor of 3.0x. Defaults to `false`. Cannot be changed after the field is created. Refer to [Auto-calibration for `bbq_disk`](/reference/elasticsearch/mapping-reference/bbq.md#bbq-auto-calibration) for details.

server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java

@@ -39,33 +39,48 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException
 
     public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException {
         if (state.context.context() == IOContext.Context.DEFAULT && useDirectIO && canUseDirectIO(state)) {
-            // only override the context for the random-access use case
+            // only wrap readers opened for searching (DEFAULT context); the wrapper adds a
+            // lazily-created, merge-hinted direct I/O reader for merges
             SegmentReadState directIOState = new SegmentReadState(
                 state.directory,
                 state.segmentInfo,
                 state.fieldInfos,
                 new DirectIOContext(state.context.hints()),
                 state.segmentSuffix
             );
-            // Use mmap for merges and direct I/O for searches.
-            return new MergeReaderWrapper(createReader(directIOState), createReader(state));
+            SegmentReadState mergeDirectIOState = new SegmentReadState(
+                state.directory,
+                state.segmentInfo,
+                state.fieldInfos,
+                new DirectIOContext(IOContext.Context.MERGE, state.context.hints()),
+                state.segmentSuffix
+            );
+            // Use direct I/O for merges too, so merge reads do not evict hotter data from the
+            // page cache. The MERGE context routes those reads to a merge-sized buffer.
+            return new MergeReaderWrapper(createReader(directIOState), () -> createReader(mergeDirectIOState));
         } else {
             return createReader(state);
         }
     }
 
     protected static class DirectIOContext implements IOContext {
 
+        private final Context context;
         final Set<FileOpenHint> hints;
 
         public DirectIOContext(Set<FileOpenHint> hints) {
+            this(Context.DEFAULT, hints);
+        }
+
+        public DirectIOContext(Context context, Set<FileOpenHint> hints) {
+            this.context = context;
             // always add DirectIOHint to the hints given
             this.hints = Sets.union(hints, Set.of(DirectIOHint.INSTANCE));
         }
 
         @Override
         public Context context() {
-            return Context.DEFAULT;
+            return context;
         }
 
         @Override
@@ -85,7 +100,7 @@ public Set<FileOpenHint> hints() {
 
         @Override
         public IOContext withHints(FileOpenHint... hints) {
-            return new DirectIOContext(Set.of(hints));
+            return new DirectIOContext(context, Set.of(hints));
         }
     }
 }

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

@@ -16,22 +16,31 @@
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.store.AlreadyClosedException;
 import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.IOSupplier;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.elasticsearch.core.IOUtils;
 
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Map;
 
+/**
+ * A {@link FlatVectorsReader} that serves searches from one reader and merges from a second,
+ * lazily-created reader, so that the two can use different I/O strategies (such as different
+ * direct I/O buffer sizes) without sharing input state.
+ */
 public class MergeReaderWrapper extends FlatVectorsReader {
 
     private final FlatVectorsReader mainReader;
-    private final FlatVectorsReader mergeReader;
+    private final IOSupplier<FlatVectorsReader> mergeReaderSupplier;
+    private FlatVectorsReader mergeReader;
+    private boolean closed;
 
-    public MergeReaderWrapper(FlatVectorsReader mainReader, FlatVectorsReader mergeReader) {
+    public MergeReaderWrapper(FlatVectorsReader mainReader, IOSupplier<FlatVectorsReader> mergeReaderSupplier) {
         this.mainReader = mainReader;
-        this.mergeReader = mergeReader;
+        this.mergeReaderSupplier = mergeReaderSupplier;
     }
 
     @Override
@@ -75,8 +84,19 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Accep
     }
 
     @Override
-    public FlatVectorsReader getMergeInstance() {
-        return mergeReader;
+    public FlatVectorsReader getMergeInstance() throws IOException {
+        // only the single thread running the merge calls this
+        if (closed) {
+            throw new AlreadyClosedException("this MergeReaderWrapper is closed");
+        }
+        // created lazily: most segments are never merged during a reader's lifetime, and the
+        // merge reader holds direct I/O resources
+        if (mergeReader == null) {
+            mergeReader = mergeReaderSupplier.get();
+        }
+        // delegate so the reader can prepare itself for merging, e.g. Lucene99FlatVectorsReader
+        // switches its data input to sequential read advice
+        return mergeReader.getMergeInstance();
     }
 
     @Override
@@ -98,6 +118,10 @@ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
 
     @Override
     public void close() throws IOException {
+        if (closed) {
+            return;
+        }
+        closed = true;
         IOUtils.close(mainReader, mergeReader);
     }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es93/DirectIOCapableLucene99FlatVectorsFormat.java

@@ -15,9 +15,7 @@
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
-import org.apache.lucene.store.IOContext;
 import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
-import org.elasticsearch.index.codec.vectors.MergeReaderWrapper;
 
 import java.io.IOException;
 
@@ -47,25 +45,4 @@ protected FlatVectorsReader createReader(SegmentReadState state) throws IOExcept
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new Lucene99FlatVectorsWriter(state, vectorsScorer);
     }
-
-    @Override
-    public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException {
-        if (state.context.context() == IOContext.Context.DEFAULT && useDirectIO && canUseDirectIO(state)) {
-            // only override the context for the random-access use case
-            SegmentReadState directIOState = new SegmentReadState(
-                state.directory,
-                state.segmentInfo,
-                state.fieldInfos,
-                new DirectIOContext(state.context.hints()),
-                state.segmentSuffix
-            );
-            // Use mmap for merges and direct I/O for searches.
-            return new MergeReaderWrapper(
-                new Lucene99FlatVectorsReader(directIOState, vectorsScorer),
-                new Lucene99FlatVectorsReader(state, vectorsScorer)
-            );
-        } else {
-            return new Lucene99FlatVectorsReader(state, vectorsScorer);
-        }
-    }
 }

server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java

@@ -177,21 +177,36 @@ private static int getBlockSize(Path path) throws IOException {
     public static final class HybridDirectory extends NIOFSDirectory {
         private final MMapDirectory delegate;
         private final DirectIODirectory directIODelegate;
+        private final DirectIODirectory mergeDirectIODelegate;
 
         public HybridDirectory(LockFactory lockFactory, MMapDirectory delegate, int asyncPrefetchLimit) throws IOException {
             super(delegate.getDirectory(), lockFactory);
             this.delegate = delegate;
 
             DirectIODirectory directIO;
+            DirectIODirectory mergeDirectIO;
             try {
-                // use 8kB buffer (two pages) to guarantee it can load all of an un-page-aligned 1024-dim float vector
-                directIO = new AlwaysDirectIODirectory(delegate, 8192, DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, asyncPrefetchLimit);
+                directIO = new AlwaysDirectIODirectory(
+                    delegate,
+                    AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE,
+                    DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT,
+                    asyncPrefetchLimit
+                );
+                // merge reads are long sequential streams: they get a merge-sized buffer and no async prefetch
+                mergeDirectIO = new AlwaysDirectIODirectory(
+                    delegate,
+                    DirectIODirectory.DEFAULT_MERGE_BUFFER_SIZE,
+                    DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT,
+                    0
+                );
             } catch (Exception e) {
                 // directio not supported
                 Log.warn("Could not initialize DirectIO access", e);
                 directIO = null;
+                mergeDirectIO = null;
             }
             this.directIODelegate = directIO;
+            this.mergeDirectIODelegate = mergeDirectIO;
         }
 
         @Override
@@ -202,8 +217,9 @@ public IndexInput openInput(String name, IOContext context) throws IOException {
                 ensureCanRead(name);
                 try {
                     Log.debug("Opening {} with direct IO", name);
-                    return directIODelegate.openInput(name, context);
-                } catch (FileSystemException e) {
+                    DirectIODirectory dio = context.context() == IOContext.Context.MERGE ? mergeDirectIODelegate : directIODelegate;
+                    return dio.openInput(name, context);
+                } catch (FileSystemException | UnsupportedOperationException e) {
                     Log.debug(() -> Strings.format("Could not open %s with direct IO", name), e);
                     directIOException = e;
                     // and fallthrough to normal opening below
@@ -306,13 +322,18 @@ MMapDirectory getDelegate() {
     }
 
     public static final class AlwaysDirectIODirectory extends DirectIODirectory {
+        // two pages, guaranteeing a single buffer can load all of an un-page-aligned 1024-dim float vector
+        public static final int RANDOM_ACCESS_BUFFER_SIZE = 8192;
+
         private final int blockSize;
+        private final int bufferSize;
         private final int asyncPrefetchLimit;
 
-        public AlwaysDirectIODirectory(FSDirectory delegate, int mergeBufferSize, long minBytesDirect, int asyncPrefetchLimit)
+        public AlwaysDirectIODirectory(FSDirectory delegate, int bufferSize, long minBytesDirect, int asyncPrefetchLimit)
             throws IOException {
-            super(delegate, mergeBufferSize, minBytesDirect);
+            super(delegate, bufferSize, minBytesDirect);
             blockSize = getBlockSize(delegate.getDirectory());
+            this.bufferSize = bufferSize;
             this.asyncPrefetchLimit = asyncPrefetchLimit;
         }
 
@@ -325,8 +346,9 @@ protected boolean useDirectIO(String name, IOContext context, OptionalLong fileL
         public IndexInput openInput(String name, IOContext context) throws IOException {
             ensureOpen();
             if (asyncPrefetchLimit > 0) {
-                return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, 8192, asyncPrefetchLimit);
+                return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, bufferSize, asyncPrefetchLimit);
             } else {
+                // no async prefetching: a plain direct-IO input at this instance's buffer size
                 return super.openInput(name, context);
             }
         }