Fix Tests/CI, refactor merge to call CAGRA's merge(), implement CAGRA prefiltering

ci/build_java.sh

@@ -32,6 +32,27 @@ set -u
 
 rapids-print-env
 
+# Locates the libcuvs.so file path and appends it to LD_LIBRARY_PATH
+rapids-logger "Find libcuvs so file and prepend paths to LD_LIBRARY_PATH"
+
+CONDA_PKG_CACHE_DIR="/opt/conda/pkgs" # comes from `conda info`. Dont know if this ever changes.
+if [ -d "$CONDA_PKG_CACHE_DIR" ]; then
+  echo "==> Directory '$CONDA_PKG_CACHE_DIR' exists."
+  LIBCUVS_SO_FILE="libcuvs.so"
+  LIBCUVS_PATH=$(find $CONDA_PKG_CACHE_DIR -name $LIBCUVS_SO_FILE)
+  if [ -z "$LIBCUVS_PATH" ]; then
+    echo "==> Could not find the so file. Not updating LD_LIBRARY_PATH"
+    exit 1
+  else
+    LIBCUVS_DIR=$(dirname "$LIBCUVS_PATH")
+    export LD_LIBRARY_PATH="$LIBCUVS_DIR:$LD_LIBRARY_PATH"
+    echo "LD_LIBRARY_PATH is: $LD_LIBRARY_PATH"
+  fi
+else
+  echo "==> Directory '$CONDA_PKG_CACHE_DIR' does not exist. Not updating LD_LIBRARY_PATH"
+  exit 1
+fi
+
 rapids-logger "Run Java build"
 
 bash ./build.sh "${EXTRA_BUILD_ARGS[@]}"

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -2,6 +2,7 @@
 # To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 channels:
 - conda-forge
+- rapidsai-nightly
 dependencies:
 - cuda-cudart-dev
 - cuda-nvtx-dev
@@ -11,6 +12,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libcuvs
 - maven
 - openjdk=22.*
 name: all_cuda-129_arch-aarch64

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -2,6 +2,7 @@
 # To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 channels:
 - conda-forge
+- rapidsai-nightly
 dependencies:
 - cuda-cudart-dev
 - cuda-nvtx-dev
@@ -11,6 +12,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libcuvs
 - maven
 - openjdk=22.*
 name: all_cuda-129_arch-x86_64

conda/environments/all_cuda-130_arch-aarch64.yaml

@@ -2,6 +2,7 @@
 # To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 channels:
 - conda-forge
+- rapidsai-nightly
 dependencies:
 - cuda-cudart-dev
 - cuda-nvtx-dev
@@ -11,6 +12,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libcuvs
 - maven
 - openjdk=22.*
 name: all_cuda-130_arch-aarch64

conda/environments/all_cuda-130_arch-x86_64.yaml

@@ -2,6 +2,7 @@
 # To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 channels:
 - conda-forge
+- rapidsai-nightly
 dependencies:
 - cuda-cudart-dev
 - cuda-nvtx-dev
@@ -11,6 +12,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libcuvs
 - maven
 - openjdk=22.*
 name: all_cuda-130_arch-x86_64

dependencies.yaml

@@ -21,6 +21,7 @@ files:
       - java
 channels:
   - conda-forge
+  - rapidsai-nightly
 dependencies:
   checks:
     common:
@@ -66,6 +67,7 @@ dependencies:
           - libcurand-dev
           - libcusolver-dev
           - libcusparse-dev
+          - libcuvs
   java:
     common:
       - output_types: conda

pom.xml

@@ -69,7 +69,7 @@
       <groupId>com.nvidia.cuvs</groupId>
       <artifactId>cuvs-java</artifactId>
       <!-- Note: This is temporary measure (artifact coming from https://maven.searchscale.com/snapshots, for now) and will get updated in subsequent PRs. -->
-      <version>25.8.0-4f53f-SNAPSHOT</version>
+      <version>25.10.0-2c0e1-SNAPSHOT</version>
     </dependency>
   </dependencies>
 

src/main/java/com/nvidia/cuvs/lucene/CuVSIndex.java

@@ -15,8 +15,6 @@
  */
 package com.nvidia.cuvs.lucene;
 
-import static com.nvidia.cuvs.lucene.CuVSVectorsReader.handleThrowable;
-
 import com.nvidia.cuvs.BruteForceIndex;
 import com.nvidia.cuvs.CagraIndex;
 import com.nvidia.cuvs.HnswIndex;
@@ -103,16 +101,16 @@ public void close() throws IOException {
   private void destroyIndices() throws IOException {
     try {
       if (cagraIndex != null) {
-        cagraIndex.destroyIndex();
+        cagraIndex.close();
       }
       if (bruteforceIndex != null) {
-        bruteforceIndex.destroyIndex();
+        bruteforceIndex.close();
       }
       if (hnswIndex != null) {
-        hnswIndex.destroyIndex();
+        hnswIndex.close();
       }
     } catch (Throwable t) {
-      handleThrowable(t);
+      Utils.handleThrowable(t);
     }
   }
 }

src/main/java/com/nvidia/cuvs/lucene/CuVSKnnFloatVectorQuery.java

@@ -19,6 +19,7 @@
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.KnnFloatVectorQuery;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.knn.KnnCollectorManager;
 import org.apache.lucene.util.Bits;
@@ -29,8 +30,9 @@ public class CuVSKnnFloatVectorQuery extends KnnFloatVectorQuery {
   private final int iTopK;
   private final int searchWidth;
 
-  public CuVSKnnFloatVectorQuery(String field, float[] target, int k, int iTopK, int searchWidth) {
-    super(field, target, k);
+  public CuVSKnnFloatVectorQuery(
+      String field, float[] target, int k, Query filter, int iTopK, int searchWidth) {
+    super(field, target, k, filter);
     this.iTopK = iTopK;
     this.searchWidth = searchWidth;
   }
@@ -46,7 +48,7 @@ protected TopDocs approximateSearch(
     PerLeafCuVSKnnCollector results = new PerLeafCuVSKnnCollector(k, iTopK, searchWidth);
 
     LeafReader reader = context.reader();
-    reader.searchNearestVectors(field, this.getTargetCopy(), results, null);
+    reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs);
     return results.topDocs();
   }
 }

src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsFormat.java

@@ -86,6 +86,12 @@ public CuVSVectorsFormat(
   }
 
   private static CuVSResources cuVSResourcesOrNull() {
+    try {
+      System.loadLibrary(
+          "cudart"); // nocommit: this is here so as to pass CI, should goto cuvs-java
+    } catch (UnsatisfiedLinkError e) {
+      LOG.warning("Could not load CUDA runtime library: " + e.getMessage());
+    }
     try {
       resources = CuVSResources.create();
       return resources;

src/main/java/com/nvidia/cuvs/lucene/CuVSVectorsReader.java

@@ -271,7 +271,7 @@ private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException {
         }
       }
     } catch (Throwable t) {
-      handleThrowable(t);
+      Utils.handleThrowable(t);
     }
     return new CuVSIndex(cagraIndex, bruteForceIndex, hnswIndex);
   }
@@ -367,7 +367,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
       try {
         searchResult = cagraIndex.search(query).getResults();
       } catch (Throwable t) {
-        handleThrowable(t);
+        Utils.handleThrowable(t);
       }
       // List expected to have only one entry because of single query "target".
       assert searchResult.size() == 1;
@@ -385,7 +385,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
       try {
         searchResult = bruteforceIndex.search(query).getResults();
       } catch (Throwable t) {
-        handleThrowable(t);
+        Utils.handleThrowable(t);
       }
       assert searchResult.size() == 1;
       result = searchResult.getFirst();
@@ -472,12 +472,15 @@ static void checkVersion(int versionMeta, int versionVectorData, IndexInput in)
     }
   }
 
-  static void handleThrowable(Throwable t) throws IOException {
-    switch (t) {
-      case IOException ioe -> throw ioe;
-      case Error error -> throw error;
-      case RuntimeException re -> throw re;
-      case null, default -> throw new RuntimeException("UNEXPECTED: exception type", t);
-    }
+  public FieldInfos getFieldInfos() {
+    return fieldInfos;
+  }
+
+  public IntObjectHashMap<CuVSIndex> getCuvsIndexes() {
+    return cuvsIndices;
+  }
+
+  public IntObjectHashMap<FieldEntry> getFieldEntries() {
+    return fields;
   }
 }