Merge CAGRA indexes on the GPU instead of rebuilding from vectors

java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java

@@ -8,6 +8,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.file.Path;
+import java.util.BitSet;
 import java.util.Objects;
 
 /**
@@ -192,6 +193,13 @@ public StandardDataset() {}
    */
   long getGraphDegree();
 
+  /**
+   * Returns the number of vectors in this index.
+   *
+   * @return the number of rows of the indexed dataset
+   */
+  long size();
+
   /**
    * A method to persist a CAGRA index using an instance of {@link OutputStream}
    * for writing index bytes.
@@ -310,7 +318,7 @@ static Builder newBuilder(CuVSResources cuvsResources) {
    * @throws Throwable if an error occurs during the merge operation
    */
   static CagraIndex merge(CagraIndex[] indexes) throws Throwable {
-    return merge(indexes, null);
+    return merge(indexes, null, null);
   }
 
   /**
@@ -322,6 +330,28 @@ static CagraIndex merge(CagraIndex[] indexes) throws Throwable {
    * @throws Throwable if an error occurs during the merge operation
    */
   static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) throws Throwable {
+    return merge(indexes, mergeParams, null);
+  }
+
+  /**
+   * Merges multiple CAGRA indexes into a single index, keeping only the rows selected by
+   * {@code rowFilter}.
+   *
+   * <p>The merge concatenates the input datasets in the order the indexes are given, so bit
+   * {@code i} of the filter refers to row {@code i} of that concatenation: bits {@code 0} to
+   * {@code indexes[0].size() - 1} address the first index, the bits that follow address the second,
+   * and so on. A <b>set</b> bit keeps the row; a clear bit drops it. The rows that survive keep
+   * their relative order and are packed together, so the merged index has one row per set bit.
+   *
+   * @param indexes Array of CAGRA indexes to merge
+   * @param mergeParams Parameters to control the merge operation, or null to use defaults
+   * @param rowFilter The rows to keep, or null to keep all of them
+   * @return A new merged CAGRA index
+   * @throws IllegalArgumentException if {@code rowFilter} has a bit set beyond the last row
+   * @throws Throwable if an error occurs during the merge operation
+   */
+  static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams, BitSet rowFilter)
+      throws Throwable {
     if (indexes == null || indexes.length == 0) {
       throw new IllegalArgumentException("At least one index must be provided for merging");
     }
@@ -333,7 +363,7 @@ static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) thro
       }
     }
 
-    return CuVSProvider.provider().mergeCagraIndexes(indexes, mergeParams);
+    return CuVSProvider.provider().mergeCagraIndexes(indexes, mergeParams, rowFilter);
   }
 
   /**

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java

@@ -9,6 +9,7 @@
 import java.lang.invoke.MethodType;
 import java.nio.file.Path;
 import java.time.Duration;
+import java.util.BitSet;
 import java.util.List;
 
 /**
@@ -165,29 +166,6 @@ HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, Cu
   TieredIndex.Builder newTieredIndexBuilder(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, CagraIndexParams mergeParams)
-      throws Throwable {
-    // Default implementation falls back to the method without parameters
-    return mergeCagraIndexes(indexes);
-  }
-
   /**
    * Reports whether the rows of {@code dataset} already sit at the row stride CAGRA requires, which
    * is the row length in bytes rounded up to a 16 byte boundary.
@@ -207,6 +185,20 @@ default boolean isCagraPaddedDataset(CuVSMatrix dataset) {
         "Padded layout detection is not supported by " + getClass().getName());
   }
 
+  /**
+   * Merges multiple CAGRA indexes into a single index, keeping only the rows selected by
+   * {@code rowFilter}. See {@link CagraIndex#merge(CagraIndex[], CagraIndexParams, BitSet)} for the
+   * meaning of the filter.
+   *
+   * @param indexes Array of CAGRA indexes to merge
+   * @param mergeParams Parameters to control the merge operation, or null to use defaults
+   * @param rowFilter The rows to keep, or null to keep all of them
+   * @return A new merged CAGRA index
+   * @throws Throwable if an error occurs during the merge operation
+   */
+  CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams mergeParams, BitSet rowFilter)
+      throws Throwable;
+
   /**
    * Creates a device-backed multi-partition filter handle from the pre-packed combined bitset.
    * Per-partition bit offsets are recomputed inside cuVS from the index sizes.

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java

@@ -8,6 +8,7 @@
 import java.lang.invoke.MethodHandle;
 import java.nio.file.Path;
 import java.time.Duration;
+import java.util.BitSet;
 import java.util.List;
 import java.util.logging.Level;
 
@@ -81,12 +82,13 @@ public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) {
   }
 
   @Override
-  public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) {
+  public boolean isCagraPaddedDataset(CuVSMatrix dataset) {
     throw new UnsupportedOperationException(reasons);
   }
 
   @Override
-  public boolean isCagraPaddedDataset(CuVSMatrix dataset) {
+  public CagraIndex mergeCagraIndexes(
+      CagraIndex[] indexes, CagraIndexParams mergeParams, BitSet rowFilter) {
     throw new UnsupportedOperationException(reasons);
   }