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 @@ $$$dense-vector-index-options$$$
 ::::
 
 `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 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. For `bbq_hnsw` indices, segment merges also read the raw vector data directly from disk where the platform supports it, so that merging is less likely to evict more frequently accessed data from the filesystem cache. {applies_to}`stack: ga 9.6` 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,28 +39,43 @@ 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(state.context.hints(), Set.of(DirectIOHint.INSTANCE, DirectIOMergeHint.INSTANCE)),
+                state.segmentSuffix
+            );
+            // Use direct I/O for merges too, so merge reads do not evict hotter data from the
+            // page cache. The merge hint makes those reads use a merge-sized buffer.
+            return new MergeReaderWrapper(createReader(directIOState), () -> createReader(mergeDirectIOState));
         } else {
             return createReader(state);
         }
     }
 
     protected static class DirectIOContext implements IOContext {
 
+        private final Set<FileOpenHint> stickyHints;
         final Set<FileOpenHint> hints;
 
         public DirectIOContext(Set<FileOpenHint> hints) {
-            // always add DirectIOHint to the hints given
-            this.hints = Sets.union(hints, Set.of(DirectIOHint.INSTANCE));
+            this(hints, Set.of(DirectIOHint.INSTANCE));
+        }
+
+        public DirectIOContext(Set<FileOpenHint> hints, Set<FileOpenHint> stickyHints) {
+            // the sticky hints are always added, and survive withHints() replacing the others
+            this.stickyHints = stickyHints;
+            this.hints = Sets.union(hints, stickyHints);
         }
 
         @Override
@@ -85,7 +100,7 @@ public Set<FileOpenHint> hints() {
 
         @Override
         public IOContext withHints(FileOpenHint... hints) {
-            return new DirectIOContext(Set.of(hints));
+            return new DirectIOContext(Set.of(hints), stickyHints);
         }
     }
 }

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

@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+import org.apache.lucene.store.IOContext;
+
+/**
+ * Hint that a file opened with direct I/O will be read as a long sequential stream by a merge,
+ * so reads should use a merge-sized buffer rather than the small buffer used for random access.
+ */
+public enum DirectIOMergeHint implements IOContext.FileOpenHint {
+    INSTANCE
+}

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

@@ -16,22 +16,32 @@
 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 final Object lock = new Object();
+    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 +85,20 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Accep
     }
 
     @Override
-    public FlatVectorsReader getMergeInstance() {
-        return mergeReader;
+    public FlatVectorsReader getMergeInstance() throws IOException {
+        synchronized (lock) {
+            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 +120,12 @@ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
 
     @Override
     public void close() throws IOException {
-        IOUtils.close(mainReader, mergeReader);
+        synchronized (lock) {
+            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/AsyncDirectIOIndexInput.java

@@ -401,6 +401,10 @@ public IndexInput slice(String sliceDescription, long offset, long length) throw
     }
 
     // pkg private for testing
+    int bufferCapacity() {
+        return buffer.capacity();
+    }
+
     int prefetchSlots() {
         return prefetcher.posToSlot.size();
     }
@@ -448,6 +452,10 @@ private static class DirectIOPrefetcher implements Closeable {
          */
         void prefetch(long pos, long length) {
             assert pos % blockSize == 0 : "prefetch pos [" + pos + "] must be aligned to block size [" + blockSize + "]";
+            if (maxConcurrentPrefetches == 0) {
+                // opened without prefetch slots, e.g. a merge-hinted input when async prefetch is disabled
+                return;
+            }
             // first determine how many slots we need given the length
             while (length > 0) {
                 Map.Entry<Long, Integer> floor = this.posToSlot.floorEntry(pos);

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

@@ -30,6 +30,7 @@
 import org.elasticsearch.index.IndexModule;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.StandardIOBehaviorHint;
+import org.elasticsearch.index.codec.vectors.DirectIOMergeHint;
 import org.elasticsearch.index.codec.vectors.es818.DirectIOHint;
 import org.elasticsearch.index.shard.ShardPath;
 import org.elasticsearch.logging.LogManager;
@@ -184,8 +185,12 @@ public HybridDirectory(LockFactory lockFactory, MMapDirectory delegate, int asyn
 
             DirectIODirectory directIO;
             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
+                );
             } catch (Exception e) {
                 // directio not supported
                 Log.warn("Could not initialize DirectIO access", e);
@@ -203,7 +208,7 @@ public IndexInput openInput(String name, IOContext context) throws IOException {
                 try {
                     Log.debug("Opening {} with direct IO", name);
                     return directIODelegate.openInput(name, context);
-                } catch (FileSystemException e) {
+                } 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,6 +311,9 @@ 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 asyncPrefetchLimit;
 
@@ -324,8 +332,18 @@ protected boolean useDirectIO(String name, IOContext context, OptionalLong fileL
         @Override
         public IndexInput openInput(String name, IOContext context) throws IOException {
             ensureOpen();
-            if (asyncPrefetchLimit > 0) {
-                return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, 8192, asyncPrefetchLimit);
+            // merge reads are long sequential streams, so they use a merge-sized buffer rather
+            // than the two-page buffer that random-access rescore reads use. They get no prefetch
+            // slots: the large buffer is their readahead, and per-slot prefetch buffers are sized
+            // by the read buffer, so slots would multiply the prefetch-limit setting's documented
+            // memory bound by the merge buffer size.
+            boolean forMerge = context.hints().contains(DirectIOMergeHint.INSTANCE);
+            int bufferSize = forMerge ? DEFAULT_MERGE_BUFFER_SIZE : RANDOM_ACCESS_BUFFER_SIZE;
+            int prefetchLimit = forMerge ? 0 : asyncPrefetchLimit;
+            if (asyncPrefetchLimit > 0 || forMerge) {
+                // merge-hinted opens always take this path, since the delegate's buffer size is
+                // fixed at construction and cannot honor the merge hint
+                return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, bufferSize, prefetchLimit);
             } else {
                 return super.openInput(name, context);
             }