diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index c7ea7a1..4a6bf41 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -36,3 +36,7 @@ sed_runner "s/VERSION=\".*\"/VERSION=\"${NEXT_FULL_JAVA_TAG}\"/g" build.sh sed_runner "/.*/s//${NEXT_FULL_JAVA_TAG}<\/version>/g" pom.xml sed_runner "s| CuVS [[:digit:]]\{2\}\.[[:digit:]]\{2\} | CuVS ${NEXT_SHORT_TAG} |g" README.md + +for FILE in dependencies.yaml conda/environments/*.yaml; do + sed_runner "s/libcuvs==.*/libcuvs==${NEXT_SHORT_TAG}.*/g" "${FILE}" +done diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 0b31020..05ec85b 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-129_arch-aarch64 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index a8f0444..f0a2538 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-129_arch-x86_64 diff --git a/conda/environments/all_cuda-130_arch-aarch64.yaml b/conda/environments/all_cuda-130_arch-aarch64.yaml index 50e93bd..6638808 100644 --- a/conda/environments/all_cuda-130_arch-aarch64.yaml +++ b/conda/environments/all_cuda-130_arch-aarch64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-130_arch-aarch64 diff --git a/conda/environments/all_cuda-130_arch-x86_64.yaml b/conda/environments/all_cuda-130_arch-x86_64.yaml index c1e7ab3..bb8a746 100644 --- a/conda/environments/all_cuda-130_arch-x86_64.yaml +++ b/conda/environments/all_cuda-130_arch-x86_64.yaml @@ -12,7 +12,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs +- libcuvs==25.10.* - maven - openjdk=22.* name: all_cuda-130_arch-x86_64 diff --git a/dependencies.yaml b/dependencies.yaml index ee06135..b9eb3fc 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -67,7 +67,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev - - libcuvs + - libcuvs==25.10.* java: common: - output_types: conda diff --git a/pom.xml b/pom.xml index 6400cc4..dfd9497 100644 --- a/pom.xml +++ b/pom.xml @@ -16,7 +16,7 @@ - cuvs-java + searchscale-maven SearchScale Maven https://maven.searchscale.com/snapshots @@ -40,36 +40,17 @@ 10.2.0 test - - com.opencsv - opencsv - 5.3 - commons-io commons-io - 2.15.1 - - - com.github.fommil - jniloader - 1.1 - - - com.fasterxml.jackson.core - jackson-databind - 2.17.0 - - - com.fasterxml.jackson.dataformat - jackson-dataformat-csv - 2.17.0 + 2.18.0 + test com.nvidia.cuvs cuvs-java - - 25.10.0-2c0e1-SNAPSHOT + + 25.10.0 @@ -104,6 +85,25 @@ + + org.apache.maven.plugins + maven-assembly-plugin + 3.6.0 + + + jar-with-dependencies + + + + + make-assembly + package + + single + + + + diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java similarity index 50% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index ec86176..4027d3a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -16,41 +16,57 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.FilterCodec; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene101.Lucene101Codec; -/** CuVS based codec for GPU based vector search */ -public class CuVSCodec extends FilterCodec { +/** CuVS based codec for GPU based vector search + * + * @apiNote cuVS serialization formats are in experimental phase and hence backward compatibility cannot be guaranteed. + * + * */ +public class CuVS2510GPUSearchCodec extends FilterCodec { + + private static final Logger log = Logger.getLogger(CuVS2510GPUSearchCodec.class.getName()); + private static final String NAME = "CuVS2510GPUSearchCodec"; - public CuVSCodec() { - this("CuVSCodec", new Lucene101Codec()); + private static final int DEFAULT_CUVS_WRITER_THREADS = 1; + private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + private static final int DEFAULT_GRAPH_DEGREE = 64; + private static final int DEFAULT_HNSW_LAYERS = 1; + private static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + + private KnnVectorsFormat format; + + public CuVS2510GPUSearchCodec() { + this(NAME, new Lucene101Codec()); } - public CuVSCodec(String name, Codec delegate) { + public CuVS2510GPUSearchCodec(String name, Codec delegate) { super(name, delegate); - KnnVectorsFormat format; try { - // TODO: Remove this hard coded values. - format = new CuVSVectorsFormat(1, 128, 64, IndexType.CAGRA); + format = + new CuVS2510GPUVectorsFormat( + DEFAULT_CUVS_WRITER_THREADS, + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + DEFAULT_HNSW_LAYERS, + DEFAULT_INDEX_TYPE); setKnnFormat(format); } catch (LibraryException ex) { - Logger log = Logger.getLogger(CuVSCodec.class.getName()); log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); } } - KnnVectorsFormat knnFormat = null; - @Override public KnnVectorsFormat knnVectorsFormat() { - return knnFormat; + return format; } public void setKnnFormat(KnnVectorsFormat format) { - this.knnFormat = format; + this.format = format; } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java similarity index 56% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 1f7e439..bccff5a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -15,12 +15,15 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; + import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.logging.Logger; import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; @@ -28,23 +31,24 @@ import org.apache.lucene.index.SegmentWriteState; /** CuVS based KnnVectorsFormat for GPU acceleration */ -public class CuVSVectorsFormat extends KnnVectorsFormat { +public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { - private static final Logger LOG = Logger.getLogger(CuVSVectorsFormat.class.getName()); + static final Logger log = Logger.getLogger(CuVS2510GPUVectorsFormat.class.getName()); // TODO: fix Lucene version in name, to the final targeted release, if any static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta"; - static final String CUVS_META_CODEC_EXT = "vemc"; // ""cagmf"; + static final String CUVS_META_CODEC_EXT = "vemc"; static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex"; static final String CUVS_INDEX_EXT = "vcag"; static final int VERSION_START = 0; static final int VERSION_CURRENT = VERSION_START; - public static final int DEFAULT_WRITER_THREADS = 32; - public static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; - public static final int DEFAULT_GRAPH_DEGREE = 64; - public static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + static final int DEFAULT_WRITER_THREADS = 32; + static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + static final int DEFAULT_GRAPH_DEGREE = 64; + static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + static final int HNSW_GRAPH_LAYERS = 1; static CuVSResources resources = cuVSResourcesOrNull(); @@ -56,80 +60,61 @@ public class CuVSVectorsFormat extends KnnVectorsFormat { final int cuvsWriterThreads; final int intGraphDegree; final int graphDegree; - final CuVSVectorsWriter.IndexType indexType; // the index type to build, when writing + final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion + final CuVS2510GPUVectorsWriter.IndexType indexType; // the index type to build, when writing /** - * Creates a CuVSVectorsFormat, with default values. + * Creates a CuVS2510GPUVectorsFormat, with default values. * * @throws LibraryException if the native library fails to load */ - public CuVSVectorsFormat() { + public CuVS2510GPUVectorsFormat() { this( DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, + HNSW_GRAPH_LAYERS, DEFAULT_INDEX_TYPE); } /** - * Creates a CuVSVectorsFormat, with the given threads, graph degree, etc. + * Creates a CuVS2510GPUVectorsFormat, with the given threads, graph degree, etc. * * @throws LibraryException if the native library fails to load */ - public CuVSVectorsFormat( - int cuvsWriterThreads, int intGraphDegree, int graphDegree, IndexType indexType) { - super("CuVSVectorsFormat"); + public CuVS2510GPUVectorsFormat( + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + IndexType indexType) { + super("CuVS2510GPUVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.hnswLayers = hnswLayers; this.indexType = indexType; } - private static CuVSResources cuVSResourcesOrNull() { - try { - System.loadLibrary( - "cudart"); // nocommit: this is here so as to pass CI, should goto cuvs-java - } catch (UnsatisfiedLinkError e) { - LOG.warning("Could not load CUDA runtime library: " + e.getMessage()); - } - try { - resources = CuVSResources.create(); - return resources; - } catch (UnsupportedOperationException uoe) { - LOG.warning("cuvs is not supported on this platform or java version: " + uoe.getMessage()); - } catch (Throwable t) { - if (t instanceof ExceptionInInitializerError ex) { - t = ex.getCause(); - } - LOG.warning("Exception occurred during creation of cuvs resources. " + t); - } - return null; - } - - /** Tells whether the platform supports cuvs. */ - public static boolean supported() { - return resources != null; - } - - private static void checkSupported() { - if (!supported()) { - throw new UnsupportedOperationException(); - } - } - @Override - public CuVSVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + public CuVS2510GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { checkSupported(); var flatWriter = flatVectorsFormat.fieldsWriter(state); - return new CuVSVectorsWriter( - state, cuvsWriterThreads, intGraphDegree, graphDegree, indexType, resources, flatWriter); + return new CuVS2510GPUVectorsWriter( + state, + cuvsWriterThreads, + intGraphDegree, + graphDegree, + hnswLayers, + indexType, + resources, + flatWriter); } @Override - public CuVSVectorsReader fieldsReader(SegmentReadState state) throws IOException { + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { checkSupported(); - var flatReader = flatVectorsFormat.fieldsReader(state); - return new CuVSVectorsReader(state, resources, flatReader); + return new CuVS2510GPUVectorsReader(state, resources, flatVectorsFormat.fieldsReader(state)); } @Override @@ -139,12 +124,24 @@ public int getMaxDimensions(String fieldName) { @Override public String toString() { - StringBuilder sb = new StringBuilder("CuVSVectorsFormat("); - sb.append("cuvsWriterThreads=").append(cuvsWriterThreads); + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(cuvsWriterThreads=").append(cuvsWriterThreads); sb.append("intGraphDegree=").append(intGraphDegree); sb.append("graphDegree=").append(graphDegree); + sb.append("hnswLayers=").append(hnswLayers); sb.append("resources=").append(resources); sb.append(")"); return sb.toString(); } + + /** Tells whether the platform supports cuVS. */ + public static boolean supported() { + return resources != null; + } + + public static void checkSupported() { + if (!supported()) { + throw new UnsupportedOperationException(); + } + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java similarity index 93% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 4118a0a..a9fe61a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -15,12 +15,12 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_CURRENT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_START; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_START; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.BruteForceQuery; @@ -62,19 +62,19 @@ import org.apache.lucene.util.hnsw.IntToIntFunction; /** KnnVectorsReader instance associated with CuVS format */ -public class CuVSVectorsReader extends KnnVectorsReader { +public class CuVS2510GPUVectorsReader extends KnnVectorsReader { @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(CuVSVectorsReader.class.getName()); + private static final Logger log = Logger.getLogger(CuVS2510GPUVectorsReader.class.getName()); private final CuVSResources resources; private final FlatVectorsReader flatVectorsReader; // for reading the raw vectors private final FieldInfos fieldInfos; private final IntObjectHashMap fields; - private final IntObjectHashMap cuvsIndices; + private final IntObjectHashMap cuvsIndices; private final IndexInput cuvsIndexInput; - public CuVSVectorsReader( + public CuVS2510GPUVectorsReader( SegmentReadState state, CuVSResources resources, FlatVectorsReader flatReader) throws IOException { this.resources = resources; @@ -226,8 +226,8 @@ private FieldEntry getFieldEntry(String field, VectorEncoding expectedEncoding) return fieldEntry; } - private IntObjectHashMap loadCuVSIndices() throws IOException { - var indices = new IntObjectHashMap(); + private IntObjectHashMap loadCuVSIndices() throws IOException { + var indices = new IntObjectHashMap(); for (var e : fields) { var fieldEntry = e.value; int fieldNumber = e.key; @@ -237,7 +237,7 @@ private IntObjectHashMap loadCuVSIndices() throws IOException { return indices; } - private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { + private GPUIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { CagraIndex cagraIndex = null; BruteForceIndex bruteForceIndex = null; HnswIndex hnswIndex = null; @@ -273,7 +273,7 @@ private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { } catch (Throwable t) { Utils.handleThrowable(t); } - return new CuVSIndex(cagraIndex, bruteForceIndex, hnswIndex); + return new GPUIndex(cagraIndex, bruteForceIndex, hnswIndex); } @Override @@ -335,7 +335,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits var fieldNumber = fieldInfos.fieldInfo(field).number; - CuVSIndex cuvsIndex = cuvsIndices.get(fieldNumber); + GPUIndex cuvsIndex = cuvsIndices.get(fieldNumber); if (cuvsIndex == null) { throw new IllegalStateException("not index found for field:" + field); } @@ -476,7 +476,7 @@ public FieldInfos getFieldInfos() { return fieldInfos; } - public IntObjectHashMap getCuvsIndexes() { + public IntObjectHashMap getCuvsIndexes() { return cuvsIndices; } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java similarity index 78% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 7861ad1..69d66db 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -15,11 +15,11 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_INDEX_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.CUVS_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.CuVSVectorsFormat.VERSION_CURRENT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_INDEX_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.CUVS_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -36,11 +36,9 @@ import java.io.OutputStream; import java.nio.file.Files; import java.nio.file.Path; -import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.Objects; -import java.util.function.Supplier; import java.util.logging.Logger; import java.util.stream.IntStream; import org.apache.lucene.codecs.CodecUtil; @@ -69,15 +67,16 @@ * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into * GPU */ -public class CuVSVectorsWriter extends KnnVectorsWriter { +public class CuVS2510GPUVectorsWriter extends KnnVectorsWriter { - private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(CuVSVectorsWriter.class); + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(CuVS2510GPUVectorsWriter.class); @SuppressWarnings("unused") - private static final Logger log = Logger.getLogger(CuVSVectorsWriter.class.getName()); + private static final Logger log = Logger.getLogger(CuVS2510GPUVectorsWriter.class.getName()); /** The name of the CUVS component for the info-stream * */ - public static final String CUVS_COMPONENT = "CUVS"; + private static final String CUVS_COMPONENT = "CUVS"; // The minimum number of vectors in the dataset required before // we attempt to build a Cagra index @@ -91,19 +90,24 @@ public class CuVSVectorsWriter extends KnnVectorsWriter { private final IndexType indexType; private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors - private final List fields = new ArrayList<>(); - private final IndexOutput meta, cuvsIndex; + private final List fields = new ArrayList<>(); + private IndexOutput meta = null, cuvsIndex = null; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; private final InfoStream infoStream; private boolean finished; /** The CuVS index Type. */ public enum IndexType { + /** Builds a Cagra index. */ CAGRA(true, false, false), + /** Builds a Brute Force index. */ BRUTE_FORCE(false, true, false), + /** Builds an HSNW index - suitable for searching on CPU. */ HNSW(false, false, true), + /** Builds a Cagra and a Brute Force index. */ CAGRA_AND_BRUTE_FORCE(true, true, false); private final boolean cagra, bruteForce, hnsw; @@ -127,11 +131,12 @@ public boolean hnsw() { } } - public CuVSVectorsWriter( + public CuVS2510GPUVectorsWriter( SegmentWriteState state, int cuvsWriterThreads, int intGraphDegree, int graphDegree, + int hnswLayers, IndexType indexType, CuVSResources resources, FlatVectorsWriter flatVectorsWriter) @@ -153,6 +158,7 @@ public CuVSVectorsWriter( boolean success = false; try { + meta = state.directory.createOutput(metaFileName, state.context); cuvsIndex = state.directory.createOutput(cagraFileName, state.context); CodecUtil.writeIndexHeader( @@ -167,6 +173,7 @@ public CuVSVectorsWriter( VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + success = true; } finally { if (success == false) { @@ -184,7 +191,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); @SuppressWarnings("unchecked") var flatWriter = (FlatFieldVectorsWriter) writer; - var cuvsFieldWriter = new CuVSFieldWriter(fieldInfo, flatWriter); + var cuvsFieldWriter = new GPUFieldWriter(fieldInfo, flatWriter); fields.add(cuvsFieldWriter); return writer; } @@ -213,25 +220,89 @@ private CagraIndexParams cagraIndexParams(int size) { .build(); } - static long nanosToMillis(long nanos) { - return Duration.ofNanos(nanos).toMillis(); - } - private void info(String msg) { if (infoStream.isEnabled(CUVS_COMPONENT)) { infoStream.message(CUVS_COMPONENT, msg); } } + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo); + return; + } + long cagraIndexOffset, cagraIndexLength = 0L; + long bruteForceIndexOffset, bruteForceIndexLength = 0L; + long hnswIndexOffset, hnswIndexLength = 0L; + + // workaround for the minimum number of vectors for Cagra + IndexType indexType = + this.indexType.cagra() && vectors.size() < MIN_CAGRA_INDEX_SIZE + ? IndexType.BRUTE_FORCE + : this.indexType; + + try { + + cagraIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.cagra()) { + try { + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeCagraIndex(cagraIndexOutputStream, dataset); + } catch (Throwable t) { + Utils.handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + // workaround for cuVS issue + indexType = IndexType.BRUTE_FORCE; + } + cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; + } + + bruteForceIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.bruteForce()) { + var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeBruteForceIndex(bruteForceIndexOutputStream, dataset); + bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; + } + + hnswIndexOffset = cuvsIndex.getFilePointer(); + if (indexType.hnsw()) { + var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + if (vectors.size() > MIN_CAGRA_INDEX_SIZE) { + try { + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeHNSWIndex(hnswIndexOutputStream, dataset); + } catch (Throwable t) { + Utils.handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); + } + } + hnswIndexLength = cuvsIndex.getFilePointer() - hnswIndexOffset; + } + + // Only write meta for non-HNSW_LUCENE modes + writeMeta( + fieldInfo, + vectors.size(), + cagraIndexOffset, + cagraIndexLength, + bruteForceIndexOffset, + bruteForceIndexLength, + hnswIndexOffset, + hnswIndexLength); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { if (dataset.size() < 2) { throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); } CagraIndexParams params = cagraIndexParams((int) dataset.size()); long startTime = System.nanoTime(); - var index = + CagraIndex index = CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); - long elapsedMillis = nanosToMillis(System.nanoTime() - startTime); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("Cagra index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); Path tmpFile = Files.createTempFile(resources.tempDirectory(), "tmpindex", "cag"); index.serialize(os, tmpFile); @@ -241,13 +312,12 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { BruteForceIndexParams params = new BruteForceIndexParams.Builder() - .withNumWriterThreads(32) // TODO: Make this - // configurable later. + .withNumWriterThreads(32) // TODO: Make this configurable. .build(); long startTime = System.nanoTime(); var index = BruteForceIndex.newBuilder(resources).withIndexParams(params).withDataset(dataset).build(); - long elapsedMillis = nanosToMillis(System.nanoTime() - startTime); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("bf index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); index.serialize(os); index.close(); @@ -259,9 +329,9 @@ private void writeHNSWIndex(OutputStream os, CuVSMatrix dataset) throws Throwabl } CagraIndexParams indexParams = cagraIndexParams((int) dataset.size()); long startTime = System.nanoTime(); - var index = + CagraIndex index = CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(indexParams).build(); - long elapsedMillis = nanosToMillis(System.nanoTime() - startTime); + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); info("HNSW index created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); Path tmpFile = Files.createTempFile("tmpindex", "hnsw"); index.serializeToHNSW(os, tmpFile); @@ -280,96 +350,23 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { } } - private void writeField(CuVSFieldWriter fieldData) throws IOException { - // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors - // from the flat writer to the CuVSMatrix? - List vectors = fieldData.getVectors(); - writeFieldInternal( - fieldData.fieldInfo(), - () -> Utils.createFloatMatrix(vectors, fieldData.fieldInfo().getVectorDimension()), - vectors.size()); + private void writeField(GPUFieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); } - private void writeSortingField(CuVSFieldWriter fieldData, Sorter.DocMap sortMap) + private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); - // TODO: Loading all vectors into memory is inefficient. Is there a way to stream the vectors - // from the flat writer to the CuVSMatrix? + List sortedVectors = new ArrayList(); for (int i = 0; i < fieldData.getVectors().size(); i++) { sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); } - writeFieldInternal( - fieldData.fieldInfo(), - () -> Utils.createFloatMatrix(sortedVectors, fieldData.fieldInfo().getVectorDimension()), - sortedVectors.size()); - } - private void writeFieldInternal( - FieldInfo fieldInfo, Supplier datasetSupplier, int datasetSize) - throws IOException { - if (datasetSize == 0) { - writeEmpty(fieldInfo); - return; - } - long cagraIndexOffset, cagraIndexLength = 0L; - long bruteForceIndexOffset, bruteForceIndexLength = 0L; - long hnswIndexOffset, hnswIndexLength = 0L; - - // workaround for the minimum number of vectors for Cagra - IndexType indexType = - this.indexType.cagra() && datasetSize < MIN_CAGRA_INDEX_SIZE - ? IndexType.BRUTE_FORCE - : this.indexType; - - try { - cagraIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.cagra()) { - try { - var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - writeCagraIndex(cagraIndexOutputStream, datasetSupplier.get()); - } catch (Throwable t) { - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - // workaround for cuVS issue - indexType = IndexType.BRUTE_FORCE; - } - cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; - } - - bruteForceIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.bruteForce()) { - var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - writeBruteForceIndex(bruteForceIndexOutputStream, datasetSupplier.get()); - bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; - } - - hnswIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.hnsw()) { - var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - if (datasetSize > MIN_CAGRA_INDEX_SIZE) { - try { - writeHNSWIndex(hnswIndexOutputStream, datasetSupplier.get()); - } catch (Throwable t) { - handleThrowableWithIgnore(t, CANNOT_GENERATE_CAGRA); - } - } - hnswIndexLength = cuvsIndex.getFilePointer() - hnswIndexOffset; - } - - writeMeta( - fieldInfo, - (int) datasetSize, - cagraIndexOffset, - cagraIndexLength, - bruteForceIndexOffset, - bruteForceIndexLength, - hnswIndexOffset, - hnswIndexLength); - } catch (Throwable t) { - Utils.handleThrowable(t); - } + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); } private void writeEmpty(FieldInfo fieldInfo) throws IOException { @@ -417,13 +414,6 @@ static int distFuncToOrd(VectorSimilarityFunction func) { during the norm computation between the dataset vectors\ """; - static void handleThrowableWithIgnore(Throwable t, String msg) throws IOException { - if (t.getMessage().contains(msg)) { - return; - } - Utils.handleThrowable(t); - } - private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { @@ -436,7 +426,7 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw // Access the CAGRA index for this field from the reader if (knnReader != null) { - if (knnReader instanceof CuVSVectorsReader cvr) { + if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { if (cvr != null) { totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); @@ -463,6 +453,21 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw } } + /** + * Creates List from merged vectors + * */ + private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List res = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + int ordinal = iter.index(); + float[] vector = mergedVectorValues.vectorValue(ordinal); + res.add(vector); + } + return res; + } + /** * Fallback method that rebuilds indexes from merged vectors. * Used when native CAGRA merge() is not possible. Also used @@ -473,70 +478,27 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws throw new AssertionError("Only Float32 supported"); } try { - // We need to compute the size of the number of merged documents up-front so that we can - // compute the CuVSMatrix capacity. TODO: Find a way to do this without merging twice. - final int numMergedDocs = getMergedDocsCount(fieldInfo, mergeState); - - if (numMergedDocs != 0) { - writeFieldInternal( - fieldInfo, - () -> { - try { - return createMatrixFromMergedVectors( - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues( - fieldInfo, mergeState), - numMergedDocs); - } catch (IOException e) { - throw new RuntimeException(e); - } - }, - numMergedDocs); - } else { - writeEmpty(fieldInfo); - } + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); } catch (Throwable t) { Utils.handleThrowable(t); } } - private int getMergedDocsCount(FieldInfo fieldInfo, MergeState mergeState) throws IOException { - KnnVectorValues.DocIndexIterator iter = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) - .iterator(); - int numMergedDocs = 0; - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - numMergedDocs++; - } - return numMergedDocs; - } - - /** - * Creates CuVSMatrix from merged vectors - * */ - private CuVSMatrix createMatrixFromMergedVectors( - FloatVectorValues mergedVectorValues, int numMergedDocs) throws IOException { - List vectors = new ArrayList<>(numMergedDocs); - KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - int ordinal = iter.index(); - float[] vector = mergedVectorValues.vectorValue(ordinal); - vectors.add(vector.clone()); - } - return Utils.createFloatMatrix(vectors, mergedVectorValues.dimension()); - } - /** * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. */ - private CagraIndex getCagraIndexFromReader(CuVSVectorsReader reader, String fieldName) { + private CagraIndex getCagraIndexFromReader(CuVS2510GPUVectorsReader reader, String fieldName) { try { - IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); + IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); FieldInfos fieldInfos = reader.getFieldInfos(); FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); if (fieldInfo != null) { - CuVSIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); + GPUIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); if (cuvsIndex != null) { return cuvsIndex.getCagraIndex(); } @@ -562,7 +524,6 @@ private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, mergedIndex.serialize(cagraIndexOutputStream, tmpFile); long cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; - // Write metadata (assuming no brute force or HNSW indexes for merged result) writeMeta(fieldInfo, vectorCount, cagraIndexOffset, cagraIndexLength, 0L, 0L, 0L, 0L); // Clean up the merged index @@ -618,11 +579,22 @@ public void finish() throws IOException { if (cuvsIndex != null) { CodecUtil.writeFooter(cuvsIndex); } + + { + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } } @Override public void close() throws IOException { - IOUtils.close(meta, cuvsIndex, flatVectorsWriter); + IOUtils.close(meta, cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java deleted file mode 100644 index 8a601b7..0000000 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSSegmentFile.java +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed 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 com.nvidia.cuvs.lucene; - -import java.io.IOException; -import java.io.OutputStream; -import java.util.Collections; -import java.util.HashSet; -import java.util.Set; -import java.util.logging.Logger; -import java.util.zip.Deflater; -import java.util.zip.ZipEntry; -import java.util.zip.ZipOutputStream; - -/** Methods to deal with a CuVS composite file inside a segment */ -/*package-private*/ class CuVSSegmentFile implements AutoCloseable { - private final ZipOutputStream zos; - - private Set filesAdded = new HashSet(); - - public CuVSSegmentFile(OutputStream out) { - zos = new ZipOutputStream(out); - zos.setLevel(Deflater.NO_COMPRESSION); - } - - protected Logger log = Logger.getLogger(getClass().getName()); - - public void addFile(String name, byte[] bytes) throws IOException { - ZipEntry indexFileZipEntry = new ZipEntry(name); - zos.putNextEntry(indexFileZipEntry); - zos.write(bytes, 0, bytes.length); - zos.closeEntry(); - filesAdded.add(name); - } - - public Set getFilesAdded() { - return Collections.unmodifiableSet(filesAdded); - } - - @Override - public void close() throws IOException { - zos.close(); - } -} diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java new file mode 100644 index 0000000..7235ec9 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -0,0 +1,258 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.RowView; +import java.util.ArrayList; +import java.util.List; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.NeighborArray; + +public class GPUBuiltHnswGraph extends HnswGraph { + + private final int size; + private final int dimensions; + private final int numLevels; + + // Store layers data - each layer has its own nodes and adjacency lists + private final List layerNodes; + private final List layerNeighbors; + + // Layer 0 is special - it contains all nodes + private final NeighborArray[] layer0Neighbors; + + // Multi-layer constructor that supports arbitrary number of layers + public GPUBuiltHnswGraph( + int size, int dimensions, List layerNodes, List layerAdjacencies) { + + this.size = size; + this.dimensions = dimensions; + this.numLevels = layerAdjacencies.size(); + this.layerNodes = new ArrayList<>(); + this.layerNeighbors = new ArrayList<>(); + + // Process Layer 0 (base layer with all nodes) + CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); + + // Process higher layers (1 to numLevels-1) + for (int level = 1; level < numLevels; level++) { + int[] nodes = layerNodes.get(level); + CuVSMatrix adjacency = layerAdjacencies.get(level); + this.layerNodes.add(nodes); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + } + } + + private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + NeighborArray[] neighbors = new NeighborArray[size]; + for (int i = 0; i < size; i++) { + RowView rv = adjacency.getRow(i); + if (rv != null && rv.size() > 0) { + neighbors[i] = new NeighborArray((int) rv.size(), true); + for (int j = 0; j < rv.size(); j++) { + neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + } + } else { + neighbors[i] = new NeighborArray(0, true); + } + } + return neighbors; + } + + public NodesIterator getNodesOnLevel(int level) { + if (level == 0) { + return new Level0NodesIterator(size); + } else if (level > 0 && level < numLevels) { + int[] nodes = layerNodes.get(level - 1); + return new HigherLevelNodesIterator(nodes); + } else { + return new Level0NodesIterator(0); + } + } + + public NeighborArray getNeighbors(int level, int node) { + if (level == 0 && node < size) { + return layer0Neighbors[node]; + } else if (level > 0 && level < numLevels) { + int[] nodes = layerNodes.get(level - 1); + NeighborArray[] neighbors = layerNeighbors.get(level - 1); + + // Find the index of this node in the layer + for (int i = 0; i < nodes.length; i++) { + if (nodes[i] == node) { + return neighbors[i]; + } + } + } + return null; + } + + // Implementation of abstract methods from HnswGraph + private int currentNode = -1; + private int currentLevel = -1; + private int neighborIndex = -1; + + @Override + public void seek(int level, int target) { + currentLevel = level; + currentNode = target; + neighborIndex = -1; + } + + @Override + public int nextNeighbor() { + if (currentLevel == 0 + && currentNode >= 0 + && currentNode < size + && layer0Neighbors[currentNode] != null) { + neighborIndex++; + if (neighborIndex < layer0Neighbors[currentNode].size()) { + int neighborNode = layer0Neighbors[currentNode].nodes()[neighborIndex]; + if (neighborNode >= 0 && neighborNode < size) { + return neighborNode; + } else { + return nextNeighbor(); // Skip invalid neighbor + } + } + } else if (currentLevel > 0 && currentLevel < numLevels) { + // Handle higher layers + NeighborArray neighbors = getNeighbors(currentLevel, currentNode); + if (neighbors != null) { + neighborIndex++; + if (neighborIndex < neighbors.size()) { + return neighbors.nodes()[neighborIndex]; + } + } + } + return NO_MORE_DOCS; + } + + @Override + public int entryNode() { + // Entry node should be from the highest layer + if (numLevels > 1) { + int topLevel = numLevels - 1; + int[] topLayerNodes = layerNodes.get(topLevel - 1); + if (topLayerNodes != null && topLayerNodes.length > 0) { + // Use random node from top layer with fixed seed for reproducibility + java.util.Random random = new java.util.Random(44); + int randomIndex = random.nextInt(topLayerNodes.length); + return topLayerNodes[randomIndex]; + } + } + return 0; // Default to node 0 for single-layer graphs + } + + @Override + public int maxConn() { + // Return the maximum degree across all nodes in layer 0 + int max = 0; + for (NeighborArray neighbor : layer0Neighbors) { + if (neighbor != null) { + max = Math.max(max, neighbor.size()); + } + } + return max; + } + + @Override + public int neighborCount() { + if (currentLevel == 0 + && currentNode >= 0 + && currentNode < size + && layer0Neighbors[currentNode] != null) { + return layer0Neighbors[currentNode].size(); + } else if (currentLevel > 0 && currentLevel < numLevels) { + NeighborArray neighbors = getNeighbors(currentLevel, currentNode); + return neighbors != null ? neighbors.size() : 0; + } + return 0; + } + + // NodesIterator for level 0 + private static class Level0NodesIterator extends NodesIterator { + private int current = -1; + + Level0NodesIterator(int size) { + super(size); + } + + @Override + public boolean hasNext() { + return current + 1 < size; + } + + @Override + public int nextInt() { + return ++current; + } + + @Override + public int consume(int[] dest) { + int numToCopy = Math.min(dest.length, size - (current + 1)); + for (int i = 0; i < numToCopy; i++) { + dest[i] = ++current; + } + return numToCopy; + } + } + + // NodesIterator for higher layers + private static class HigherLevelNodesIterator extends NodesIterator { + private final int[] nodeIds; + private int current = -1; + + HigherLevelNodesIterator(int[] nodeIds) { + super(nodeIds.length); + this.nodeIds = nodeIds; + } + + @Override + public boolean hasNext() { + return current + 1 < nodeIds.length; + } + + @Override + public int nextInt() { + return nodeIds[++current]; + } + + @Override + public int consume(int[] dest) { + int numToCopy = Math.min(dest.length, nodeIds.length - (current + 1)); + for (int i = 0; i < numToCopy; i++) { + dest[i] = nodeIds[++current]; + } + return numToCopy; + } + } + + public int size() { + return size; + } + + public int numLevels() { + return numLevels; + } + + public int dimensions() { + return dimensions; + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSFieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java similarity index 84% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSFieldWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java index acd1518..483e18d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSFieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java @@ -24,16 +24,16 @@ import org.apache.lucene.util.RamUsageEstimator; /** CuVS based fields writer */ -/*package-private*/ class CuVSFieldWriter extends KnnFieldVectorsWriter { +/*package-private*/ class GPUFieldWriter extends KnnFieldVectorsWriter { private static final long SHALLOW_SIZE = - RamUsageEstimator.shallowSizeOfInstance(CuVSFieldWriter.class); + RamUsageEstimator.shallowSizeOfInstance(GPUFieldWriter.class); private final FieldInfo fieldInfo; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; - public CuVSFieldWriter( + public GPUFieldWriter( FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { this.fieldInfo = fieldInfo; this.flatFieldVectorsWriter = flatFieldVectorsWriter; @@ -74,6 +74,10 @@ public long ramBytesUsed() { @Override public String toString() { - return "CuVSFieldWriter[field name=" + fieldInfo.name + ", number=" + fieldInfo.number + "]"; + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(field name=").append(fieldInfo.name); + sb.append("number=").append(fieldInfo.number); + sb.append(")"); + return sb.toString(); } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java similarity index 92% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java index 1151980..ea33c84 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java @@ -23,7 +23,7 @@ import java.util.Objects; /** This class holds references to the actual CuVS Index (Cagra, Brute force, etc.) */ -public class CuVSIndex implements Closeable { +public class GPUIndex implements Closeable { private final CagraIndex cagraIndex; private final BruteForceIndex bruteforceIndex; private final HnswIndex hnswIndex; @@ -33,7 +33,7 @@ public class CuVSIndex implements Closeable { private String segmentName; private volatile boolean closed; - public CuVSIndex( + public GPUIndex( String segmentName, String fieldName, CagraIndex cagraIndex, @@ -47,10 +47,10 @@ public CuVSIndex( throw new IllegalArgumentException("negative maxDocs:" + maxDocs); } this.maxDocs = maxDocs; - this.hnswIndex = null; // TODO: + this.hnswIndex = null; // TODO: remove hnswlib logic in a subsequent PR } - public CuVSIndex(CagraIndex cagraIndex, BruteForceIndex bruteforceIndex, HnswIndex hnswIndex) { + public GPUIndex(CagraIndex cagraIndex, BruteForceIndex bruteforceIndex, HnswIndex hnswIndex) { this.cagraIndex = cagraIndex; this.bruteforceIndex = bruteforceIndex; this.hnswIndex = hnswIndex; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java similarity index 86% rename from src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 8caf30a..64f7081 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -18,19 +18,20 @@ import java.io.IOException; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.KnnCollector; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.knn.KnnCollectorManager; import org.apache.lucene.util.Bits; -/** Query for CuVS */ -public class CuVSKnnFloatVectorQuery extends KnnFloatVectorQuery { +/** Query on GPU only */ +public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; - public CuVSKnnFloatVectorQuery( + public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { super(field, target, k, filter); this.iTopK = iTopK; @@ -45,7 +46,7 @@ protected TopDocs approximateSearch( KnnCollectorManager knnCollectorManager) throws IOException { - PerLeafCuVSKnnCollector results = new PerLeafCuVSKnnCollector(k, iTopK, searchWidth); + KnnCollector results = new GPUPerLeafCuVSKnnCollector(k, iTopK, searchWidth); LeafReader reader = context.reader(); reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs); diff --git a/src/main/java/com/nvidia/cuvs/lucene/PerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java similarity index 93% rename from src/main/java/com/nvidia/cuvs/lucene/PerLeafCuVSKnnCollector.java rename to src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 8e00557..ca30d3d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/PerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -24,7 +24,7 @@ import org.apache.lucene.search.knn.KnnSearchStrategy; /** KnnCollector for CuVS */ -/*package-private*/ class PerLeafCuVSKnnCollector implements KnnCollector { +/*package-private*/ class GPUPerLeafCuVSKnnCollector implements KnnCollector { public List scoreDocs; public int topK = 0; @@ -32,7 +32,7 @@ public int searchWidth = 1; // TODO getter, no setter public int results = 0; - public PerLeafCuVSKnnCollector(int topK, int iTopK, int searchWidth) { + public GPUPerLeafCuVSKnnCollector(int topK, int iTopK, int searchWidth) { super(); this.topK = topK; this.iTopK = iTopK; diff --git a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java index b786636..53283eb 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java +++ b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java @@ -27,7 +27,7 @@ final class IndexOutputOutputStream extends OutputStream { final IndexOutput out; final int bufferSize; final byte[] buffer; - int idx; + int pos; IndexOutputOutputStream(IndexOutput out) { this(out, DEFAULT_BUFFER_SIZE); @@ -41,16 +41,16 @@ final class IndexOutputOutputStream extends OutputStream { @Override public void write(int b) throws IOException { - buffer[idx] = (byte) b; - idx++; - if (idx == bufferSize) { + buffer[pos] = (byte) b; + pos++; + if (pos == bufferSize) { flush(); } } @Override public void write(byte[] b, int offset, int length) throws IOException { - if (idx != 0) { + if (pos != 0) { flush(); } out.writeBytes(b, offset, length); @@ -58,8 +58,8 @@ public void write(byte[] b, int offset, int length) throws IOException { @Override public void flush() throws IOException { - out.writeBytes(buffer, 0, idx); - idx = 0; + out.writeBytes(buffer, 0, pos); + pos = 0; } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java new file mode 100644 index 0000000..facd93d --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -0,0 +1,97 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.cuvs.lucene; + +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; + +import com.nvidia.cuvs.LibraryException; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene101.Lucene101Codec; + +/** CuVS based codec for GPU based vector search */ +public class Lucene101AcceleratedHNSWCodec extends FilterCodec { + + private static final Logger log = Logger.getLogger(Lucene101AcceleratedHNSWCodec.class.getName()); + + private static final int DEFAULT_CUVS_WRITER_THREADS = 1; + private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + private static final int DEFAULT_GRAPH_DEGREE = 64; + private static final int DEFAULT_HNSW_LAYERS = 1; + private static final String NAME = "Lucene101AcceleratedHNSWCodec"; + + private KnnVectorsFormat format; + + public Lucene101AcceleratedHNSWCodec() { + this(NAME, new Lucene101Codec()); + } + + public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { + super(name, delegate); + initializeFormatDefaultValues(); + } + + public Lucene101AcceleratedHNSWCodec( + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + int maxConn, + int beamWidth) { + this(NAME, new Lucene101Codec()); + initializeFormat( + cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, maxConn, beamWidth); + } + + private void initializeFormatDefaultValues() { + initializeFormat( + DEFAULT_CUVS_WRITER_THREADS, + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + DEFAULT_HNSW_LAYERS, + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH); + } + + private void initializeFormat( + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + int maxConn, + int beamWidth) { + try { + format = + new Lucene99AcceleratedHNSWVectorsFormat( + cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, maxConn, beamWidth); + setKnnFormat(format); + } catch (LibraryException ex) { + log.severe("Couldn't load native library, possible classloader issue. " + ex.getMessage()); + } + } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return format; + } + + public void setKnnFormat(KnnVectorsFormat format) { + this.format = format; + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java new file mode 100644 index 0000000..3a18aae --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -0,0 +1,161 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER; + +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.LibraryException; +import java.io.IOException; +import java.util.logging.Logger; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** CuVS based KnnVectorsFormat for GPU acceleration */ +public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { + + private static final Logger log = + Logger.getLogger(Lucene99AcceleratedHNSWVectorsFormat.class.getName()); + + static final int DEFAULT_WRITER_THREADS = 32; + static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; + static final int DEFAULT_GRAPH_DEGREE = 64; + static final int DEFAULT_HNSW_GRAPH_LAYERS = 1; + + static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; + static final String HNSW_META_CODEC_EXT = "vem"; + static final String HNSW_INDEX_CODEC_NAME = "Lucene99HnswVectorsFormatIndex"; + static final String HNSW_INDEX_EXT = "vex"; + + private static CuVSResources resources = cuVSResourcesOrNull(); + + /** The format for storing, reading, and merging raw vectors on disk. */ + private static final FlatVectorsFormat flatVectorsFormat = + new Lucene99FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE); + + private final int maxDimensions = 4096; + private final int cuvsWriterThreads; + private final int intGraphDegree; + private final int graphDegree; + private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion + + private final int maxConn; + private final int beamWidth; + + /** + * Creates a Lucene99AcceleratedHNSWVectorsFormat, with default values. + * + * @throws LibraryException if the native library fails to load + */ + public Lucene99AcceleratedHNSWVectorsFormat() { + this( + DEFAULT_WRITER_THREADS, + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + DEFAULT_HNSW_GRAPH_LAYERS, + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH); + } + + /** + * Creates a Lucene99AcceleratedHNSWVectorsFormat, with the given threads, graph degree, etc. + * + * @throws LibraryException if the native library fails to load + */ + public Lucene99AcceleratedHNSWVectorsFormat( + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + int maxConn, + int beamWidth) { + super("Lucene99AcceleratedHNSWVectorsFormat"); + this.cuvsWriterThreads = cuvsWriterThreads; + this.intGraphDegree = intGraphDegree; + this.graphDegree = graphDegree; + this.hnswLayers = hnswLayers; + this.maxConn = maxConn; + this.beamWidth = beamWidth; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + var flatWriter = flatVectorsFormat.fieldsWriter(state); + if (supported()) { + log.info("cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); + return new Lucene99AcceleratedHNSWVectorsWriter( + state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); + } else { + log.warning( + "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); + // TODO: Make num merge workers configurable. + return new Lucene99HnswVectorsWriter( + state, maxConn, beamWidth, flatWriter, DEFAULT_NUM_MERGE_WORKER, null); + } + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state)); + } + + @Override + public int getMaxDimensions(String fieldName) { + return maxDimensions; + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()); + sb.append("(cuvsWriterThreads=").append(cuvsWriterThreads); + sb.append("intGraphDegree=").append(intGraphDegree); + sb.append("graphDegree=").append(graphDegree); + sb.append("hnswLayers=").append(hnswLayers); + sb.append("resources=").append(resources); + sb.append(")"); + return sb.toString(); + } + + public static CuVSResources getResources() { + return resources; + } + + public static void setResources(CuVSResources resources) { + Lucene99AcceleratedHNSWVectorsFormat.resources = resources; + } + + /** Tells whether the platform supports cuVS. */ + public static boolean supported() { + return resources != null; + } + + public static void checkSupported() { + if (!supported()) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java new file mode 100644 index 0000000..e15b743 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -0,0 +1,702 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; + +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.Random; +import java.util.SortedSet; +import java.util.TreeSet; +import java.util.logging.Logger; +import java.util.stream.IntStream; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntObjectHashMap; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.apache.lucene.util.packed.DirectMonotonicWriter; + +/** + * KnnVectorsWriter for CuVS, responsible for merge and flush of vectors into + * GPU + */ +public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_RAM_BYTES_USED = + shallowSizeOfInstance(Lucene99AcceleratedHNSWVectorsWriter.class); + + @SuppressWarnings("unused") + private static final Logger log = + Logger.getLogger(Lucene99AcceleratedHNSWVectorsWriter.class.getName()); + + /** The name of the CUVS component for the info-stream * */ + private static final String CUVS_COMPONENT = "CUVS"; + + private final int cuvsWriterThreads; + private final int intGraphDegree; + private final int graphDegree; + private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion + private final CuVSResources resources; + private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors + private final List fields = new ArrayList<>(); + private final InfoStream infoStream; + private IndexOutput cuvsIndex = null; + private IndexOutput hnswMeta = null, hnswVectorIndex = null; + private boolean finished; + private String vemFileName; + private String vexFileName; + + public Lucene99AcceleratedHNSWVectorsWriter( + SegmentWriteState state, + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + int hnswLayers, + CuVSResources resources, + FlatVectorsWriter flatVectorsWriter) + throws IOException { + super(); + this.cuvsWriterThreads = cuvsWriterThreads; + this.intGraphDegree = intGraphDegree; + this.graphDegree = graphDegree; + this.hnswLayers = hnswLayers; + this.resources = resources; + this.flatVectorsWriter = flatVectorsWriter; + this.infoStream = state.infoStream; + + vemFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); + + vexFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); + + boolean success = false; + try { + + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + + CodecUtil.writeIndexHeader( + hnswMeta, + HNSW_META_CODEC_NAME, + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + HNSW_INDEX_CODEC_NAME, + Lucene99HnswVectorsFormat.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + var encoding = fieldInfo.getVectorEncoding(); + if (encoding != FLOAT32) { + throw new IllegalArgumentException("expected float32, got:" + encoding); + } + var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); + @SuppressWarnings("unchecked") + var flatWriter = (FlatFieldVectorsWriter) writer; + var cuvsFieldWriter = new GPUFieldWriter(fieldInfo, flatWriter); + fields.add(cuvsFieldWriter); + return writer; + } + + static String indexMsg(int size, int... args) { + StringBuilder sb = new StringBuilder("cagra index params"); + sb.append(": size=").append(size); + sb.append(", intGraphDegree=").append(args[0]); + sb.append(", actualIntGraphDegree=").append(args[1]); + sb.append(", graphDegree=").append(args[2]); + sb.append(", actualGraphDegree=").append(args[3]); + return sb.toString(); + } + + private CagraIndexParams cagraIndexParams() { + return new CagraIndexParams.Builder() + .withNumWriterThreads(cuvsWriterThreads) + .withIntermediateGraphDegree(intGraphDegree) + .withGraphDegree(graphDegree) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .build(); + } + + private void info(String msg) { + if (infoStream.isEnabled(CUVS_COMPONENT)) { + infoStream.message(CUVS_COMPONENT, msg); + } + } + + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + + if (vectors.size() == 0) { + writeEmpty(fieldInfo); + return; + } + + try { + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + + if (dataset.size() < 2) { + throw new IllegalArgumentException(dataset.size() + " vectors, less than min [2] required"); + } + + long startTime = System.nanoTime(); + CagraIndexParams params = cagraIndexParams(); + CagraIndex cagraIndex = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build(); + + // Get the adjacency list from CAGRA index + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + + int size = (int) dataset.size(); + int dimensions = fieldInfo.getVectorDimension(); + + // Create multi-layer HNSW graph from CAGRA + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, size, dimensions, adjacencyListMatrix, vectors, hnswLayers); + + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + + // Write the graph to the vector index + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets); + + long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); + info("HNSW graph created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); + + cagraIndex.close(); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Creates a multi-layer HNSW graph with dynamic number of layers. + * M = cagraGraphDegree/2 + * Each layer contains 1/M nodes from the previous layer + * Creates layers until the highest layer has ≤ M nodes + */ + private GPUBuiltHnswGraph createMultiLayerHnswGraph( + FieldInfo fieldInfo, + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers) + throws Throwable { + + // Calculate M as cagraGraphDegree/2 + int M = graphDegree / 2; + + // Store all layers data + List layerNodes = new ArrayList<>(); + List layerAdjacencies = new ArrayList<>(); + + // Layer 0: Use full CAGRA adjacency list + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list + layerAdjacencies.add(adjacencyListMatrix); + + int currentLayerSize = size; + int layerIndex = 1; + Random random = new Random(); + + while (layerIndex < hnswLayers && currentLayerSize > 1) { + // Calculate size for next layer (1/M of current layer) + int nextLayerSize = Math.max(2, currentLayerSize / M); + // Select nodes for this layer + SortedSet selectedNodesSet = new TreeSet<>(); + + if (layerIndex == 1) { + // Select from all nodes (Layer 0) + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(random.nextInt(size)); + } + } else { + // Select from previous layer nodes + int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); + while (selectedNodesSet.size() < nextLayerSize) { + int idx = random.nextInt(prevLayerNodes.length); + selectedNodesSet.add(prevLayerNodes[idx]); + } + } + + // Convert to sorted array + int[] selectedNodes = + selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); + + layerNodes.add(selectedNodes); + + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = vectors.get(selectedNodes[i]); + } + + // Build CAGRA graph for this layer + layerAdjacencies.add(buildCagraGraphForSubset(selectedVectors, selectedNodes)); + + // Update for next iteration + currentLayerSize = nextLayerSize; + layerIndex++; + + // Use different seed for each layer + random = new Random(new Random().nextLong()); + } + + // Create the multi-layer graph with all layers + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } + + /** + * Builds a CAGRA graph for a subset of vectors + */ + private CuVSMatrix buildCagraGraphForSubset(float[][] vectors, int[] selectedNodes) + throws Throwable { + // Create CuVSMatrix from the subset vectors + CuVSMatrix subsetDataset = CuVSMatrix.ofArray(vectors); + + // Build CAGRA index for the subset + CagraIndexParams params = cagraIndexParams(); + CagraIndex subsetIndex = + CagraIndex.newBuilder(resources).withDataset(subsetDataset).withIndexParams(params).build(); + + // Get adjacency list from subset CAGRA index + CuVSMatrix cagraGraph = subsetIndex.getGraph(); + + long numNodes = cagraGraph.size(); + long degree = cagraGraph.columns(); + + // Create a re-mapped adjacency list + int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; + + for (int i = 0; i < numNodes; i++) { + RowView rv = cagraGraph.getRow(i); + for (int j = 0; j < degree && j < rv.size(); j++) { + int subsetIndex1 = rv.getAsInt(j); + // Map subset index to original node ID + if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { + remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; + } else { + // Invalid index, use self-reference + remappedAdjacency[i][j] = selectedNodes[i]; + } + } + } + + subsetIndex.close(); + return CuVSMatrix.ofArray(remappedAdjacency); + } + + private void writeMeta( + IndexOutput vectorIndex, + IndexOutput meta, + FieldInfo field, + long vectorIndexOffset, + long vectorIndexLength, + int count, + HnswGraph graph, + int[][] graphLevelNodeOffsets) + throws IOException { + + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); + meta.writeVLong(vectorIndexOffset); + meta.writeVLong(vectorIndexLength); + meta.writeVInt(field.getVectorDimension()); + meta.writeInt(count); + meta.writeVInt(graphDegree / 2); // M = cagraGraphDegree/2 + + // write graph nodes on each level + if (graph == null) { + meta.writeVInt(0); + } else { + meta.writeVInt(graph.numLevels()); + long valueCount = 0; + for (int level = 0; level < graph.numLevels(); level++) { + NodesIterator nodesOnLevel = graph.getNodesOnLevel(level); + valueCount += nodesOnLevel.size(); + if (level > 0) { + int[] nol = new int[nodesOnLevel.size()]; + int numberConsumed = nodesOnLevel.consume(nol); + Arrays.sort(nol); + assert numberConsumed == nodesOnLevel.size(); + meta.writeVInt(nol.length); // number of nodes on a level + for (int i = nodesOnLevel.size() - 1; i > 0; --i) { + nol[i] -= nol[i - 1]; + } + for (int n : nol) { + meta.writeVInt(n); + } + } else { + assert nodesOnLevel.size() == count : "Level 0 expects to have all nodes"; + } + } + + long start = vectorIndex.getFilePointer(); + meta.writeLong(start); + meta.writeVInt(16); // DIRECT_MONOTONIC_BLOCK_SHIFT); + + final DirectMonotonicWriter memoryOffsetsWriter = + DirectMonotonicWriter.getInstance(meta, vectorIndex, valueCount, 16); + long cumulativeOffsetSum = 0; + for (int[] levelOffsets : graphLevelNodeOffsets) { + for (int v : levelOffsets) { + memoryOffsetsWriter.add(cumulativeOffsetSum); + cumulativeOffsetSum += v; + } + } + + memoryOffsetsWriter.finish(); + + meta.writeLong(vectorIndex.getFilePointer() - start); + } + } + + private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { + // write vectors' neighbors on each level into the vectorIndex file + int countOnLevel0 = graph.size(); + int[][] offsets = new int[graph.numLevels()][]; + int[] scratch = new int[graph.maxConn() * 2]; + for (int level = 0; level < graph.numLevels(); level++) { + int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); + offsets[level] = new int[sortedNodes.length]; + int nodeOffsetId = 0; + + for (int node : sortedNodes) { + // Get node neighbors + NeighborArray neighbors = graph.getNeighbors(level, node); + // Get the size of the neighbor array + int size = neighbors.size(); + // Write size in VInt as the neighbors list is typically small + long offsetStart = vectorIndex.getFilePointer(); + // Get neighbors + int[] nnodes = neighbors.nodes(); + // Sort them + Arrays.sort(nnodes, 0, size); + // Now that we have sorted, do delta encoding to minimize the required bits to store the + // information + int actualSize = 0; + if (size > 0) { + scratch[0] = nnodes[0]; + actualSize = 1; + } + // De-duplication + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + // Sorting step helps here + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + } + // Write the size after duplicates are removed + vectorIndex.writeVInt(actualSize); + // Write de-duplicated neighbors + for (int i = 0; i < actualSize; i++) { + vectorIndex.writeVInt(scratch[i]); + } + offsets[level][nodeOffsetId++] = + Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); + } + } + // Return offsets (information written while writing the meta info) + return offsets; + } + + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + flatVectorsWriter.flush(maxDoc, sortMap); + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } + } + + private void writeField(GPUFieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); + } + + private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) + throws IOException { + + DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord + mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); + + List sortedVectors = new ArrayList(); + for (int i = 0; i < fieldData.getVectors().size(); i++) { + sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); + } + + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); + } + + private void writeEmpty(FieldInfo fieldInfo) throws IOException { + writeMeta(null, hnswMeta, fieldInfo, 0, 0, 0, null, null); + } + + static int distFuncToOrd(VectorSimilarityFunction func) { + for (int i = 0; i < SIMILARITY_FUNCTIONS.size(); i++) { + if (SIMILARITY_FUNCTIONS.get(i).equals(func)) { + return (byte) i; + } + } + throw new IllegalArgumentException("invalid distance function: " + func); + } + + private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + + List cagraIndexes = new ArrayList<>(); + // We need this count so that the merged segment's meta information has the vector count. + int totalVectorCount = 0; + + for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { + KnnVectorsReader knnReader = mergeState.knnVectorsReaders[i]; + // Access the CAGRA index for this field from the reader + + if (knnReader != null) { + if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { + if (cvr != null) { + totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); + CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); + if (cagraIndex != null) { + cagraIndexes.add(cagraIndex); + } + } + } else { + // This should never happen + throw new RuntimeException( + "Reader is not of CuVSVectorsReader type. Instead it is: " + knnReader.getClass()); + } + } + } + assert cagraIndexes.size() > 1; + + CagraIndex mergedIndex = + CagraIndex.merge(cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()])); + writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + info("Successfully merged " + cagraIndexes.size() + " CAGRA indexes using native merge API"); + + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Creates List from merged vectors + * */ + private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List vectors = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + float[] vector = mergedVectorValues.vectorValue(iter.index()); + vectors.add(vector); + } + return vectors; + } + + /** + * Fallback method that rebuilds indexes from merged vectors. + * Used when native CAGRA merge() is not possible. Also used + * when non-CAGRA index types are used (for e.g. Brute Force index). + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + if (fieldInfo.getVectorEncoding() != FLOAT32) { + throw new AssertionError("Only Float32 supported"); + } + try { + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /** + * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. + */ + private CagraIndex getCagraIndexFromReader(CuVS2510GPUVectorsReader reader, String fieldName) { + try { + IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); + FieldInfos fieldInfos = reader.getFieldInfos(); + FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); + + if (fieldInfo != null) { + GPUIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); + if (cuvsIndex != null) { + return cuvsIndex.getCagraIndex(); + } + } + } catch (Exception e) { + e.printStackTrace(); + info("Failed to extract CAGRA index for field " + fieldName + ": " + e.getMessage()); + } + return null; + } + + /** + * Writes a pre-built merged CAGRA index to the output. + */ + private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, int vectorCount) + throws IOException { + try { + var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); + + // Serialize the merged index + Path tmpFile = Files.createTempFile(resources.tempDirectory(), "mergedindex", "cag"); + mergedIndex.serialize(cagraIndexOutputStream, tmpFile); + + // TODO: Path to writeFieldInternal missing. Fix this. + + // Clean up the merged index + mergedIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + flatVectorsWriter.mergeOneField(fieldInfo, mergeState); + + // Since CAGRA merge does not support merging of indexes with purging of deletes, + // we fallback to vector-based re-indexing. Issue: + // https://github.com/rapidsai/cuvs/issues/1253 + boolean hasDeletions = + IntStream.range(0, mergeState.liveDocs.length) + .anyMatch( + i -> + mergeState.liveDocs[i] == null + || IntStream.range(0, mergeState.maxDocs[i]) + .anyMatch(j -> !mergeState.liveDocs[i].get(j))); + + if (mergeState.knnVectorsReaders.length > 1 && !hasDeletions) { + mergeCagraIndexes(fieldInfo, mergeState); + } else { + // CAGRA's merge API does not handle the trivial case of merging 1 index. + vectorBasedMerge(fieldInfo, mergeState); + } + } + + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + flatVectorsWriter.finish(); + + if (cuvsIndex != null) { + CodecUtil.writeFooter(cuvsIndex); + } + + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } + } + + @Override + public void close() throws IOException { + IOUtils.close(cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); + } + + @Override + public long ramBytesUsed() { + long total = SHALLOW_RAM_BYTES_USED; + for (var field : fields) { + total += field.ramBytesUsed(); + } + return total; + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 8d1d4bd..dfb9aa0 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -16,11 +16,16 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; import java.io.IOException; +import java.time.Duration; import java.util.List; +import java.util.logging.Logger; public class Utils { + static final Logger log = Logger.getLogger(Utils.class.getName()); + static void handleThrowable(Throwable t) throws IOException { switch (t) { case IOException ioe -> throw ioe; @@ -45,4 +50,34 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions) { float[][] vectors = data.toArray(new float[0][]); return CuVSMatrix.ofArray(vectors); } + + static long nanosToMillis(long nanos) { + return Duration.ofNanos(nanos).toMillis(); + } + + static CuVSResources cuVSResourcesOrNull() { + try { + System.loadLibrary("cudart"); + } catch (UnsatisfiedLinkError e) { + log.warning("Could not load CUDA runtime library: " + e.getMessage()); + } + try { + return CuVSResources.create(); + } catch (UnsupportedOperationException uoe) { + log.warning("cuVS is not supported on this platform or java version: " + uoe.getMessage()); + } catch (Throwable t) { + if (t instanceof ExceptionInInitializerError ex) { + t = ex.getCause(); + } + log.warning("Exception occurred during creation of cuVS resources. " + t); + } + return null; + } + + static void handleThrowableWithIgnore(Throwable t, String msg) throws IOException { + if (t.getMessage().contains(msg)) { + return; + } + handleThrowable(t); + } } diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec new file mode 100644 index 0000000..9865a33 --- /dev/null +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -0,0 +1,2 @@ +com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec +com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 6e486ee..747a157 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -15,4 +15,5 @@ org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat -com.nvidia.cuvs.lucene.CuVSVectorsFormat +com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat +com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java new file mode 100644 index 0000000..4f5f7ab --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -0,0 +1,167 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { + + private static Logger log = + Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); + private static Random random; + private static Path indexDirPath; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); + // Fixed seed so that we can validate against the same result. + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @Test + public void testCagraToHnswSerializationAndSearch() throws IOException { + Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath)) { + try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.info("Successfully opened index"); + + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(knnValues); + log.info( + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + } + assertTrue("Dataset size mismatch", vectorCount == numDocs); + + log.info("Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.info("Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.info("Search results (" + results.totalHits + " total hits):"); + Integer[] expected = new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}; + HashSet expectedIds = new HashSet(Arrays.asList(expected)); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.info( + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue( + "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + + } catch (Exception e) { + e.printStackTrace(); + } + } + } + + @AfterClass + public static void afterClass() throws Exception { + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java new file mode 100644 index 0000000..792aca9 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -0,0 +1,173 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed 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 com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraToHnswSerializationAndSearchWithFallbackWriter extends LuceneTestCase { + + private static Logger log = + Logger.getLogger(TestCagraToHnswSerializationAndSearchWithFallbackWriter.class.getName()); + + private static Random random; + private static Path indexDirPath; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); + // Set resources to null to simulate that cuVS is not supported. + Lucene99AcceleratedHNSWVectorsFormat.setResources(null); + // Fixed seed so that we can validate against the same result. + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @Test + public void testCagraToHnswSerializationAndSearchWithFallbackWriter() throws IOException { + Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath)) { + try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.info("Successfully opened index"); + + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(knnValues); + log.info( + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + } + assertTrue("Dataset size mismatch", vectorCount == numDocs); + + log.info("Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.info("Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.info("Search results (" + results.totalHits + " total hits):"); + Integer[] expected = new Integer[] {1869, 1411, 1497, 351, 554}; + HashSet expectedIds = new HashSet(Arrays.asList(expected)); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.info( + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue( + "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + + } catch (Exception e) { + e.printStackTrace(); + } + } + } + + @AfterClass + public static void afterClass() throws Exception { + // Reset resources for other tests to work + Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull()); + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index ef17136..1774587 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -51,17 +51,17 @@ import org.junit.BeforeClass; import org.junit.Test; -@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +@SuppressSysoutChecks(bugUrl = "prints info from within cuVS") public class TestCuVSDeletedDocuments extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); private static Random random; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); random = random(); } @@ -194,7 +194,7 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Test filtered search with deletions Query filter = new TermQuery(new Term("category", "A")); Query filteredQuery = - new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; for (ScoreDoc hit : filteredHits) { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index e27568b..bba34bf 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -48,12 +48,12 @@ import org.junit.BeforeClass; import org.junit.Test; -@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +@SuppressSysoutChecks(bugUrl = "prints info from within cuVS") public class TestCuVSGaps extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); static IndexSearcher searcher; static IndexReader reader; static Directory directory; @@ -70,7 +70,7 @@ public class TestCuVSGaps extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assertTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); random = random(); @@ -120,7 +120,7 @@ public static void afterClass() throws Exception { @Test public void testVectorSearchWithAlternatingDocuments() throws IOException { - assertTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; @@ -153,7 +153,7 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { @Test public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; @@ -163,7 +163,7 @@ public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOExcepti // This should further restrict our results to even numbers 0, 2, 4, 6, 8 Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 - Query filteredQuery = new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; // Should only get document 8 (the only one that matches the filter and has a vector) diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 8802373..add0fe2 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -51,12 +51,12 @@ import org.junit.BeforeClass; import org.junit.Test; -@SuppressSysoutChecks(bugUrl = "prints info from within cuvs") +@SuppressSysoutChecks(bugUrl = "prints info from within cuVS") public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); static IndexSearcher searcher; static IndexReader reader; static Directory directory; @@ -69,7 +69,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { - assertTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); RandomIndexWriter writer = @@ -184,7 +184,7 @@ private static List> generateExpectedResults( @Test public void testVectorSearchWithFilter() throws IOException { - assertTrue("cuvs not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); Random random = random(); int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); @@ -206,7 +206,7 @@ public void testVectorSearchWithFilter() throws IOException { Query filter = new TermQuery(new Term("id", targetDocId)); // Test the new constructor with filter - Query filteredQuery = new CuVSKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 69aeb29..cf78b05 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -41,12 +41,12 @@ public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { @BeforeClass public static void beforeClass() { - assertTrue("cuvs is not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported()); } @Override protected Codec getCodec() { - return TestUtil.alwaysKnnVectorsFormat(new CuVSVectorsFormat()); + return TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index a424492..0ad4861 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -17,7 +17,7 @@ import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; -import com.nvidia.cuvs.lucene.CuVSVectorsWriter.IndexType; +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -70,7 +70,7 @@ public class TestMerge extends LuceneTestCase { @BeforeClass public static void beforeClass() { - assertTrue("cuVS is not supported", CuVSVectorsFormat.supported()); + assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported()); } private Directory directory; @@ -128,7 +128,7 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -252,7 +252,7 @@ public void testMergeWithIndexSorting() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setIndexSort(indexSort) // This automatically enables sorting during merges .setMergePolicy(mergePolicy) .setMaxBufferedDocs(maxBufferedDocs) @@ -458,7 +458,7 @@ public void testMergeWithMissingVectors() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -597,7 +597,7 @@ public void testMergeWithDeletions() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); @@ -730,11 +730,12 @@ public void testMergeBruteForceIndex() throws IOException { + vectorProbability); // Configure with brute force index type - CuVSVectorsFormat bruteForceFormat = - new CuVSVectorsFormat( + CuVS2510GPUVectorsFormat bruteForceFormat = + new CuVS2510GPUVectorsFormat( 32, // writer threads 128, // intermediate graph degree 64, // graph degree + 1, IndexType.BRUTE_FORCE); // Use brute force index IndexWriterConfig config = @@ -881,11 +882,12 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { + vectorProbability); // Configure with CAGRA + brute force combined index type - CuVSVectorsFormat combinedFormat = - new CuVSVectorsFormat( + CuVS2510GPUVectorsFormat combinedFormat = + new CuVS2510GPUVectorsFormat( 32, // writer threads 128, // intermediate graph degree 64, // graph degree + 1, IndexType.CAGRA_AND_BRUTE_FORCE); // Use combined CAGRA + brute force IndexWriterConfig config = @@ -1056,7 +1058,7 @@ public void testLargeScaleMerge() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVSVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH);