Use a boolean fields for directIO or not

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/DirectIOCapableFlatVectorsFormat.java

@@ -0,0 +1,19 @@
+/*
+ * 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.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.index.SegmentReadState;
+
+import java.io.IOException;
+
+public interface DirectIOCapableFlatVectorsFormat {
+    FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException;
+}

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

@@ -18,11 +18,9 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
-import org.elasticsearch.index.codec.vectors.es93.DirectIOLucene99FlatVectorsFormat;
 
 import java.io.IOException;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Codec format for Inverted File Vector indexes. This index expects to break the dimensional space
@@ -59,18 +57,10 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     public static final int VERSION_START = 0;
     public static final int VERSION_CURRENT = VERSION_START;
 
-    private static final FlatVectorsFormat defaultRawVectorFormat = new Lucene99FlatVectorsFormat(
+    private static final Lucene99FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
-    private static final FlatVectorsFormat directIORawVectorFormat = new DirectIOLucene99FlatVectorsFormat(
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-    );
-    private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(
-        defaultRawVectorFormat.getName(),
-        defaultRawVectorFormat,
-        directIORawVectorFormat.getName(),
-        directIORawVectorFormat
-    );
+    private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat);
 
     // This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
     // useful when searching with 'efSearch' type parameters instead of requiring a specific ratio.
@@ -116,26 +106,12 @@ public ES920DiskBBQVectorsFormat() {
         this(DEFAULT_VECTORS_PER_CLUSTER, DEFAULT_CENTROIDS_PER_PARENT_CLUSTER);
     }
 
-    private final AtomicReference<FlatVectorsFormat> flatVectorsFormat = new AtomicReference<>();
-
-    public void useDirectIO() {
-        if (flatVectorsFormat.compareAndSet(null, directIORawVectorFormat) == false) {
-            throw new IllegalStateException("Flat vector format has already been set");
-        }
-    }
-
     @Override
     public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-        var rawFormat = flatVectorsFormat.get();
-        if (rawFormat == null) {
-            rawFormat = defaultRawVectorFormat;
-            flatVectorsFormat.set(rawFormat);
-        }
-
         return new ES920DiskBBQVectorsWriter(
-            rawFormat.getName(),
+            rawVectorFormat.getName(),
             state,
-            rawFormat.fieldsWriter(state),
+            rawVectorFormat.fieldsWriter(state),
             vectorPerCluster,
             centroidsPerParentCluster
         );

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

@@ -1,21 +1,10 @@
 /*
- * @notice
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- *
- * Modifications copyright (C) 2024 Elasticsearch B.V.
+ * 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.es93;
 
@@ -31,29 +20,22 @@
 import org.apache.lucene.store.MergeInfo;
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.index.codec.vectors.AbstractFlatVectorsFormat;
+import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.MergeReaderWrapper;
 import org.elasticsearch.index.codec.vectors.es818.DirectIOHint;
 import org.elasticsearch.index.store.FsDirectoryFactory;
 
 import java.io.IOException;
 import java.util.Set;
 
-/**
- * Copied from Lucene99FlatVectorsFormat.
- *
- * This is copied to change the implementation of {@link #fieldsReader} only.
- */
-public class DirectIOLucene99FlatVectorsFormat extends AbstractFlatVectorsFormat {
-
-    static final String NAME = "DirectIOLucene99FlatVectorsFormat";
+public class DirectIOCapableLucene99FlatVectorsFormat extends AbstractFlatVectorsFormat implements DirectIOCapableFlatVectorsFormat {
 
-    public static final int VERSION_START = 0;
-    public static final int VERSION_CURRENT = VERSION_START;
+    static final String NAME = "Lucene99FlatVectorsFormat";
 
     private final FlatVectorsScorer vectorsScorer;
 
     /** Constructs a format */
-    public DirectIOLucene99FlatVectorsFormat(FlatVectorsScorer vectorsScorer) {
+    public DirectIOCapableLucene99FlatVectorsFormat(FlatVectorsScorer vectorsScorer) {
         super(NAME);
         this.vectorsScorer = vectorsScorer;
     }
@@ -68,13 +50,18 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
         return new Lucene99FlatVectorsWriter(state, vectorsScorer);
     }
 
-    static boolean shouldUseDirectIO(SegmentReadState state) {
+    static boolean canUseDirectIO(SegmentReadState state) {
         return FsDirectoryFactory.isHybridFs(state.directory);
     }
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        if (shouldUseDirectIO(state) && state.context.context() == IOContext.Context.DEFAULT) {
+        return fieldsReader(state, false);
+    }
+
+    @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,
@@ -84,7 +71,6 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException
                 state.segmentSuffix
             );
             // 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(
                 new Lucene99FlatVectorsReader(directIOState, vectorsScorer),
                 new Lucene99FlatVectorsReader(state, vectorsScorer)

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

@@ -20,14 +20,13 @@
 package org.elasticsearch.index.codec.vectors.es93;
 
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
-import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 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.AbstractFlatVectorsFormat;
+import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryFlatVectorsScorer;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsReader;
@@ -88,21 +87,19 @@
   *  <li>The sparse vector information, if required, mapping vector ordinal to doc ID
   * </ul>
  */
-public class ES93BinaryQuantizedVectorsFormat extends AbstractFlatVectorsFormat {
+public class ES93BinaryQuantizedVectorsFormat extends AbstractFlatVectorsFormat implements DirectIOCapableFlatVectorsFormat {
 
     public static final String NAME = "ES93BinaryQuantizedVectorsFormat";
 
-    private final FlatVectorsFormat rawVectorFormat;
+    private final DirectIOCapableLucene99FlatVectorsFormat rawVectorFormat;
 
     private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
-    public ES93BinaryQuantizedVectorsFormat(boolean useDirectIO) {
-        super(useDirectIO ? "DirectIO" + NAME : NAME);
-        rawVectorFormat = useDirectIO
-            ? new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
-            : new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
+    public ES93BinaryQuantizedVectorsFormat() {
+        super(NAME);
+        rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
     }
 
     @Override
@@ -119,4 +116,9 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
         return new ES818BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), scorer);
     }
+
+    @Override
+    public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException {
+        return new ES818BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state, useDirectIO), scorer);
+    }
 }

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

@@ -19,7 +19,6 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
-import org.apache.lucene.util.IOFunction;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.elasticsearch.core.IOUtils;
 
@@ -35,7 +34,12 @@ class ES93GenericFlatVectorsReader extends FlatVectorsReader {
 
     private final FlatVectorsReader vectorsReader;
 
-    ES93GenericFlatVectorsReader(SegmentReadState state, IOFunction<String, FlatVectorsReader> getFormatReader) throws IOException {
+    @FunctionalInterface
+    interface GetFormatReader {
+        FlatVectorsReader getReader(String formatName, boolean useDirectIO) throws IOException;
+    }
+
+    ES93GenericFlatVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException {
         super(null);    // Hacks ahoy!
         // read in the meta information
         final String metaFileName = IndexFileNames.segmentFileName(
@@ -57,9 +61,12 @@ class ES93GenericFlatVectorsReader extends FlatVectorsReader {
                     state.segmentSuffix
                 );
                 String innerFormatName = metaIn.readString();
-                reader = getFormatReader.apply(innerFormatName);
+                byte useDirectIO = metaIn.readByte();
+                reader = getFormatReader.getReader(innerFormatName, useDirectIO == 1);
                 if (reader == null) {
-                    throw new IllegalStateException("Cannot find knn vector format: " + innerFormatName);
+                    throw new IllegalStateException(
+                        "Cannot find knn vector format [" + innerFormatName + "]" + (useDirectIO == 1 ? " with directIO" : "")
+                    );
                 }
             } catch (Throwable exception) {
                 priorE = exception;

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

@@ -33,7 +33,8 @@ class ES93GenericFlatVectorsWriter extends FlatVectorsWriter {
     private final FlatVectorsWriter rawVectorWriter;
 
     @SuppressWarnings("this-escape")
-    ES93GenericFlatVectorsWriter(String knnFormatName, SegmentWriteState state, FlatVectorsWriter rawWriter) throws IOException {
+    ES93GenericFlatVectorsWriter(String knnFormatName, boolean useDirectIOReads, SegmentWriteState state, FlatVectorsWriter rawWriter)
+        throws IOException {
         super(rawWriter.getFlatVectorScorer());
         this.rawVectorWriter = rawWriter;
         final String metaFileName = IndexFileNames.segmentFileName(
@@ -46,6 +47,7 @@ class ES93GenericFlatVectorsWriter extends FlatVectorsWriter {
             CodecUtil.writeIndexHeader(metaOut, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix);
             // write the format name used for this segment
             metaOut.writeString(knnFormatName);
+            metaOut.writeByte(useDirectIOReads ? (byte) 1 : 0);
         } catch (Throwable t) {
             IOUtils.closeWhileHandlingException(this);
             throw t;

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

@@ -28,6 +28,7 @@
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.elasticsearch.index.codec.vectors.AbstractHnswVectorsFormat;
+import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
 
 import java.io.IOException;
 import java.util.Map;
@@ -61,16 +62,19 @@ protected final FlatVectorsFormat flatVectorsFormat() {
 
     protected abstract FlatVectorsFormat writeFlatVectorsFormat();
 
+    protected abstract boolean useDirectIOReads();
+
     protected abstract Map<String, FlatVectorsFormat> supportedReadFlatVectorsFormats();
 
     @Override
     public final KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         var flatFormat = writeFlatVectorsFormat();
+        boolean directIO = useDirectIOReads();
         return new Lucene99HnswVectorsWriter(
             state,
             maxConn,
             beamWidth,
-            new ES93GenericFlatVectorsWriter(flatFormat.getName(), state, flatFormat.fieldsWriter(state)),
+            new ES93GenericFlatVectorsWriter(flatFormat.getName(), directIO, state, flatFormat.fieldsWriter(state)),
             numMergeWorkers,
             mergeExec
         );
@@ -79,10 +83,16 @@ public final KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExc
     @Override
     public final KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
         var readFormats = supportedReadFlatVectorsFormats();
-        return new Lucene99HnswVectorsReader(state, new ES93GenericFlatVectorsReader(state, f -> {
+        return new Lucene99HnswVectorsReader(state, new ES93GenericFlatVectorsReader(state, (f, dio) -> {
             var format = readFormats.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);
+            }
         }));
     }