Merge CAGRA indexes on the GPU instead of rebuilding from vectors

mergeOneField now hands the segments' CAGRA indexes to the cuVS merge API,
falling back to the vector based merge when the merged rows cannot be lined
up with the merged flat vectors. Deleted documents are dropped through a
row, which needs a new CagraIndex.merge overload taking a BitSet. Also
cleans up some merge overloads on CuVSProvider and CagraIndex.

Closes #1078

Author: Igor Motov

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);
   }