Add a Faiss codec for KNN searches

gradle/testing/defaults-tests.gradle

@@ -142,6 +142,7 @@ allprojects {
               ':lucene:core',
               ':lucene:codecs',
               ":lucene:distribution.tests",
+              ':lucene:sandbox',
               ":lucene:test-framework"
       ] ? 'ALL-UNNAMED' : 'org.apache.lucene.core')
 

gradle/testing/randomization/policies/tests.policy

@@ -80,6 +80,9 @@ grant {
   permission java.io.FilePermission "${hunspell.corpora}${/}-", "read";
   permission java.io.FilePermission "${hunspell.dictionaries}", "read";
   permission java.io.FilePermission "${hunspell.dictionaries}${/}-", "read";
+
+  // allow loading specific library for the Faiss codec
+  permission java.lang.RuntimePermission "loadLibrary.faiss_c";
 };
 
 // Permissions for jacoco code coverage

lucene/sandbox/src/java/module-info.java

@@ -22,6 +22,7 @@
   requires org.apache.lucene.facet;
 
   exports org.apache.lucene.payloads;
+  exports org.apache.lucene.sandbox.codecs.faiss;
   exports org.apache.lucene.sandbox.codecs.idversion;
   exports org.apache.lucene.sandbox.codecs.quantization;
   exports org.apache.lucene.sandbox.document;
@@ -39,4 +40,6 @@
 
   provides org.apache.lucene.codecs.PostingsFormat with
       org.apache.lucene.sandbox.codecs.idversion.IDVersionPostingsFormat;
+  provides org.apache.lucene.codecs.KnnVectorsFormat with
+      org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat;
 }

lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissKnnVectorsFormat.java

@@ -0,0 +1,91 @@
+/*
+ * 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.
+ */
+package org.apache.lucene.sandbox.codecs.faiss;
+
+import java.io.IOException;
+import java.util.Locale;
+import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.codecs.KnnVectorsReader;
+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;
+
+/**
+ * A format which uses <a href="https://github.com/facebookresearch/faiss">Faiss</a> to create and
+ * search vector indexes, using {@link LibFaissC} to interact with the native library.
+ *
+ * <p>A separate Faiss index is created per-segment, and uses the following files:
+ *
+ * <ul>
+ *   <li><code>.faissm</code> (metadata file): stores field number, offset and length of actual
+ *       Faiss index in data file.
+ *   <li><code>.faissd</code> (data file): stores concatenated Faiss indexes for all fields.
+ *   <li>All files required by {@link Lucene99FlatVectorsFormat} for storing raw vectors.
+ * </ul>
+ *
+ * @lucene.experimental
+ */
+public final class FaissKnnVectorsFormat extends KnnVectorsFormat {
+  public static final String NAME = FaissKnnVectorsFormat.class.getSimpleName();
+  static final int VERSION_START = 0;
+  static final int VERSION_CURRENT = VERSION_START;
+  static final String META_CODEC_NAME = NAME + "Meta";
+  static final String DATA_CODEC_NAME = NAME + "Data";
+  static final String META_EXTENSION = "faissm";
+  static final String DATA_EXTENSION = "faissd";
+
+  private final String description;
+  private final String indexParams;
+  private final FlatVectorsFormat rawVectorsFormat;
+
+  public FaissKnnVectorsFormat() {
+    this("IDMap,HNSW32", "efConstruction=200");
+  }
+
+  public FaissKnnVectorsFormat(String description, String indexParams) {
+    super(NAME);
+    this.description = description;
+    this.indexParams = indexParams;
+    this.rawVectorsFormat =
+        new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
+  }
+
+  @Override
+  public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+    return new FaissKnnVectorsWriter(
+        description, indexParams, state, rawVectorsFormat.fieldsWriter(state));
+  }
+
+  @Override
+  public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
+    return new FaissKnnVectorsReader(state, rawVectorsFormat.fieldsReader(state));
+  }
+
+  @Override
+  public int getMaxDimensions(String fieldName) {
+    return DEFAULT_MAX_DIMENSIONS;
+  }
+
+  @Override
+  public String toString() {
+    return String.format(
+        Locale.ROOT, "%s(description=%s indexParams=%s)", NAME, description, indexParams);
+  }
+}

lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissKnnVectorsReader.java

@@ -0,0 +1,188 @@
+/*
+ * 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.
+ */
+package org.apache.lucene.sandbox.codecs.faiss;
+
+import static org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat.DATA_CODEC_NAME;
+import static org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat.DATA_EXTENSION;
+import static org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat.META_CODEC_NAME;
+import static org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat.META_EXTENSION;
+import static org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat.VERSION_CURRENT;
+import static org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat.VERSION_START;
+import static org.apache.lucene.sandbox.codecs.faiss.LibFaissC.indexRead;
+import static org.apache.lucene.sandbox.codecs.faiss.LibFaissC.indexSearch;
+
+import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.ReadAdvice;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IOUtils;
+
+/**
+ * Read per-segment Faiss indexes and associated metadata.
+ *
+ * @lucene.experimental
+ */
+public final class FaissKnnVectorsReader extends KnnVectorsReader {
+  private final FlatVectorsReader rawVectorsReader;
+  private final IndexInput meta, data;
+  private final Map<String, IndexEntry> indexMap;
+  private final Arena arena;
+  private boolean closed;
+
+  public FaissKnnVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader)
+      throws IOException {
+    this.rawVectorsReader = rawVectorsReader;
+    this.indexMap = new HashMap<>();
+    this.arena = Arena.ofShared();
+    this.closed = false;
+
+    boolean failure = true;
+    try {
+      meta =
+          openInput(
+              state,
+              META_EXTENSION,
+              META_CODEC_NAME,
+              VERSION_START,
+              VERSION_CURRENT,
+              state.context);
+      data =
+          openInput(
+              state,
+              DATA_EXTENSION,
+              DATA_CODEC_NAME,
+              VERSION_START,
+              VERSION_CURRENT,
+              state.context.withReadAdvice(ReadAdvice.RANDOM));
+
+      Map.Entry<String, IndexEntry> entry;
+      while ((entry = parseNextField(state)) != null) {
+        this.indexMap.put(entry.getKey(), entry.getValue());
+      }
+
+      failure = false;
+    } finally {
+      if (failure) {
+        IOUtils.closeWhileHandlingException(this);
+      }
+    }
+  }
+
+  @SuppressWarnings("SameParameterValue")
+  private IndexInput openInput(
+      SegmentReadState state,
+      String extension,
+      String codecName,
+      int versionStart,
+      int versionEnd,
+      IOContext context)
+      throws IOException {
+
+    String fileName =
+        IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, extension);
+    IndexInput input = state.directory.openInput(fileName, context);
+    CodecUtil.checkIndexHeader(
+        input, codecName, versionStart, versionEnd, state.segmentInfo.getId(), state.segmentSuffix);
+    return input;
+  }
+
+  private Map.Entry<String, IndexEntry> parseNextField(SegmentReadState state) throws IOException {
+    int fieldNumber = meta.readInt();
+    if (fieldNumber == -1) {
+      return null;
+    }
+
+    FieldInfo fieldInfo = state.fieldInfos.fieldInfo(fieldNumber);
+    if (fieldInfo == null) {
+      throw new IllegalStateException("Invalid field");
+    }
+
+    long dataOffset = meta.readLong();
+    long dataLength = meta.readLong();
+
+    // See flags defined in c_api/index_io_c.h
+    int ioFlags = 3;
+
+    // Read index into memory
+    MemorySegment indexPointer =
+        indexRead(data.slice(fieldInfo.name, dataOffset, dataLength), ioFlags)
+            // Ensure timely cleanup
+            .reinterpret(arena, LibFaissC::freeIndex);
+
+    return Map.entry(
+        fieldInfo.name, new IndexEntry(indexPointer, fieldInfo.getVectorSimilarityFunction()));
+  }
+
+  @Override
+  public void checkIntegrity() throws IOException {
+    rawVectorsReader.checkIntegrity();
+    CodecUtil.checksumEntireFile(meta);
+    CodecUtil.checksumEntireFile(data);
+  }
+
+  @Override
+  public FloatVectorValues getFloatVectorValues(String field) throws IOException {
+    return rawVectorsReader.getFloatVectorValues(field);
+  }
+
+  @Override
+  public ByteVectorValues getByteVectorValues(String field) {
+    // TODO: Support using SQ8 quantization, see
+    // https://github.com/opensearch-project/k-NN/pull/2425
+    throw new UnsupportedOperationException("Byte vectors not supported");
+  }
+
+  @Override
+  public void search(String field, float[] vector, KnnCollector knnCollector, Bits acceptDocs) {
+    IndexEntry entry = indexMap.get(field);
+    if (entry != null) {
+      indexSearch(entry.indexPointer, entry.function, vector, knnCollector, acceptDocs);
+    }
+  }
+
+  @Override
+  public void search(String field, byte[] vector, KnnCollector knnCollector, Bits acceptDocs) {
+    // TODO: Support using SQ8 quantization, see
+    // https://github.com/opensearch-project/k-NN/pull/2425
+    throw new UnsupportedOperationException("Byte vectors not supported");
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (closed == false) {
+      IOUtils.close(rawVectorsReader, arena::close, meta, data);
+      closed = true;
+    }
+  }
+
+  private record IndexEntry(MemorySegment indexPointer, VectorSimilarityFunction function) {}
+}