-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Backport Faiss-based vector format to 10.x #14843
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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: [ '21' ] | ||
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} |
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Think we should not duplicate that file. Either refer to it from the other module (seems most non-invasive) and especially compatible with main branch or move the code that accesses foreign APIs to main branch. In ideal world only the LibFaissC should access native APIs and add abstractions for all code to get rid of MemorySegment. Then we could place a provider to get the LibFaissC fclass with tha API abstractions from the core's org.apache.lucene.internal package (which is only exported to sandbox but not to public) and the KNN vector codec would reside in the sandbox package and only call methods which use public Java classes. I would prefer also on main branch to put all restricted code in LibFaissC only. This would also solve @rmuir's complaints about duoing "restricted" chnages to MemorySegments. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Makes sense, I've addressed these in the refactor! |
Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,88 @@ | ||||||||||||||||||||||||||||||||||||||||||||||
/* | ||||||||||||||||||||||||||||||||||||||||||||||
* 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.lang.invoke.MethodHandle; | ||||||||||||||||||||||||||||||||||||||||||||||
import java.lang.invoke.MethodHandles; | ||||||||||||||||||||||||||||||||||||||||||||||
import java.lang.invoke.MethodType; | ||||||||||||||||||||||||||||||||||||||||||||||
import java.util.Arrays; | ||||||||||||||||||||||||||||||||||||||||||||||
import java.util.stream.Collectors; | ||||||||||||||||||||||||||||||||||||||||||||||
import org.apache.lucene.codecs.KnnVectorsFormat; | ||||||||||||||||||||||||||||||||||||||||||||||
import org.apache.lucene.codecs.KnnVectorsReader; | ||||||||||||||||||||||||||||||||||||||||||||||
import org.apache.lucene.codecs.KnnVectorsWriter; | ||||||||||||||||||||||||||||||||||||||||||||||
import org.apache.lucene.index.SegmentReadState; | ||||||||||||||||||||||||||||||||||||||||||||||
import org.apache.lucene.index.SegmentWriteState; | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
/** | ||||||||||||||||||||||||||||||||||||||||||||||
* Provides a Faiss-based vector format, see corresponding classes in {@code java21/} for docs! | ||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. the docs of private classes in java21 are not visible in javadocs. So need to stay here. |
||||||||||||||||||||||||||||||||||||||||||||||
* | ||||||||||||||||||||||||||||||||||||||||||||||
* @lucene.experimental | ||||||||||||||||||||||||||||||||||||||||||||||
*/ | ||||||||||||||||||||||||||||||||||||||||||||||
public class FaissKnnVectorsFormatProvider extends KnnVectorsFormat { | ||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this class should be named Instead the FaissKnnVectorsFormat in the java21 part should be package private (it is now public) and have a different name, e.g. 'FaissKnnVectorsFormatImpl' There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Makes sense, I've addressed this in the refactor on |
||||||||||||||||||||||||||||||||||||||||||||||
private final KnnVectorsFormat delegate; | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
public FaissKnnVectorsFormatProvider() { | ||||||||||||||||||||||||||||||||||||||||||||||
this(lookup()); | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
public FaissKnnVectorsFormatProvider(String description, String indexParams) { | ||||||||||||||||||||||||||||||||||||||||||||||
this(lookup(description, indexParams)); | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
private FaissKnnVectorsFormatProvider(KnnVectorsFormat delegate) { | ||||||||||||||||||||||||||||||||||||||||||||||
super(delegate.getName()); | ||||||||||||||||||||||||||||||||||||||||||||||
this.delegate = delegate; | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
private static KnnVectorsFormat lookup(Object... args) { | ||||||||||||||||||||||||||||||||||||||||||||||
try { | ||||||||||||||||||||||||||||||||||||||||||||||
MethodHandles.Lookup lookup = MethodHandles.lookup(); | ||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. we have a private lookup here, so we can lookup a package private impl. |
||||||||||||||||||||||||||||||||||||||||||||||
Class<?> cls = | ||||||||||||||||||||||||||||||||||||||||||||||
lookup.findClass("org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat"); | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
MethodType type = | ||||||||||||||||||||||||||||||||||||||||||||||
MethodType.methodType( | ||||||||||||||||||||||||||||||||||||||||||||||
void.class, | ||||||||||||||||||||||||||||||||||||||||||||||
Arrays.stream(args).map(Object::getClass).collect(Collectors.toUnmodifiableList())); | ||||||||||||||||||||||||||||||||||||||||||||||
MethodHandle constr = lookup.findConstructor(cls, type); | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
return (KnnVectorsFormat) constr.invokeWithArguments(args); | ||||||||||||||||||||||||||||||||||||||||||||||
} catch (ClassNotFoundException e) { | ||||||||||||||||||||||||||||||||||||||||||||||
throw new LinkageError("FaissKnnVectorsFormat is missing from JAR file", e); | ||||||||||||||||||||||||||||||||||||||||||||||
} catch (IllegalAccessException | NoSuchMethodException e) { | ||||||||||||||||||||||||||||||||||||||||||||||
throw new LinkageError("FaissKnnVectorsFormat is missing correctly typed constructor", e); | ||||||||||||||||||||||||||||||||||||||||||||||
} catch (Throwable t) { | ||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please check how it is done here: lucene/lucene/core/src/java/org/apache/lucene/store/MMapDirectory.java Lines 399 to 420 in e610d0e
As said before carefully differentiate between Exception types. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In the original code we have a separate try/ctach around the invoke call, this is unsafe as it may catch too many exceptions from unrelated places. So please add a try-catch around only the call. An please use a correctly typed method handle..... This is too lazy as not statically typed. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I wozuld only code the multi-arg ctor here statically and delegate from no-arg ctor of this class by calling There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Addressed in the refactor! |
||||||||||||||||||||||||||||||||||||||||||||||
throw new RuntimeException(t); | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
@Override | ||||||||||||||||||||||||||||||||||||||||||||||
public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { | ||||||||||||||||||||||||||||||||||||||||||||||
return delegate.fieldsWriter(state); | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
@Override | ||||||||||||||||||||||||||||||||||||||||||||||
public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { | ||||||||||||||||||||||||||||||||||||||||||||||
return delegate.fieldsReader(state); | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||
@Override | ||||||||||||||||||||||||||||||||||||||||||||||
public int getMaxDimensions(String fieldName) { | ||||||||||||||||||||||||||||||||||||||||||||||
return delegate.getMaxDimensions(fieldName); | ||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,22 @@ | ||
/* | ||
* 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. | ||
*/ | ||
/** | ||
* Provides a Faiss-based vector format, see corresponding classes in {@code java21/} for docs! | ||
* | ||
* @lucene.experimental | ||
*/ | ||
package org.apache.lucene.sandbox.codecs.faiss; |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. make package-private There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. +rename (see above) |
||
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); | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
unfortunately this causes also the vector API be replicated in the sandbox module.... I don't like that. It also adds useless JAR files again to the code.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Makes sense, I'll try to reuse the
apijar
fromlucene.core