rapidsai · rapids-bot · May 29, 2025 · May 21, 2025 · May 21, 2025 · May 22, 2025
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
@@ -170,6 +170,40 @@ static Builder newBuilder(CuVSResources cuvsResources) {
         return CuVSProvider.provider().newCagraIndexBuilder(cuvsResources);
     }
 
+    /**
+     * Merges multiple CAGRA indexes into a single index using default merge parameters.
+     *
+     * @param indexes Array of CAGRA indexes to merge
+     * @return A new merged CAGRA index
+     * @throws Throwable if an error occurs during the merge operation
+     */
+    static CagraIndex merge(CagraIndex[] indexes) throws Throwable {
+        return merge(indexes, null);
+    }
+
+    /**
+     * Merges multiple CAGRA indexes into a single index with the specified merge parameters.
+     *
+     * @param indexes Array of CAGRA indexes to merge
+     * @param mergeParams Parameters to control the merge operation, or null to use defaults
+     * @return A new merged CAGRA index
+     * @throws Throwable if an error occurs during the merge operation
+     */
+    static CagraIndex merge(CagraIndex[] indexes, CagraMergeParams mergeParams) throws Throwable {
+        if (indexes == null || indexes.length == 0) {
+            throw new IllegalArgumentException("At least one index must be provided for merging");
+        }
+
+        CuVSResources resources = indexes[0].getCuVSResources();
+        for (int i = 1; i < indexes.length; i++) {
+            if (!resources.equals(indexes[i].getCuVSResources())) {
+                throw new IllegalArgumentException("All indexes must use the same CuVSResources instance");
+            }
+        }
+
+        return CuVSProvider.provider().mergeCagraIndexes(indexes, mergeParams);
+    }
+
     /**
      * Builder helps configure and create an instance of {@link CagraIndex}.
      */

diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraMergeParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraMergeParams.java
@@ -0,0 +1,105 @@
+/*
+ * 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;
+public class CagraMergeParams {
+
+    private final CagraIndexParams outputIndexParams;
+    private final MergeStrategy strategy;
+
+    /**
+     * Constructs a CagraMergeParams with the given output index parameters and merge strategy.
+     *
+     * @param outputIndexParams Index parameters for the output index
+     * @param strategy Merge strategy to use
+     */
+    private CagraMergeParams(CagraIndexParams outputIndexParams, MergeStrategy strategy) {
+        this.outputIndexParams = outputIndexParams;
+        this.strategy = strategy;
+    }
+
+    /**
+     * Gets the index parameters for the output index.
+     *
+     * @return Index parameters to use for the output index
+     */
+    public CagraIndexParams getOutputIndexParams() {
+        return outputIndexParams;
+    }
+
+    /**
+     * Gets the merge strategy to use.
+     *
+     * @return The merge strategy
+     */
+    public MergeStrategy getStrategy() {
+        return strategy;
+    }
+
+    /**
+     * Strategy to use when merging CAGRA indexes.
+     */
+    public enum MergeStrategy {
+        PHYSICAL(0),
+
+        LOGICAL(1);
+
+        public final int value;
+
+        MergeStrategy(int value) {
+            this.value = value;
+        }
+    }
+
+    /**
+     * Builder class for {@link CagraMergeParams}.
+     */
+    public static class Builder {
+        private CagraIndexParams outputIndexParams = new CagraIndexParams.Builder().build();
+        private MergeStrategy strategy = MergeStrategy.PHYSICAL; // Default to PHYSICAL
+
+        /**
+         * Sets the index parameters for the output index.
+         *
+         * @param outputIndexParams Index parameters to use for the output index
+         * @return This builder
+         */
+        public Builder withOutputIndexParams(CagraIndexParams outputIndexParams) {
+            this.outputIndexParams = outputIndexParams;
+            return this;
+        }
+
+        /**
+         * Sets the merge strategy.
+         *
+         * @param strategy The merge strategy to use
+         * @return This builder
+         */
+        public Builder withStrategy(MergeStrategy strategy) {
+            this.strategy = strategy;
+            return this;
+        }
+
+        /**
+         * Builds the {@link CagraMergeParams} object.
+         *
+         * @return The built parameters
+         */
+        public CagraMergeParams build() {
+            return new CagraMergeParams(outputIndexParams, strategy);
+        }
+    }
+}
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
@@ -21,6 +21,7 @@
 import com.nvidia.cuvs.CuVSResources;
 import com.nvidia.cuvs.Dataset;
 import com.nvidia.cuvs.HnswIndex;
+import com.nvidia.cuvs.CagraMergeParams;
 
 import java.nio.file.Path;
 
@@ -66,6 +67,28 @@ CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources)
   HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources)
       throws UnsupportedOperationException;
 
+  /**
+   * Merges multiple CAGRA indexes into a single index.
+   *
+   * @param indexes Array of CAGRA indexes to merge
+   * @return A new merged CAGRA index
+   * @throws Throwable if an error occurs during the merge operation
+   */
+  CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable;
+
+  /**
+   * Merges multiple CAGRA indexes into a single index with the specified merge parameters.
+   *
+   * @param indexes Array of CAGRA indexes to merge
+   * @param mergeParams Parameters to control the merge operation, or null to use defaults
+   * @return A new merged CAGRA index
+   * @throws Throwable if an error occurs during the merge operation
+   */
+  default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams mergeParams) throws Throwable {
+    // Default implementation falls back to the method without parameters
+    return mergeCagraIndexes(indexes);
+  }
+
   /** Retrieves the system-wide provider. */
   static CuVSProvider provider() {
     return CuVSServiceProvider.Holder.INSTANCE;

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
@@ -49,6 +49,11 @@ public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException {
     throw new UnsupportedOperationException();

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
@@ -42,6 +42,7 @@
 import com.nvidia.cuvs.CagraIndex;
 import com.nvidia.cuvs.CagraIndexParams;
 import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo;
+import com.nvidia.cuvs.CagraMergeParams;
 import com.nvidia.cuvs.CagraQuery;
 import com.nvidia.cuvs.CagraSearchParams;
 import com.nvidia.cuvs.CuVSResources;
@@ -51,6 +52,7 @@
 import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams;
 import com.nvidia.cuvs.internal.panama.cuvsCagraIndex;
 import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams;
+import com.nvidia.cuvs.internal.panama.cuvsCagraMergeParams;
 import com.nvidia.cuvs.internal.panama.cuvsCagraSearchParams;
 import com.nvidia.cuvs.internal.panama.cuvsIvfPqIndexParams;
 import com.nvidia.cuvs.internal.panama.cuvsIvfPqParams;
@@ -88,6 +90,10 @@ public class CagraIndexImpl implements CagraIndex {
   private static final MethodHandle serializeCAGRAIndexToHNSWMethodHandle = downcallHandle("serialize_cagra_index_to_hnsw",
       FunctionDescriptor.ofVoid(ADDRESS, ADDRESS, ADDRESS, ADDRESS));
 
+
+  private static final MethodHandle mergeMethodHandle = downcallHandle("merge_cagra_indexes",
+      FunctionDescriptor.ofVoid(ADDRESS, ADDRESS, ADDRESS, C_INT, ADDRESS, ADDRESS));
+
   private final float[][] vectors;
   private final Dataset dataset;
   private final CuVSResourcesImpl resources;
@@ -131,6 +137,23 @@ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) thr
     this.cagraIndexReference = deserialize(inputStream);
   }
 
+  /**
+   * Constructor for creating an index from an existing index reference.
+   * Used primarily for the merge operation.
+   *
+   * @param indexReference The reference to the existing index
+   * @param resources The resources instance
+   */
+  private CagraIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) {
+    this.vectors = null;
+    this.cagraIndexParameters = null;
+    this.cagraCompressionParams = null;
+    this.dataset = null;
+    this.resources = resources;
+    this.cagraIndexReference = indexReference;
+    this.destroyed = false;
+  }
+
   private void checkNotDestroyed() {
     if (destroyed) {
       throw new IllegalStateException("destroyed");
@@ -165,7 +188,7 @@ private IndexReference build() throws Throwable {
     long cols = dataset != null? dataset.dimensions(): (rows > 0 ? vectors[0].length : 0);
 
     MemorySegment indexParamsMemorySegment = cagraIndexParameters != null
-        ? segmentFromIndexParams(cagraIndexParameters)
+        ? segmentFromIndexParams(resources, cagraIndexParameters)
         : MemorySegment.NULL;
 
     int numWriterThreads = cagraIndexParameters != null ? cagraIndexParameters.getNumWriterThreads() : 1;
@@ -397,7 +420,7 @@ private MemorySegment segmentFromCompressionParams(CagraCompressionParams params
   /**
    * Allocates the configured index parameters in the MemorySegment.
    */
-  private MemorySegment segmentFromIndexParams(CagraIndexParams params) {
+  private static MemorySegment segmentFromIndexParams(CuVSResourcesImpl resources, CagraIndexParams params) {
     MemorySegment seg = cuvsCagraIndexParams.allocate(resources.getArena());
     cuvsCagraIndexParams.intermediate_graph_degree(seg, params.getIntermediateGraphDegree());
     cuvsCagraIndexParams.graph_degree(seg, params.getGraphDegree());
@@ -471,6 +494,80 @@ public static CagraIndex.Builder newBuilder(CuVSResources cuvsResources) {
     return new CagraIndexImpl.Builder((CuVSResourcesImpl)cuvsResources);
   }
 
+  /**
+   * Merges multiple CAGRA indexes into a single index.
+   *
+   * @param indexes Array of CAGRA indexes to merge
+   * @return A new merged CAGRA index
+   * @throws Throwable if an error occurs during the merge operation
+   */
+  public static CagraIndex merge(CagraIndex[] indexes) throws Throwable {
+    return merge(indexes, null);
+  }
+
+  /**
+   * Merges multiple CAGRA indexes into a single index with specified merge parameters.
+   *
+   * @param indexes Array of CAGRA indexes to merge
+   * @param mergeParams Parameters to control the merge operation, or null to use defaults
+   * @return A new merged CAGRA index
+   * @throws Throwable if an error occurs during the merge operation
+   */
+  public static CagraIndex merge(CagraIndex[] indexes, CagraMergeParams mergeParams) throws Throwable {
+    CuVSResourcesImpl resources = (CuVSResourcesImpl) indexes[0].getCuVSResources();
+    IndexReference mergedIndexReference = new IndexReference(resources);
+
+    try (var arena = Arena.ofConfined()) {
+      MemorySegment indexesSegment = arena.allocate(indexes.length * ADDRESS.byteSize());
+      for (int i = 0; i < indexes.length; i++) {
+        CagraIndexImpl indexImpl = (CagraIndexImpl) indexes[i];
+        indexesSegment.setAtIndex(ADDRESS, i, indexImpl.cagraIndexReference.getMemorySegment());
+      }
+
+      MemorySegment returnValue = arena.allocate(C_INT);
+
+      MemorySegment mergeParamsSegment = MemorySegment.NULL;
+      if (mergeParams != null) {
+        mergeParamsSegment = createMergeParamsSegment(mergeParams, resources);
+      }
+
+      mergeMethodHandle.invokeExact(
+        resources.getMemorySegment(),
+        indexesSegment,
+        mergedIndexReference.getMemorySegment(),
+        indexes.length,
+        returnValue,
+        mergeParamsSegment
+      );
+
+      checkError(returnValue.get(C_INT, 0L), "mergeMethodHandle");
+    }
+
+    return new CagraIndexImpl(mergedIndexReference, resources);
+  }
+
+  /**
+   * Creates a memory segment for merge parameters.
+   *
+   * @param mergeParams The merge parameters
+   * @param resources The CuVS resources
+   * @return A memory segment with the merge parameters
+   */
+  private static MemorySegment createMergeParamsSegment(CagraMergeParams mergeParams, CuVSResourcesImpl resources) {
+    MemorySegment seg = cuvsCagraMergeParams.allocate(resources.getArena());
+
+    if (mergeParams.getOutputIndexParams() != null) {
+      MemorySegment outputIndexParamsSeg = segmentFromIndexParams(resources, mergeParams.getOutputIndexParams());
+      cuvsCagraMergeParams.output_index_params(seg, outputIndexParamsSeg);
+    } else {
+      cuvsCagraMergeParams.output_index_params(seg, MemorySegment.NULL);
+    }
+
+    cuvsCagraMergeParams.strategy(seg, mergeParams.getStrategy().value);
+
+    return seg;
+  }
+
   /**
    * Builder helps configure and create an instance of {@link CagraIndex}.
    */
@@ -533,7 +630,7 @@ public CagraIndexImpl build() throws Throwable {
   /**
    * Holds the memory reference to a CAGRA index.
    */
-  protected static class IndexReference {
+  public static class IndexReference {
 
     private final MemorySegment memorySegment;
 

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
@@ -21,6 +21,7 @@
 import com.nvidia.cuvs.CuVSResources;
 import com.nvidia.cuvs.Dataset;
 import com.nvidia.cuvs.HnswIndex;
+import com.nvidia.cuvs.CagraMergeParams;
 import com.nvidia.cuvs.internal.BruteForceIndexImpl;
 import com.nvidia.cuvs.internal.CagraIndexImpl;
 import com.nvidia.cuvs.internal.CuVSResourcesImpl;
@@ -60,6 +61,22 @@ public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) {
     return HnswIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources));
   }
 
+  @Override
+  public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable {
+    if (indexes == null || indexes.length == 0) {
+      throw new IllegalArgumentException("At least one index must be provided for merging");
+    }
+    return CagraIndexImpl.merge(indexes);
+  }
+
+  @Override
+  public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams mergeParams) throws Throwable {
+    if (indexes == null || indexes.length == 0) {
+      throw new IllegalArgumentException("At least one index must be provided for merging");
+    }
+    return CagraIndexImpl.merge(indexes, mergeParams);
+  }
+
   @Override
   public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException {
       return new DatasetImpl(size, dimensions);