You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ka...@apache.org on 2026/05/26 14:13:15 UTC
(lucene) branch main updated: FaissKnnVectorsFormat COSINE similarity support (#16065)
This is an automated email from the ASF dual-hosted git repository.
kaivalnp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new 366dc3c0396 FaissKnnVectorsFormat COSINE similarity support (#16065)
366dc3c0396 is described below
commit 366dc3c0396bba2ba634670c1592926c75bcf8ae
Author: Prithvi S <pr...@gmail.com>
AuthorDate: Tue May 26 19:43:08 2026 +0530
FaissKnnVectorsFormat COSINE similarity support (#16065)
Query and document vectors are normalized inside Faiss, and indexed using `METRIC_INNER_PRODUCT` similarity.
---
lucene/CHANGES.txt | 2 +
.../codecs/faiss/FaissKnnVectorsReader.java | 5 +-
.../lucene/sandbox/codecs/faiss/FaissLibrary.java | 2 +-
.../codecs/faiss/FaissLibraryNativeImpl.java | 58 +++++++++++++---------
.../sandbox/codecs/faiss/FaissNativeWrapper.java | 21 ++++++--
.../codecs/faiss/TestFaissKnnVectorsFormat.java | 5 +-
6 files changed, 62 insertions(+), 31 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 6859547f32f..75ef7e8db29 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -107,6 +107,8 @@ New Features
* GITHUB#14178: Add a Faiss-based vector format in the sandbox module. (Kaival Parikh)
+* GITHUB#16065: Add COSINE similarity support to FaissKnnVectorsFormat. (Prithvi S)
+
* GITHUB#15508: Use native vectorization in Lucene. (Ankur Goel, Shubham Chaudhary, Dawid Weiss)
* GITHUB#15818: Add BM25 k3 query-term frequency saturation to BM25Similarity. (Sagar Upadhyaya)
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissKnnVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissKnnVectorsReader.java
index 02275e5b596..1f791299091 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissKnnVectorsReader.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissKnnVectorsReader.java
@@ -119,7 +119,10 @@ final class FaissKnnVectorsReader extends KnnVectorsReader {
throw new CorruptIndexException("Duplicate field: " + fieldMeta.name, meta);
}
IndexInput indexInput = data.slice(fieldMeta.name, fieldMeta.offset, fieldMeta.length);
- indexMap.put(fieldMeta.name, FaissLibrary.INSTANCE.readIndex(indexInput));
+ FieldInfo fi = state.fieldInfos.fieldInfo(fieldMeta.name);
+ indexMap.put(
+ fieldMeta.name,
+ FaissLibrary.INSTANCE.readIndex(indexInput, fi.getVectorSimilarityFunction()));
}
} catch (Throwable t) {
IOUtils.closeWhileSuppressingExceptions(t, this);
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibrary.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibrary.java
index ad0715bcb33..7bdc393e681 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibrary.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibrary.java
@@ -54,5 +54,5 @@ interface FaissLibrary {
FloatVectorValues floatVectorValues,
IntToIntFunction oldToNewDocId);
- Index readIndex(IndexInput input);
+ Index readIndex(IndexInput input, VectorSimilarityFunction function);
}
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibraryNativeImpl.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibraryNativeImpl.java
index 08ec8264131..2ec3b564af6 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibraryNativeImpl.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissLibraryNativeImpl.java
@@ -20,6 +20,7 @@ import static java.lang.foreign.ValueLayout.ADDRESS;
import static java.lang.foreign.ValueLayout.JAVA_BYTE;
import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
import static java.lang.foreign.ValueLayout.JAVA_LONG;
+import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
import static org.apache.lucene.sandbox.codecs.faiss.FaissNativeWrapper.Exception.handleException;
@@ -36,7 +37,6 @@ import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.nio.ByteOrder;
import java.util.Map;
-import java.util.stream.Collectors;
import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.KnnVectorValues;
import org.apache.lucene.index.VectorSimilarityFunction;
@@ -140,7 +140,11 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
private static final Map<VectorSimilarityFunction, Integer> FUNCTION_TO_METRIC =
Map.of(
// Mapped from faiss/MetricType.h
+ // COSINE uses METRIC_INNER_PRODUCT with normalized vectors.
+ // See
+ // https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-index-vectors-for-cosine-distance
DOT_PRODUCT, 0,
+ COSINE, 0,
EUCLIDEAN, 1);
private static int functionToMetric(VectorSimilarityFunction function) {
@@ -151,19 +155,6 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
return metric;
}
- // Invert FUNCTION_TO_METRIC
- private static final Map<Integer, VectorSimilarityFunction> METRIC_TO_FUNCTION =
- FUNCTION_TO_METRIC.entrySet().stream()
- .collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
-
- private static VectorSimilarityFunction metricToFunction(int metric) {
- VectorSimilarityFunction function = METRIC_TO_FUNCTION.get(metric);
- if (function == null) {
- throw new UnsupportedOperationException("Metric not supported: " + metric);
- }
- return function;
- }
-
@Override
public FaissLibrary.Index createIndex(
String description,
@@ -220,6 +211,13 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
docsOffset += perDocByteSize;
}
+ // NOTE: Faiss does not have a native cosine metric, so we L2-normalize vectors and use
+ // METRIC_INNER_PRODUCT. This adds overhead at indexing and query time. If vectors are
+ // already normalized and raw vector retrieval is not needed, DOT_PRODUCT is preferred.
+ if (function == COSINE) {
+ wrapper.faiss_fvec_renorm_L2(dimension, size, docs);
+ }
+
// Train index
int isTrained = wrapper.faiss_Index_is_trained(indexPointer);
if (isTrained == 0) {
@@ -229,7 +227,7 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
// Add docs to index
handleException(wrapper.faiss_Index_add_with_ids(indexPointer, size, docs, ids));
- return new Index(indexPointer);
+ return new Index(indexPointer, function);
} catch (IOException e) {
throw new UncheckedIOException(e);
@@ -241,7 +239,7 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
private static final int FAISS_IO_FLAG_READ_ONLY = 2;
@Override
- public FaissLibrary.Index readIndex(IndexInput input) {
+ public FaissLibrary.Index readIndex(IndexInput input, VectorSimilarityFunction function) {
try (Arena temp = Arena.ofConfined()) {
MethodHandle readerHandle = READ_BYTES_HANDLE.bindTo(input);
MemorySegment readerStub =
@@ -262,7 +260,7 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
customIOReaderPointer, FAISS_IO_FLAG_MMAP | FAISS_IO_FLAG_READ_ONLY, pointer));
MemorySegment indexPointer = pointer.get(ADDRESS, 0);
- return new Index(indexPointer);
+ return new Index(indexPointer, function);
}
}
@@ -275,31 +273,34 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
private final Arena arena;
private final MemorySegment indexPointer;
private final FloatToFloatFunction scaler;
+ private final VectorSimilarityFunction function;
+ private final int dimension;
private boolean closed;
- private Index(MemorySegment indexPointer) {
+ private Index(MemorySegment indexPointer, VectorSimilarityFunction function) {
this.arena = Arena.ofShared();
this.indexPointer =
indexPointer
// Ensure timely cleanup
.reinterpret(arena, wrapper::faiss_Index_free);
-
- // Get underlying function
- int metricType = wrapper.faiss_Index_metric_type(indexPointer);
- VectorSimilarityFunction function = metricToFunction(metricType);
+ this.function = function;
+ this.dimension = wrapper.faiss_Index_d(indexPointer);
// Scale Faiss distances to Lucene scores, see VectorSimilarityFunction.java
this.scaler =
switch (function) {
- case DOT_PRODUCT ->
+ case DOT_PRODUCT, COSINE ->
// distance in Faiss === dotProduct in Lucene
+ // For COSINE, vectors are normalized so inner product == cosine similarity
distance -> Math.max((1 + distance) / 2, 0);
case EUCLIDEAN ->
// distance in Faiss === squareDistance in Lucene
distance -> 1 / (1 + distance);
- case COSINE, MAXIMUM_INNER_PRODUCT -> throw new AssertionError("Should not reach here");
+ case MAXIMUM_INNER_PRODUCT ->
+ throw new UnsupportedOperationException(
+ "Similarity function not supported: " + function);
};
this.closed = false;
@@ -326,6 +327,10 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
// Allocate queries in native memory
MemorySegment queries = temp.allocateFrom(JAVA_FLOAT, query);
+ // See indexing-time comment about COSINE overhead.
+ if (function == COSINE) {
+ wrapper.faiss_fvec_renorm_L2(dimension, 1, queries);
+ }
// Faiss knn search
int k = knnCollector.k();
@@ -373,6 +378,7 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
}
// Record hits
+ int resultCount = 0;
for (int i = 0; i < k; i++) {
int id = (int) idsPointer.getAtIndex(JAVA_LONG, i);
@@ -382,9 +388,13 @@ final class FaissLibraryNativeImpl implements FaissLibrary {
}
// Collect result
+ resultCount++;
float distance = distancesPointer.getAtIndex(JAVA_FLOAT, i);
knnCollector.collect(id, scaler.scale(distance));
}
+ if (resultCount > 0) {
+ knnCollector.incVisitedCount(resultCount);
+ }
} catch (IOException e) {
throw new RuntimeException(e);
}
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissNativeWrapper.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissNativeWrapper.java
index 575ebf95322..2e76d3cd5c2 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissNativeWrapper.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/FaissNativeWrapper.java
@@ -197,12 +197,25 @@ final class FaissNativeWrapper {
}
}
- private final MethodHandle faiss_Index_metric_type$MH =
- getHandle("faiss_Index_metric_type", FunctionDescriptor.of(JAVA_INT, ADDRESS));
+ private final MethodHandle faiss_Index_d$MH =
+ getHandle("faiss_Index_d", FunctionDescriptor.of(JAVA_INT, ADDRESS));
- int faiss_Index_metric_type(MemorySegment indexPointer) {
+ int faiss_Index_d(MemorySegment indexPointer) {
try {
- return (int) faiss_Index_metric_type$MH.invokeExact(indexPointer);
+ return (int) faiss_Index_d$MH.invokeExact(indexPointer);
+ } catch (RuntimeException | Error e) {
+ throw e;
+ } catch (Throwable t) {
+ throw new AssertionError(t);
+ }
+ }
+
+ private final MethodHandle faiss_fvec_renorm_L2$MH =
+ getHandle("faiss_fvec_renorm_L2", FunctionDescriptor.ofVoid(JAVA_LONG, JAVA_LONG, ADDRESS));
+
+ void faiss_fvec_renorm_L2(long dimension, long numVectors, MemorySegment vectors) {
+ try {
+ faiss_fvec_renorm_L2$MH.invokeExact(dimension, numVectors, vectors);
} catch (RuntimeException | Error e) {
throw e;
} catch (Throwable t) {
diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/faiss/TestFaissKnnVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/faiss/TestFaissKnnVectorsFormat.java
index bd4bcf9c70f..ec38af32454 100644
--- a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/faiss/TestFaissKnnVectorsFormat.java
+++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/faiss/TestFaissKnnVectorsFormat.java
@@ -17,6 +17,7 @@
package org.apache.lucene.sandbox.codecs.faiss;
import static org.apache.lucene.index.VectorEncoding.FLOAT32;
+import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
@@ -45,7 +46,9 @@ public class TestFaissKnnVectorsFormat extends BaseKnnVectorsFormatTestCase {
private static final String FAISS_RUN_TESTS = "tests.faiss.run";
private static final VectorEncoding[] SUPPORTED_ENCODINGS = {FLOAT32};
- private static final VectorSimilarityFunction[] SUPPORTED_FUNCTIONS = {DOT_PRODUCT, EUCLIDEAN};
+ private static final VectorSimilarityFunction[] SUPPORTED_FUNCTIONS = {
+ DOT_PRODUCT, EUCLIDEAN, COSINE
+ };
@BeforeClass
public static void maybeSuppress() {