Enable direct IO for disk BBQ

Author: thecoop

server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java

@@ -74,7 +74,7 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
 
     private String indexVectors(boolean directIO) {
         String indexName = "test-vectors-" + directIO;
-        String type = randomFrom("bbq_hnsw"/*, "bbq_disk"*/);
+        String type = randomFrom("bbq_hnsw", "bbq_disk");
         assertAcked(
             prepareCreate(indexName).setSettings(Settings.builder().put(InternalSettingsPlugin.USE_COMPOUND_FILE.getKey(), false))
                 .setMapping(Strings.format("""

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java

@@ -14,10 +14,11 @@
 import org.apache.lucene.codecs.KnnVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
-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.DirectIOCapableFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
+import org.elasticsearch.index.codec.vectors.es93.DirectIOCapableLucene99FlatVectorsFormat;
 
 import java.io.IOException;
 import java.util.Map;
@@ -55,9 +56,10 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     static final String IVF_META_EXTENSION = "mivf";
 
     public static final int VERSION_START = 0;
-    public static final int VERSION_CURRENT = VERSION_START;
+    public static final int VERSION_DIRECT_IO = 1;
+    public static final int VERSION_CURRENT = VERSION_DIRECT_IO;
 
-    private static final Lucene99FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
+    private static final FlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
     private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat);
@@ -74,8 +76,13 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
 
     private final int vectorPerCluster;
     private final int centroidsPerParentCluster;
+    private final boolean useDirectIO;
 
     public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) {
+        this(vectorPerCluster, centroidsPerParentCluster, false);
+    }
+
+    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO) {
         super(NAME);
         if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) {
             throw new IllegalArgumentException(
@@ -99,6 +106,7 @@ public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentClu
         }
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
+        this.useDirectIO = useDirectIO;
     }
 
     /** Constructs a format using the given graph construction parameters and scalar quantization. */
@@ -110,6 +118,7 @@ public ES920DiskBBQVectorsFormat() {
     public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new ES920DiskBBQVectorsWriter(
             rawVectorFormat.getName(),
+            useDirectIO,
             state,
             rawVectorFormat.fieldsWriter(state),
             vectorPerCluster,
@@ -119,10 +128,16 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
 
     @Override
     public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new ES920DiskBBQVectorsReader(state, f -> {
+        return new ES920DiskBBQVectorsReader(state, (f, dio) -> {
             var format = supportedFormats.get(f);
             if (format == null) return null;
-            return format.fieldsReader(state);
+
+            if (format instanceof DirectIOCapableFlatVectorsFormat diof) {
+                return diof.fieldsReader(state, dio);
+            } else {
+                assert dio == false : format + " is not DirectIO capable";
+                return format.fieldsReader(state);
+            }
         });
     }
 

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java

@@ -9,14 +9,12 @@
 
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
-import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.Bits;
-import org.apache.lucene.util.IOFunction;
 import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue;
@@ -40,7 +38,7 @@
  */
 public class ES920DiskBBQVectorsReader extends IVFVectorsReader {
 
-    public ES920DiskBBQVectorsReader(SegmentReadState state, IOFunction<String, FlatVectorsReader> getFormatReader) throws IOException {
+    ES920DiskBBQVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException {
         super(state, getFormatReader);
     }
 

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java

@@ -53,12 +53,13 @@ public class ES920DiskBBQVectorsWriter extends IVFVectorsWriter {
 
     public ES920DiskBBQVectorsWriter(
         String rawVectorFormatName,
+        boolean useDirectIOReads,
         SegmentWriteState state,
         FlatVectorsWriter rawVectorDelegate,
         int vectorPerCluster,
         int centroidsPerParentCluster
     ) throws IOException {
-        super(state, rawVectorFormatName, rawVectorDelegate);
+        super(state, rawVectorFormatName, useDirectIOReads, rawVectorDelegate);
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
     }

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java

@@ -29,7 +29,6 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.Bits;
-import org.apache.lucene.util.IOFunction;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.search.vectors.IVFKnnSearchStrategy;
 
@@ -43,6 +42,7 @@
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.DYNAMIC_VISIT_RATIO;
+import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO;
 
 /**
  * Reader for IVF vectors. This reader is used to read the IVF vectors from the index.
@@ -55,8 +55,13 @@ public abstract class IVFVectorsReader extends KnnVectorsReader {
     protected final IntObjectHashMap<FieldEntry> fields;
     private final Map<String, FlatVectorsReader> rawVectorReaders;
 
+    @FunctionalInterface
+    protected interface GetFormatReader {
+        FlatVectorsReader getReader(String formatName, boolean useDirectIO) throws IOException;
+    }
+
     @SuppressWarnings("this-escape")
-    protected IVFVectorsReader(SegmentReadState state, IOFunction<String, FlatVectorsReader> getFormatReader) throws IOException {
+    protected IVFVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException {
         this.state = state;
         this.fieldInfos = state.fieldInfos;
         this.fields = new IntObjectHashMap<>();
@@ -80,7 +85,7 @@ protected IVFVectorsReader(SegmentReadState state, IOFunction<String, FlatVector
                     state.segmentInfo.getId(),
                     state.segmentSuffix
                 );
-                readers = readFields(ivfMeta, getFormatReader);
+                readers = readFields(ivfMeta, getFormatReader, versionMeta);
             } catch (Throwable exception) {
                 priorE = exception;
             } finally {
@@ -153,7 +158,7 @@ private static IndexInput openDataInput(
         }
     }
 
-    private Map<String, FlatVectorsReader> readFields(ChecksumIndexInput meta, IOFunction<String, FlatVectorsReader> loadReader)
+    private Map<String, FlatVectorsReader> readFields(ChecksumIndexInput meta, GetFormatReader loadReader, int versionMeta)
         throws IOException {
         Map<String, FlatVectorsReader> readers = new HashMap<>();
         for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) {
@@ -162,11 +167,11 @@ private Map<String, FlatVectorsReader> readFields(ChecksumIndexInput meta, IOFun
                 throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta);
             }
 
-            FieldEntry fieldEntry = readField(meta, info);
+            FieldEntry fieldEntry = readField(meta, info, versionMeta);
 
             FlatVectorsReader reader = readers.get(fieldEntry.rawVectorFormatName);
             if (reader == null) {
-                reader = loadReader.apply(fieldEntry.rawVectorFormatName);
+                reader = loadReader.getReader(fieldEntry.rawVectorFormatName, fieldEntry.useDirectIOReads);
                 if (reader == null) {
                     throw new IllegalStateException("Cannot find flat vector format: " + fieldEntry.rawVectorFormatName);
                 }
@@ -178,8 +183,9 @@ private Map<String, FlatVectorsReader> readFields(ChecksumIndexInput meta, IOFun
         return readers;
     }
 
-    private FieldEntry readField(IndexInput input, FieldInfo info) throws IOException {
+    private FieldEntry readField(IndexInput input, FieldInfo info, int versionMeta) throws IOException {
         final String rawVectorFormat = input.readString();
+        final boolean useDirectIOReads = versionMeta >= VERSION_DIRECT_IO && input.readByte() == 1;
         final VectorEncoding vectorEncoding = readVectorEncoding(input);
         final VectorSimilarityFunction similarityFunction = readSimilarityFunction(input);
         if (similarityFunction != info.getVectorSimilarityFunction()) {
@@ -207,6 +213,7 @@ private FieldEntry readField(IndexInput input, FieldInfo info) throws IOExceptio
         }
         return new FieldEntry(
             rawVectorFormat,
+            useDirectIOReads,
             similarityFunction,
             vectorEncoding,
             numCentroids,
@@ -387,6 +394,7 @@ public void close() throws IOException {
 
     protected record FieldEntry(
         String rawVectorFormatName,
+        boolean useDirectIOReads,
         VectorSimilarityFunction similarityFunction,
         VectorEncoding vectorEncoding,
         int numCentroids,

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsWriter.java

@@ -53,12 +53,18 @@ public abstract class IVFVectorsWriter extends KnnVectorsWriter {
     private final IndexOutput ivfCentroids, ivfClusters;
     private final IndexOutput ivfMeta;
     private final String rawVectorFormatName;
+    private final boolean useDirectIOReads;
     private final FlatVectorsWriter rawVectorDelegate;
 
     @SuppressWarnings("this-escape")
-    protected IVFVectorsWriter(SegmentWriteState state, String rawVectorFormatName, FlatVectorsWriter rawVectorDelegate)
-        throws IOException {
+    protected IVFVectorsWriter(
+        SegmentWriteState state,
+        String rawVectorFormatName,
+        boolean useDirectIOReads,
+        FlatVectorsWriter rawVectorDelegate
+    ) throws IOException {
         this.rawVectorFormatName = rawVectorFormatName;
+        this.useDirectIOReads = useDirectIOReads;
         this.rawVectorDelegate = rawVectorDelegate;
         final String metaFileName = IndexFileNames.segmentFileName(
             state.segmentInfo.name,
@@ -497,6 +503,7 @@ private void writeMeta(
     ) throws IOException {
         ivfMeta.writeInt(field.number);
         ivfMeta.writeString(rawVectorFormatName);
+        ivfMeta.writeByte(useDirectIOReads ? (byte) 1 : 0);
         ivfMeta.writeInt(field.getVectorEncoding().ordinal());
         ivfMeta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction()));
         ivfMeta.writeInt(numCentroids);

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -2174,7 +2174,11 @@ static class BBQIVFIndexOptions extends QuantizedIndexOptions {
         @Override
         KnnVectorsFormat getVectorsFormat(ElementType elementType) {
             assert elementType == ElementType.FLOAT;
-            return new ES920DiskBBQVectorsFormat(clusterSize, ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER);
+            return new ES920DiskBBQVectorsFormat(
+                clusterSize,
+                ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
+                onDiskRescore
+            );
         }
 
         @Override