Copy Lucene99FlatVectorsReader allowing direct IO to be specified directly

Author: thecoop

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIODirectory.java

@@ -0,0 +1,23 @@
+/*
+ * 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.es818;
+
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+
+import java.io.IOException;
+
+/**
+ * A hook for {@link ES818FlatVectorsReader} to specify the input should be opened using DirectIO.
+ * Remove when IOContext allows more extensible payloads to be specified.
+ */
+public interface DirectIODirectory {
+    IndexInput openInputDirect(String name, IOContext context) throws IOException;
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java

@@ -23,7 +23,6 @@
 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;
 
@@ -97,7 +96,7 @@ public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
     static final String VECTOR_DATA_EXTENSION = "veb";
     static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
-    private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
+    private static final ES818FlatVectorsFormat rawVectorFormat = new ES818FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818FlatVectorsFormat.java

@@ -0,0 +1,67 @@
+/*
+ * @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.
+ */
+package org.elasticsearch.index.codec.vectors.es818;
+
+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.Lucene99FlatVectorsWriter;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+
+import java.io.IOException;
+
+/** Copied from Lucene99FlatVectorsFormat in Lucene 10.1 */
+public class ES818FlatVectorsFormat extends FlatVectorsFormat {
+
+    static final String NAME = "Lucene99FlatVectorsFormat";
+    static final String META_CODEC_NAME = "Lucene99FlatVectorsFormatMeta";
+    static final String VECTOR_DATA_CODEC_NAME = "Lucene99FlatVectorsFormatData";
+    static final String META_EXTENSION = "vemf";
+    static final String VECTOR_DATA_EXTENSION = "vec";
+
+    public static final int VERSION_START = 0;
+    public static final int VERSION_CURRENT = VERSION_START;
+
+    static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
+    private final FlatVectorsScorer vectorsScorer;
+
+    /** Constructs a format */
+    public ES818FlatVectorsFormat(FlatVectorsScorer vectorsScorer) {
+        super(NAME);
+        this.vectorsScorer = vectorsScorer;
+    }
+
+    @Override
+    public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+        return new Lucene99FlatVectorsWriter(state, vectorsScorer);
+    }
+
+    @Override
+    public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
+        return new ES818FlatVectorsReader(state, vectorsScorer);
+    }
+
+    @Override
+    public String toString() {
+        return "ES818FlatVectorsFormat(" + "vectorsScorer=" + vectorsScorer + ')';
+    }
+}