diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 5959124870..00da741932 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-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. @@ -574,6 +574,41 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res, cuvsCagraIndex_t index, bool include_dataset); +/** + * @defgroup cagra_c_serialize CAGRA C-API serialize functions + * @{ + */ +/** + * Save the index to file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.c} + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsError_t res_create_status = cuvsResourcesCreate(&res); + * + * // create an index with `cuvsCagraBuild` + * size_t length = ...; + * void* buffer = malloc(sizeof(char) * length); + * cuvsCagraSerializeToMemory(res, buffer, length, index, true); + * + * // read buffer + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] buffer pointer to a buffer + * @param[in/out] length the size of the buffer; the function will update it with the number of + * bytes written + * @param[in] index CAGRA index + * @param[in] include_dataset Whether or not to write out the dataset to the file. + * + */ +cuvsError_t cuvsCagraSerializeToMemory( + cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset); + /** * Save the CAGRA index to file in hnswlib format. * NOTE: The saved index can only be read by the hnswlib wrapper in cuVS, diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 656724826e..8d3deb4885 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -233,6 +233,25 @@ void _serialize(cuvsResources_t res, cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset); } +struct _membuf : std::streambuf { + _membuf(char* p, size_t size) { setp(p, p + size); } + size_t written() { return pptr() - pbase(); } +}; + +template +void _serialize( + cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index->addr); + + _membuf stream_buffer((char*)buffer, *length); + std::ostream out(&stream_buffer); + + cuvs::neighbors::cagra::serialize(*res_ptr, out, *index_ptr, include_dataset); + *length = stream_buffer.written(); +} + template void _serialize_to_hnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) { @@ -652,6 +671,24 @@ extern "C" cuvsError_t cuvsCagraSerialize(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsCagraSerializeToMemory( + cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _serialize(res, buffer, length, index, include_dataset); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + _serialize(res, buffer, length, index, include_dataset); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _serialize(res, buffer, length, index, include_dataset); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _serialize(res, buffer, length, index, include_dataset); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) diff --git a/java/.gitignore b/java/.gitignore index 2d3b14dd31..a240aca207 100644 --- a/java/.gitignore +++ b/java/.gitignore @@ -1,9 +1,8 @@ +*.iml +target/ +jextract-22/ +openjdk-22-jextract* # cuvs-java -/cuvs-java/target/ /cuvs-java/bin/ /cuvs-java/src/main/java22/com/nvidia/cuvs/internal/panama/ /cuvs-java/*.cag -jextract-22/ -openjdk-22-jextract* -# examples -/examples/target/ diff --git a/java/README.md b/java/README.md index b568a1ffa0..b913ebe709 100644 --- a/java/README.md +++ b/java/README.md @@ -25,10 +25,14 @@ do `./build.sh java` in the top level directory or just do `./build.sh` in this Run `./build.sh --run-java-tests` from this directory. -To run a single test: +To run a single test suite: ```sh cd cuvs-java/ -mvn verify -Dintegration-test=com.nvidia.cuvs.CagraBuildAndSearchIT +mvn clean integration-test -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT +``` +or, for a single test: +```sh +mvn clean integration-test -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT#testMergeStrategies ``` Be sure to set (manually, if needed) your `LD_LIBRARY_PATH` to include the directory with the appropriate (matching) version of `libcuvs.so`. diff --git a/java/benchmarks/README.md b/java/benchmarks/README.md new file mode 100644 index 0000000000..1740e52811 --- /dev/null +++ b/java/benchmarks/README.md @@ -0,0 +1,27 @@ +# CuVS Java API benchmarks + +This maven project contains JMH benchmarks for the CAGRA Java API. + +## Prerequisites +- [CuVS libraries](https://docs.rapids.ai/api/cuvs/stable/build/#build-from-source) +- Build the CuVS-Java API (`./build.sh` from the parent directory) + +## Run benchmarks + +Build: +```shell +mvn clean verify +``` +Run: +```shell +export RAFT_DEBUG_LOG_FILE=/dev/null +java -jar target/benchmarks.jar +``` +The environment variable is needed to silence RAFT logging; RAFT emits some logs at INFO level when +building indices and queries, and writing them to stdout (the default) influences benchmark results. + +It is possible to change the dataset size and the vectors dimension via 2 parameters: +```shell +java -jar target/benchmarks.jar -p size=4 -p dims=4 +``` +Use `java -jar target/benchmarks.jar -h` for details on the options to fine-tune your benchmark runs. diff --git a/java/benchmarks/pom.xml b/java/benchmarks/pom.xml new file mode 100644 index 0000000000..f5c77a14bc --- /dev/null +++ b/java/benchmarks/pom.xml @@ -0,0 +1,154 @@ + + + 4.0.0 + + com.nvidia.cuvs + benchmarks + 25.08.0 + jar + + cuvs-java-benchmarks + + + + org.openjdk.jmh + jmh-core + ${jmh.version} + + + org.openjdk.jmh + jmh-generator-annprocess + ${jmh.version} + provided + + + com.nvidia.cuvs + cuvs-java + 25.08.0 + jar + + + + + UTF-8 + 1.37 + 22 + benchmarks + + + + + + org.apache.maven.plugins + maven-compiler-plugin + 3.8.0 + + ${javac.target} + ${javac.target} + ${javac.target} + + -proc:full + + + + + org.apache.maven.plugins + maven-shade-plugin + 3.2.1 + + + package + + shade + + + ${uberjar.name} + + + org.openjdk.jmh.Main + + true + true + + + + + + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + + + + + + maven-clean-plugin + 2.5 + + + maven-deploy-plugin + 2.8.1 + + + maven-install-plugin + 2.5.1 + + + maven-jar-plugin + 2.4 + + + maven-javadoc-plugin + 2.9.1 + + + maven-resources-plugin + 2.6 + + + maven-site-plugin + 3.3 + + + maven-source-plugin + 2.2.1 + + + maven-surefire-plugin + 2.17 + + + + + + diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java new file mode 100644 index 0000000000..f1fab12e63 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java @@ -0,0 +1,127 @@ +/* + * 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; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.infra.Blackhole; + +import java.lang.foreign.*; +import java.util.concurrent.TimeUnit; + +import static com.nvidia.cuvs.Utils.createSampleData; +import static com.nvidia.cuvs.Utils.createSampleDataSegment; + +@Fork(value = 1, warmups = 0) +@State(Scope.Benchmark) +public class CagraIndexingBenchmarks { + + @Param({ "1024" }) + private int dims; + + @Param({ "100" }) + private int size; + + private float[][] arrayDataset; + + private Arena arena; + + private MemorySegment memorySegmentDataset; + + @Setup + public void initialize() { + arena = Arena.ofShared(); + arrayDataset = createSampleData(size, dims); + memorySegmentDataset = createSampleDataSegment(arena, arrayDataset, size, dims); + } + + @TearDown + public void cleanUp() { + if (arena != null) { + arena.close(); + } + } + + @Benchmark + public void testIndexingFromHeap(Blackhole blackhole) throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(arrayDataset) + .withIndexParams(indexParams) + .build(); + blackhole.consume(index); + } + } + + @Benchmark + public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) + .withIndexParams(indexParams) + .build(); + blackhole.consume(index); + } + } + + @Benchmark + @OutputTimeUnit(TimeUnit.NANOSECONDS) + @BenchmarkMode(Mode.AverageTime) + public void testDatasetFromHeap(Blackhole blackhole) throws Throwable { + try (var dataset = Dataset.ofArray(arrayDataset)) { + blackhole.consume(dataset); + } + } + + @Benchmark + @OutputTimeUnit(TimeUnit.NANOSECONDS) + @BenchmarkMode(Mode.AverageTime) + public void testDatasetFromMemorySegment(Blackhole blackhole) throws Throwable { + try (var dataset = Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) { + blackhole.consume(dataset); + } + } +} diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java new file mode 100644 index 0000000000..2056f56c14 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java @@ -0,0 +1,79 @@ +package com.nvidia.cuvs; + +import org.openjdk.jmh.annotations.*; + +import java.lang.foreign.Arena; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.UUID; + +import static com.nvidia.cuvs.Utils.createSampleData; + +@Fork(value = 1, warmups = 0) +@State(Scope.Benchmark) +public class CagraSerializationBenchmarks { + + @Param({ "1024" }) + private int dims; + + @Param({ "100" }) + private int size; + + private Arena arena; + + private CuVSResources resources; + + private CagraIndex index; + + @Setup + public void initialize() throws Throwable { + resources = CuVSResources.create(); + arena = Arena.ofShared(); + var arrayDataset = createSampleData(size, dims); + + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + index = CagraIndex.newBuilder(resources) + .withDataset(arrayDataset) + .withIndexParams(indexParams) + .build(); + } + + @TearDown + public void cleanUp() throws Throwable { + if (resources != null) { + resources.close(); + } + if (arena != null) { + arena.close(); + } + if (index != null) { + index.destroyIndex(); + } + } + + @Benchmark + public void testSerializeToFile() throws Throwable { + var indexFilePath = Path.of(UUID.randomUUID() + ".cag"); + try (var outputStream = Files.newOutputStream(indexFilePath)) { + index.serialize(outputStream); + } + Files.deleteIfExists(indexFilePath); + } + + @Benchmark + public void testSerializeToMemory() throws Throwable { + try (var arena = Arena.ofConfined()) { + var buffer = arena.allocate(1024 * 1024); + index.serialize((Object) buffer); + } + } +} diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java new file mode 100644 index 0000000000..4a0c1200e6 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java @@ -0,0 +1,31 @@ +package com.nvidia.cuvs; + +import java.lang.foreign.*; +import java.util.Random; + +class Utils { + private static final Random random = new Random(); + + static float[][] createSampleData(int size, int dimensions) { + var array = new float[size][dimensions]; + for (int i = 0; i < size; ++i) { + for (int j = 0; j < dimensions; ++j) { + array[i][j] = random.nextFloat(); + } + } + return array; + } + + static MemorySegment createSampleDataSegment(Arena arena, float[][] array, int size, int dimensions) { + final ValueLayout.OfFloat C_FLOAT = (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float"); + + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT); + + var segment = arena.allocate(dataMemoryLayout); + for (int i = 0; i < size; ++i) { + var vector = array[i]; + MemorySegment.copy(vector, 0, segment, C_FLOAT, (i * dimensions * C_FLOAT.byteSize()), dimensions); + } + return segment; + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index e878f77109..c868bd4cc4 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -97,6 +97,16 @@ default void serialize(OutputStream outputStream, Path tempFile) throws Throwabl */ void serialize(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable; + /** + * A method to persist a CAGRA index using a {@link java.lang.foreign.MemorySegment} + * for writing index bytes. + * + * @param memoryStream an instance of {@link java.lang.foreign.MemorySegment} to write the index + * bytes into + * @return a MemorySegment containing the serialized data + */ + Object serialize(Object memoryStream); + /** * A method to create and persist HNSW index from CAGRA index using an instance * of {@link OutputStream} and path to the intermediate temporary file. @@ -143,13 +153,6 @@ default void serializeToHNSW(OutputStream outputStream, Path tempFile) throws Th */ void serializeToHNSW(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable; - /** - * Gets an instance of {@link CagraIndexParams} - * - * @return an instance of {@link CagraIndexParams} - */ - CagraIndexParams getCagraIndexParameters(); - /** * Gets an instance of {@link CuVSResources} * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java index 7bc7dbd797..1b83cce5c5 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java @@ -27,12 +27,25 @@ */ public interface Dataset extends AutoCloseable { + /** + * Creates a dataset from a on-heap array of vectors + * + * @since 25.08 + */ + static Dataset ofArray(float[][] vectors) { + return CuVSProvider.provider().newArrayDataset(vectors); + } + + static Dataset ofMemorySegment(Object memorySegment, int size, int dimensions) { + return CuVSProvider.provider().newMemoryDataset(memorySegment, size, dimensions); + } + /** * Add a single vector to the dataset. * * @param vector A float array of as many elements as the dimensions */ - public void addVector(float[] vector); + void addVector(float[] vector); /** * Create a new instance of a dataset @@ -50,12 +63,12 @@ static Dataset create(int size, int dimensions) { * * @return Size of the dataset */ - public int size(); + int size(); /** * Gets the dimensions of the vectors in this dataset * * @return Dimensions of the vectors in the dataset */ - public int dimensions(); + int dimensions(); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 12c0a99902..dab80a3fdf 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -52,6 +52,11 @@ default Path nativeLibraryPath() { /** Create a {@link Dataset} instance **/ Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException; + Dataset newMemoryDataset(Object memorySegment, int size, int dimensions); + + /** Create a {@link Dataset} backed by a on-heap array **/ + Dataset newArrayDataset(float[][] vectors); + /** Creates a new BruteForceIndex Builder. */ BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index b9b5997905..a9312b2d53 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -56,4 +56,14 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException { throw new UnsupportedOperationException(); } + + @Override + public Dataset newMemoryDataset(Object memorySegment, int size, int dimensions) { + throw new UnsupportedOperationException(); + } + + @Override + public Dataset newArrayDataset(float[][] vectors) { + throw new UnsupportedOperationException(); + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index 362b946e3a..5cd5ba17c8 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -48,7 +48,6 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.SearchResults; -import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsBruteForceIndex; import com.nvidia.cuvs.internal.panama.cuvsFilter; import java.io.FileInputStream; @@ -74,8 +73,6 @@ */ public class BruteForceIndexImpl implements BruteForceIndex { - private final float[][] vectors; - private final Dataset dataset; private final CuVSResourcesImpl resources; private final IndexReference bruteForceIndexReference; private final BruteForceIndexParams bruteForceIndexParams; @@ -91,16 +88,15 @@ public class BruteForceIndexImpl implements BruteForceIndex { * holding the index parameters */ private BruteForceIndexImpl( - float[][] vectors, - Dataset dataset, - CuVSResourcesImpl resources, - BruteForceIndexParams bruteForceIndexParams) - throws Throwable { - this.vectors = vectors; - this.dataset = dataset; - this.resources = resources; - this.bruteForceIndexParams = bruteForceIndexParams; - this.bruteForceIndexReference = build(); + Dataset dataset, CuVSResourcesImpl resources, BruteForceIndexParams bruteForceIndexParams) + throws Exception { + Objects.requireNonNull(dataset); + try (dataset) { + this.resources = resources; + this.bruteForceIndexParams = bruteForceIndexParams; + assert dataset instanceof DatasetImpl; + this.bruteForceIndexReference = build((DatasetImpl) dataset); + } } /** @@ -112,8 +108,6 @@ private BruteForceIndexImpl( private BruteForceIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { this.bruteForceIndexParams = null; - this.vectors = null; - this.dataset = null; this.resources = resources; this.bruteForceIndexReference = deserialize(inputStream); } @@ -137,7 +131,6 @@ public void destroyIndex() throws Throwable { } finally { destroyed = true; } - if (dataset != null) dataset.close(); } /** @@ -147,16 +140,13 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() throws Throwable { + private IndexReference build(DatasetImpl dataset) { try (var localArena = Arena.ofConfined()) { - long rows = dataset != null ? dataset.size() : vectors.length; - long cols = dataset != null ? dataset.dimensions() : (rows > 0 ? vectors[0].length : 0); + long rows = dataset.size(); + long cols = dataset.dimensions(); Arena arena = resources.getArena(); - MemorySegment datasetMemSegment = - dataset != null - ? ((DatasetImpl) dataset).seg - : Util.buildMemorySegment(resources.getArena(), vectors); + MemorySegment datasetMemSegment = dataset.asMemorySegment(); long cuvsResources = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); @@ -176,7 +166,7 @@ private IndexReference build() throws Throwable { cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION); - long datasetShape[] = {rows, cols}; + long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = prepareTensor(arena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 2, 1); @@ -228,9 +218,9 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { BitSet[] prefilters = cuvsQuery.getPrefilters(); if (prefilters != null && prefilters.length > 0) { BitSet concatenatedFilters = concatenate(prefilters, cuvsQuery.getNumDocs()); - long filters[] = concatenatedFilters.toLongArray(); + long[] filters = concatenatedFilters.toLongArray(); prefilterDataMemorySegment = buildMemorySegment(arena, filters); - prefilterDataLength = cuvsQuery.getNumDocs() * prefilters.length; + prefilterDataLength = (long) cuvsQuery.getNumDocs() * prefilters.length; } MemorySegment querySeg = buildMemorySegment(arena, cuvsQuery.getQueryVectors()); @@ -411,7 +401,6 @@ public static BruteForceIndex.Builder newBuilder(CuVSResources cuvsResources) { */ public static class Builder implements BruteForceIndex.Builder { - private float[][] vectors; private Dataset dataset; private final CuVSResourcesImpl cuvsResources; private BruteForceIndexParams bruteForceIndexParams; @@ -460,7 +449,7 @@ public Builder from(InputStream inputStream) { */ @Override public Builder withDataset(float[][] vectors) { - this.vectors = vectors; + this.dataset = Dataset.ofArray(vectors); return this; } @@ -486,7 +475,7 @@ public BruteForceIndexImpl build() throws Throwable { if (inputStream != null) { return new BruteForceIndexImpl(inputStream, cuvsResources); } else { - return new BruteForceIndexImpl(vectors, dataset, cuvsResources, bruteForceIndexParams); + return new BruteForceIndexImpl(dataset, cuvsResources, bruteForceIndexParams); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 43bb92ff5c..54197e1fe9 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -38,6 +38,7 @@ import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSearch; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSerialize; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSerializeToHnswlib; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSerializeToMemory; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResources_t; @@ -57,7 +58,6 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.SearchResults; -import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams; import com.nvidia.cuvs.internal.panama.cuvsCagraIndex; import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; @@ -95,10 +95,7 @@ * @since 25.02 */ public class CagraIndexImpl implements CagraIndex { - private final float[][] vectors; - private final Dataset dataset; private final CuVSResourcesImpl resources; - private final CagraIndexParams cagraIndexParameters; private final IndexReference cagraIndexReference; private boolean destroyed; @@ -111,16 +108,14 @@ public class CagraIndexImpl implements CagraIndex { * @param resources an instance of {@link CuVSResources} */ private CagraIndexImpl( - CagraIndexParams indexParameters, - float[][] vectors, - Dataset dataset, - CuVSResourcesImpl resources) - throws Throwable { - this.cagraIndexParameters = indexParameters; - this.vectors = vectors; - this.dataset = dataset; - this.resources = resources; - this.cagraIndexReference = build(); + CagraIndexParams indexParameters, Dataset dataset, CuVSResourcesImpl resources) + throws Exception { + Objects.requireNonNull(dataset); + try (dataset) { + this.resources = resources; + assert dataset instanceof DatasetImpl; + this.cagraIndexReference = build(indexParameters, (DatasetImpl) dataset); + } } /** @@ -130,9 +125,6 @@ private CagraIndexImpl( * @param resources an instance of {@link CuVSResources} */ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { - this.cagraIndexParameters = null; - this.vectors = null; - this.dataset = null; this.resources = resources; this.cagraIndexReference = deserialize(inputStream); } @@ -145,9 +137,6 @@ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) thr * @param resources The resources instance */ private CagraIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) { - this.vectors = null; - this.cagraIndexParameters = null; - this.dataset = null; this.resources = resources; this.cagraIndexReference = indexReference; this.destroyed = false; @@ -163,7 +152,7 @@ private void checkNotDestroyed() { * Invokes the native destroy_cagra_index to de-allocate the CAGRA index */ @Override - public void destroyIndex() throws Throwable { + public void destroyIndex() { checkNotDestroyed(); try (var arena = Arena.ofConfined()) { int returnValue = cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()); @@ -171,7 +160,6 @@ public void destroyIndex() throws Throwable { } finally { destroyed = true; } - if (dataset != null) dataset.close(); } /** @@ -181,26 +169,22 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() throws Throwable { + private IndexReference build(CagraIndexParams indexParameters, DatasetImpl dataset) { try (var localArena = Arena.ofConfined()) { - long rows = dataset != null ? dataset.size() : vectors.length; - long cols = dataset != null ? dataset.dimensions() : (rows > 0 ? vectors[0].length : 0); + long rows = dataset.size(); + long cols = dataset.dimensions(); MemorySegment indexParamsMemorySegment = - cagraIndexParameters != null - ? segmentFromIndexParams(resources, cagraIndexParameters) + indexParameters != null + ? segmentFromIndexParams(resources, indexParameters) : MemorySegment.NULL; - int numWriterThreads = - cagraIndexParameters != null ? cagraIndexParameters.getNumWriterThreads() : 1; + int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; omp_set_num_threads(numWriterThreads); - MemorySegment dataSeg = - dataset != null - ? ((DatasetImpl) dataset).seg - : Util.buildMemorySegment(resources.getArena(), vectors); - Arena arena = resources.getArena(); + MemorySegment dataSeg = dataset.asMemorySegment(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); var returnValue = cuvsStreamGet(cuvsRes, stream); @@ -386,6 +370,22 @@ public SearchResults search(CagraQuery query) throws Throwable { } } + @Override + public Object serialize(Object memorySegment) { + assert memorySegment instanceof MemorySegment; + var buffer = (MemorySegment) memorySegment; + checkNotDestroyed(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + try (var arena = Arena.ofConfined()) { + var length = arena.allocateFrom(ValueLayout.JAVA_LONG, buffer.byteSize()); + var returnValue = + cuvsCagraSerializeToMemory( + cuvsRes, buffer, length, cagraIndexReference.getMemorySegment(), true); + checkCuVSError(returnValue, "cuvsCagraSerializeToMemory"); + return buffer.reinterpret(length.get(ValueLayout.JAVA_LONG, 0)); + } + } + @Override public void serialize(OutputStream outputStream) throws Throwable { Path path = @@ -509,16 +509,6 @@ private IndexReference deserialize(InputStream inputStream, int bufferLength) th return indexReference; } - /** - * Gets an instance of {@link CagraIndexParams} - * - * @return an instance of {@link CagraIndexParams} - */ - @Override - public CagraIndexParams getCagraIndexParameters() { - return cagraIndexParameters; - } - /** * Gets an instance of {@link CuVSResources} * @@ -751,7 +741,6 @@ private static MemorySegment createMergeParamsSegment( */ public static class Builder implements CagraIndex.Builder { - private float[][] vectors; private Dataset dataset; private CagraIndexParams cagraIndexParams; private CuVSResourcesImpl cuvsResources; @@ -769,7 +758,7 @@ public Builder from(InputStream inputStream) { @Override public Builder withDataset(float[][] vectors) { - this.vectors = vectors; + this.dataset = Dataset.ofArray(vectors); return this; } @@ -790,11 +779,7 @@ public CagraIndexImpl build() throws Throwable { if (inputStream != null) { return new CagraIndexImpl(inputStream, cuvsResources); } else { - if (vectors != null && dataset != null) { - throw new IllegalArgumentException( - "Please specify only one type of dataset (a float[] or a Dataset instance)"); - } - return new CagraIndexImpl(cagraIndexParams, vectors, dataset, cuvsResources); + return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java index 5e26dcd73a..162eb9407e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java @@ -33,22 +33,29 @@ public DatasetImpl(int size, int dimensions) { this.size = size; this.dimensions = dimensions; - MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout(size * dimensions, C_FLOAT); + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long) size * dimensions, C_FLOAT); this.arena = Arena.ofShared(); - seg = arena.allocate(dataMemoryLayout); + this.seg = arena.allocate(dataMemoryLayout); + } + + public DatasetImpl(Arena arena, MemorySegment memorySegment, int size, int dimensions) { + this.arena = arena; + this.seg = memorySegment; + this.size = size; + this.dimensions = dimensions; } @Override public void addVector(float[] vector) { if (current >= size) throw new ArrayIndexOutOfBoundsException(); MemorySegment.copy( - vector, 0, seg, C_FLOAT, ((current++) * dimensions * C_FLOAT.byteSize()), (int) dimensions); + vector, 0, seg, C_FLOAT, ((current++) * dimensions * C_FLOAT.byteSize()), dimensions); } @Override public void close() { - if (!arena.scope().isAlive()) { + if (arena != null) { arena.close(); } } @@ -62,4 +69,8 @@ public int size() { public int dimensions() { return dimensions; } + + public MemorySegment asMemorySegment() { + return seg; + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index d5a348e1fc..67b12aac42 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -26,6 +26,9 @@ import com.nvidia.cuvs.internal.CuVSResourcesImpl; import com.nvidia.cuvs.internal.DatasetImpl; import com.nvidia.cuvs.internal.HnswIndexImpl; +import com.nvidia.cuvs.internal.common.Util; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; import java.nio.file.Files; import java.nio.file.Path; import java.util.Objects; @@ -80,4 +83,23 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merge public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException { return new DatasetImpl(size, dimensions); } + + @Override + public Dataset newMemoryDataset(Object memorySegment, int size, int dimensions) { + return new DatasetImpl(null, (MemorySegment) memorySegment, size, dimensions); + } + + @Override + public Dataset newArrayDataset(float[][] vectors) { + Objects.requireNonNull(vectors); + if (vectors.length == 0) { + throw new IllegalArgumentException("vectors should not be empty"); + } + int size = vectors.length; + int dimensions = vectors[0].length; + + Arena arena = Arena.ofShared(); + var memorySegment = Util.buildMemorySegment(arena, vectors); + return new DatasetImpl(arena, memorySegment, size, dimensions); + } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index 54e3815ef3..9e9a9e0919 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -16,6 +16,7 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.fail; @@ -28,7 +29,12 @@ import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.InputStream; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; import java.lang.invoke.MethodHandles; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.BitSet; import java.util.List; @@ -70,6 +76,7 @@ private static void runConcurrently(int nThreads, Supplier runnableSup CompletableFuture.allOf(futures) .exceptionally( t -> { + log.error("Exception while executing runnable", t); fail("Exception while executing runnable: " + t); return null; }) @@ -238,6 +245,55 @@ public void testPrefilteringReducesResults() throws Throwable { } } + @Test + public void testIndexSerialization() throws Throwable { + float[][] dataset = createSampleData(); + + try (CuVSResources resources = CuVSResources.create()) { + + // Configure index parameters + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + var indexFile = Path.of(UUID.randomUUID() + ".cag"); + try { + // Saving the index on to the disk. + try (var output = Files.newOutputStream(indexFile)) { + index.serialize(output); + } + + var indexFromFile = Files.readAllBytes(indexFile); + byte[] indexFromMemory; + try (var arena = Arena.ofConfined()) { + var buffer = arena.allocate(1024 * 1024); + var serialized = (MemorySegment) index.serialize((Object) buffer); + indexFromMemory = serialized.toArray(ValueLayout.JAVA_BYTE); + } + + assertNotNull(indexFromFile); + assertNotNull(indexFromMemory); + assertArrayEquals(indexFromFile, indexFromMemory); + + } finally { + Files.deleteIfExists(indexFile); + index.destroyIndex(); + } + } + } + private void indexAndQueryOnce( float[][] dataset, List map,