Skip to content

Enable Faiss-based vector format to index larger number of vectors in a single segment #14847

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.

Sign up for GitHub

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

Merged
mikemccand merged 7 commits into from
Jul 5, 2025
Merged

Enable Faiss-based vector format to index larger number of vectors in a single segment #14847

Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d4b09a1
Enable Faiss-based vector format to index larger number of vectors in…
kaivalnp Jun 25, 2025
c7dbca2
Add test
kaivalnp Jun 26, 2025
f70051d
Mark test as "monster"
kaivalnp Jun 26, 2025
5cf3c60
Allocate filtered doc bits without ByteBuffer
kaivalnp Jul 3, 2025
4c9b052
Merge branch 'main' into faiss-larger-vectors
kaivalnp Jul 3, 2025
7b3cdb5
Merge branch 'main' into faiss-larger-vectors
kaivalnp Jul 5, 2025
2de7c99
Add CHANGES.txt entry
kaivalnp Jul 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 19 additions & 33 deletions lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/faiss/LibFaissC.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package org.apache.lucene.sandbox.codecs.faiss;

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_INT;
import static java.lang.foreign.ValueLayout.JAVA_LONG;
Expand All @@ -32,8 +33,6 @@
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.LongBuffer;
import java.util.Arrays;
import java.util.Locale;
import org.apache.lucene.index.FloatVectorValues;
Expand Down Expand Up @@ -221,16 +220,22 @@ public static MemorySegment createIndex(

// Allocate docs in native memory
MemorySegment docs = temp.allocate(JAVA_FLOAT, (long) size * dimension);
FloatBuffer docsBuffer = docs.asByteBuffer().order(ByteOrder.nativeOrder()).asFloatBuffer();
long docsOffset = 0;
long perDocByteSize = dimension * JAVA_FLOAT.byteSize();

// Allocate ids in native memory
MemorySegment ids = temp.allocate(JAVA_LONG, size);
LongBuffer idsBuffer = ids.asByteBuffer().order(ByteOrder.nativeOrder()).asLongBuffer();
int idsIndex = 0;

KnnVectorValues.DocIndexIterator iterator = floatVectorValues.iterator();
for (int i = iterator.nextDoc(); i != NO_MORE_DOCS; i = iterator.nextDoc()) {
idsBuffer.put(oldToNewDocId.apply(i));
docsBuffer.put(floatVectorValues.vectorValue(iterator.index()));
int id = oldToNewDocId.apply(i);
ids.setAtIndex(JAVA_LONG, idsIndex, id);
idsIndex++;

float[] vector = floatVectorValues.vectorValue(iterator.index());
MemorySegment.copy(vector, 0, docs, JAVA_FLOAT, docsOffset, vector.length);
docsOffset += perDocByteSize;
}

// Train index
Expand All @@ -254,18 +259,12 @@ private static long writeBytes(
inputPointer = inputPointer.reinterpret(size);

if (size <= BUFFER_SIZE) { // simple case, avoid buffering
byte[] bytes = new byte[(int) size];
inputPointer.asSlice(0, size).asByteBuffer().order(ByteOrder.nativeOrder()).get(bytes);
output.writeBytes(bytes, bytes.length);
output.writeBytes(inputPointer.toArray(JAVA_BYTE), (int) size);
} else { // copy buffered number of bytes repeatedly
byte[] bytes = new byte[BUFFER_SIZE];
for (long offset = 0; offset < size; offset += BUFFER_SIZE) {
int length = (int) Math.min(size - offset, BUFFER_SIZE);
inputPointer
.asSlice(offset, length)
.asByteBuffer()
.order(ByteOrder.nativeOrder())
.get(bytes, 0, length);
MemorySegment.copy(inputPointer, JAVA_BYTE, offset, bytes, 0, length);
output.writeBytes(bytes, length);
}
}
Expand All @@ -282,21 +281,13 @@ private static long readBytes(
if (size <= BUFFER_SIZE) { // simple case, avoid buffering
byte[] bytes = new byte[(int) size];
input.readBytes(bytes, 0, bytes.length);
outputPointer
.asSlice(0, bytes.length)
.asByteBuffer()
.order(ByteOrder.nativeOrder())
.put(bytes);
MemorySegment.copy(bytes, 0, outputPointer, JAVA_BYTE, 0, bytes.length);
} else { // copy buffered number of bytes repeatedly
byte[] bytes = new byte[BUFFER_SIZE];
for (long offset = 0; offset < size; offset += BUFFER_SIZE) {
int length = (int) Math.min(size - offset, BUFFER_SIZE);
input.readBytes(bytes, 0, length);
outputPointer
.asSlice(offset, length)
.asByteBuffer()
.order(ByteOrder.nativeOrder())
.put(bytes, 0, length);
MemorySegment.copy(bytes, 0, outputPointer, JAVA_BYTE, offset, length);
}
}
return numItems;
Expand Down Expand Up @@ -411,8 +402,7 @@ public static void indexSearch(
};

// Allocate queries in native memory
MemorySegment queries = temp.allocate(JAVA_FLOAT, query.length);
queries.asByteBuffer().order(ByteOrder.nativeOrder()).asFloatBuffer().put(query);
MemorySegment queries = temp.allocateFrom(JAVA_FLOAT, query);

// Faiss knn search
int k = knnCollector.k();
Expand Down Expand Up @@ -458,13 +448,9 @@ public static void indexSearch(
idsPointer);
}

// Retrieve scores
float[] distances = new float[k];
distancesPointer.asByteBuffer().order(ByteOrder.nativeOrder()).asFloatBuffer().get(distances);

// Retrieve ids
long[] ids = new long[k];
idsPointer.asByteBuffer().order(ByteOrder.nativeOrder()).asLongBuffer().get(ids);
// Retrieve scores and ids
float[] distances = distancesPointer.toArray(JAVA_FLOAT);
long[] ids = idsPointer.toArray(JAVA_LONG);

// Record hits
for (int i = 0; i < k; i++) {
Expand Down
Loading