[WIP][Java] Use MemorySegment for index serialization

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 <cuvs/neighbors/cagra.h>
+ *
+ * // 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,

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 <typename T>
+void _serialize(
+  cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(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 <typename T>
 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<float>(res, buffer, length, index, include_dataset);
+    } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+      _serialize<half>(res, buffer, length, index, include_dataset);
+    } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+      _serialize<int8_t>(res, buffer, length, index, include_dataset);
+    } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+      _serialize<uint8_t>(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)

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/

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`.

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.

java/benchmarks/pom.xml

@@ -0,0 +1,154 @@
+<!--
+/*
+ * 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.
+ */
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>com.nvidia.cuvs</groupId>
+    <artifactId>benchmarks</artifactId>
+    <!--CUVS_JAVA#VERSION_UPDATE_MARKER_START--><version>25.08.0</version><!--CUVS_JAVA#VERSION_UPDATE_MARKER_END-->
+    <packaging>jar</packaging>
+
+    <name>cuvs-java-benchmarks</name>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.openjdk.jmh</groupId>
+            <artifactId>jmh-core</artifactId>
+            <version>${jmh.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.openjdk.jmh</groupId>
+            <artifactId>jmh-generator-annprocess</artifactId>
+            <version>${jmh.version}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.nvidia.cuvs</groupId>
+            <artifactId>cuvs-java</artifactId>
+            <!--CUVS_JAVA#VERSION_UPDATE_MARKER_START--><version>25.08.0</version><!--CUVS_JAVA#VERSION_UPDATE_MARKER_END-->
+            <type>jar</type>
+        </dependency>
+    </dependencies>
+
+    <properties>
+        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+        <jmh.version>1.37</jmh.version>
+        <javac.target>22</javac.target>
+        <uberjar.name>benchmarks</uberjar.name>
+    </properties>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <version>3.8.0</version>
+                <configuration>
+                    <compilerVersion>${javac.target}</compilerVersion>
+                    <source>${javac.target}</source>
+                    <target>${javac.target}</target>
+                    <compilerArgs>
+                        <arg>-proc:full</arg>
+                    </compilerArgs>
+                </configuration>
+            </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <version>3.2.1</version>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>shade</goal>
+                        </goals>
+                        <configuration>
+                            <finalName>${uberjar.name}</finalName>
+                            <transformers>
+                                <transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
+                                    <mainClass>org.openjdk.jmh.Main</mainClass>
+                                    <manifestEntries>
+                                        <Multi-Release>true</Multi-Release>
+                                        <addClasspath>true</addClasspath>
+                                    </manifestEntries>
+                                </transformer>
+                                <transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+                            </transformers>
+                            <filters>
+                                <filter>
+                                    <!--
+                                        Shading signed JARs will fail without this.
+                                        http://stackoverflow.com/questions/999489/invalid-signature-file-when-attempting-to-run-a-jar
+                                    -->
+                                    <artifact>*:*</artifact>
+                                    <excludes>
+                                        <exclude>META-INF/*.SF</exclude>
+                                        <exclude>META-INF/*.DSA</exclude>
+                                        <exclude>META-INF/*.RSA</exclude>
+                                    </excludes>
+                                </filter>
+                            </filters>
+                        </configuration>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+        <pluginManagement>
+            <plugins>
+                <plugin>
+                    <artifactId>maven-clean-plugin</artifactId>
+                    <version>2.5</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-deploy-plugin</artifactId>
+                    <version>2.8.1</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-install-plugin</artifactId>
+                    <version>2.5.1</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-jar-plugin</artifactId>
+                    <version>2.4</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-javadoc-plugin</artifactId>
+                    <version>2.9.1</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-resources-plugin</artifactId>
+                    <version>2.6</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-site-plugin</artifactId>
+                    <version>3.3</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-source-plugin</artifactId>
+                    <version>2.2.1</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-surefire-plugin</artifactId>
+                    <version>2.17</version>
+                </plugin>
+            </plugins>
+        </pluginManagement>
+    </build>
+
+</project>