Add a Faiss codec for KNN searches

.github/workflows/run-special-checks-sandbox.yml

@@ -0,0 +1,58 @@
+name: "Run special checks: module lucene/sandbox"
+
+on:
+  workflow_dispatch:
+
+  pull_request:
+    branches:
+      - '*'
+
+  push:
+    branches:
+      - 'main'
+      - 'branch_10x'
+
+jobs:
+  faiss-tests:
+    name: tests for the Faiss codec (v${{ matrix.faiss-version }} with JDK ${{ matrix.java }} on ${{ matrix.os }})
+    timeout-minutes: 15
+
+    strategy:
+      matrix:
+        os: [ ubuntu-latest ]
+        java: [ '24' ]
+        faiss-version: [ '1.11.0' ]
+
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Install Mamba
+        uses: conda-incubator/setup-miniconda@835234971496cad1653abb28a638a281cf32541f #v3.2.0
+        with:
+          miniforge-version: 'latest'
+          auto-activate-base: 'false'
+          activate-environment: 'faiss-env'
+          # TODO: Use only conda-forge if possible, see https://github.com/conda-forge/faiss-split-feedstock/pull/88
+          channels: 'pytorch,conda-forge'
+          conda-remove-defaults: 'true'
+
+      - name: Install Faiss
+        run: mamba install faiss-cpu=${{ matrix.faiss-version }}
+
+      - name: Checkout Lucene
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+
+      - name: Prepare Lucene workspace
+        uses: ./.github/actions/prepare-for-build
+
+      - name: Run tests for Faiss codec
+        run: >
+          LD_LIBRARY_PATH=$CONDA_PREFIX/lib
+          ./gradlew -p lucene/sandbox
+          -Dtests.faiss.run=true
+          test
+          --tests "org.apache.lucene.sandbox.codecs.faiss.*"
+
+    defaults:
+      run:
+        shell: bash -leo pipefail {0}

build-tools/build-infra/src/main/groovy/lucene.java.tests-and-randomization.gradle

@@ -147,6 +147,9 @@ buildOptions.addOption("tests.file.encoding", "Sets the default file.encoding on
   ])
 })
 
+buildOptions.addBooleanOption("tests.faiss.run", "Explicitly run tests for the Faiss codec.", false)
+optionsInheritedAsProperties += ["tests.faiss.run"]
+
 // TODO: do we still use these?
 // Test data file used.
 // [propName: 'tests.linedocsfile', value: 'europarl.lines.txt.gz', description: "Test data file path."],

lucene/CHANGES.txt

@@ -30,6 +30,8 @@ New Features
 ---------------------
 * GITHUB#14097: Binary partitioning merge policy over float-valued vector field. (Mike Sokolov)
 
+* GITHUB#14178: Add a Faiss-based vector format in the sandbox module. (Kaival Parikh)
+
 Improvements
 ---------------------
 

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,120 @@
+/*
+ * 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.util.hnsw.HnswGraphBuilder.DEFAULT_BEAM_WIDTH;
+import static org.apache.lucene.util.hnsw.HnswGraphBuilder.DEFAULT_MAX_CONN;
+
+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 Faiss-based format to create and search vector indexes, using {@link LibFaissC} to interact
+ * with the native library.
+ *
+ * <p>The Faiss index is configured using its flexible <a
+ * href="https://github.com/facebookresearch/faiss/wiki/The-index-factory">index factory</a>, which
+ * allows creating arbitrary indexes by "describing" them. These indexes can be tuned by <a
+ * href="https://github.com/facebookresearch/faiss/wiki/Index-IO,-cloning-and-hyper-parameter-tuning">setting
+ * relevant parameters</a>.
+ *
+ * <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>
+ *
+ * <p>Note: Set the {@code $OMP_NUM_THREADS} environment variable to control <a
+ * href="https://github.com/facebookresearch/faiss/wiki/Threads-and-asynchronous-calls">internal
+ * threading</a>.
+ *
+ * <p>TODO: There is no guarantee of backwards compatibility!
+ *
+ * @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;
+
+  /**
+   * Constructs an HNSW-based format using default {@code maxConn}={@value
+   * org.apache.lucene.util.hnsw.HnswGraphBuilder#DEFAULT_MAX_CONN} and {@code beamWidth}={@value
+   * org.apache.lucene.util.hnsw.HnswGraphBuilder#DEFAULT_BEAM_WIDTH}.
+   */
+  public FaissKnnVectorsFormat() {
+    this(
+        String.format(Locale.ROOT, "IDMap,HNSW%d", DEFAULT_MAX_CONN),
+        String.format(Locale.ROOT, "efConstruction=%d", DEFAULT_BEAM_WIDTH));
+  }
+
+  /**
+   * Constructs a format using the specified index factory string and index parameters (see class
+   * docs for more information).
+   *
+   * @param description the index factory string to initialize Faiss indexes.
+   * @param indexParams the index params to set on Faiss indexes.
+   */
+  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);
+  }
+}