Add PyLucene integration and CPU/GPU end-to-end tests

java/cuvs-lucene/README.md

@@ -12,14 +12,14 @@ This is a project for using [cuVS](https://github.com/rapidsai/cuvs), NVIDIA's G
 
 ## What is cuvs-lucene?
 
-`cuvs-lucene` provides a pluggable [KnnVectorsFormat](https://lucene.apache.org/core/10_2_0/core/org/apache/lucene/codecs/KnnVectorsFormat.html) that uses cuVS to offload vector index build — and optionally search — to NVIDIA GPUs. Because it plugs in through a standard Lucene codec, existing Lucene applications can take advantage of GPU acceleration with minimal code changes and gracefully fall back to the default CPU codec when no GPU is present.
+`cuvs-lucene` provides a pluggable [KnnVectorsFormat](https://lucene.apache.org/core/10_2_0/core/org/apache/lucene/codecs/KnnVectorsFormat.html) that uses cuVS to offload vector index build — and optionally search — to NVIDIA GPUs. The accelerated-HNSW codecs can fall back to Lucene's CPU HNSW writer when cuVS is unavailable; the GPU-search codec requires cuVS. This development line is compiled and tested against the Lucene 10.2.0 runtime ABI.
 
 Four codecs are currently provided:
 
-- `Lucene101AcceleratedHNSWCodec` — GPU-accelerated HNSW build with CPU HNSW search. The on-disk format is standard Lucene HNSW, so indexes built on the GPU can be read by any stock Lucene 10.x reader.
+- `Lucene101AcceleratedHNSWCodec` — GPU-accelerated HNSW build with CPU HNSW search. Its vector data uses Lucene's standard HNSW format and stock HNSW reader; applications still need a compatible `cuvs-lucene` codec provider to resolve the segment codec.
   - `LuceneAcceleratedHNSWScalarQuantizedCodec` — scalar-quantized vectors for a smaller index footprint.
   - `LuceneAcceleratedHNSWBinaryQuantizedCodec` — binary-quantized vectors for an even smaller index footprint.
-- `CuVS2510GPUSearchCodec` — GPU-accelerated HNSW build and GPU search
+- `CuVS2510GPUSearchCodec` — GPU CAGRA build and GPU CAGRA search
 
 ## Installing cuvs-lucene
 

java/cuvs-lucene/pom.xml

@@ -196,6 +196,24 @@ SPDX-License-Identifier: Apache-2.0
           </systemPropertyVariables>
         </configuration>
       </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-failsafe-plugin</artifactId>
+        <version>3.5.4</version>
+        <executions>
+          <execution>
+            <goals>
+              <goal>integration-test</goal>
+              <goal>verify</goal>
+            </goals>
+          </execution>
+        </executions>
+        <configuration>
+          <systemPropertyVariables>
+            <cuvs.lucene.thinJar>${project.build.directory}/${project.build.finalName}.jar</cuvs.lucene.thinJar>
+          </systemPropertyVariables>
+        </configuration>
+      </plugin>
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-source-plugin</artifactId>

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java

@@ -52,9 +52,9 @@
  * single multi-partition search to cuVS, passing one Lucene segment per cuVS partition. cuVS
  * runs the per-partition CAGRA searches, applies distance post-processing, and performs the
  * cross-partition top-k merge internally; the returned arrays are mapped to Lucene doc IDs on
- * the host. The effective CAGRA algorithm (SINGLE_CTA or MULTI_KERNEL) is selected by cuVS
- * based on {@code searchAlgo} and {@code itopk_size}, with MULTI_KERNEL handling k beyond
- * SINGLE_CTA's per-partition cap.
+ * the host. For a multi-partition search, cuVS resolves {@code AUTO} to {@code SINGLE_CTA} or
+ * {@code MULTI_CTA} from the search parameters and query/partition topology; {@code MULTI_KERNEL}
+ * is not supported by the multi-partition API.
  *
  * <p>If the query has an explicit {@code filter}, or if any segment carries live-document deletes,
  * the acceptance mask (filter ∩ liveDocs) is packed into one {@link FilterBitsetHandle} per segment
@@ -481,10 +481,9 @@ private static CuVS2510GPUVectorsReader unwrapGpuReader(LeafReaderContext ctx, S
   /**
    * Builds a {@link Query} that matches exactly the given pre-scored documents.
    *
-   * <p>Partitions {@code scoreDocs} by segment (using {@link ScoreDoc#shardIndex} as the segment
-   * offset relative to {@link LeafReaderContext#docBase}), then returns a {@link Scorer} per
-   * segment that iterates those docs in ascending doc-ID order and replays their pre-computed
-   * scores.
+   * <p>Partitions {@code scoreDocs} by each global doc ID's membership in a leaf's {@link
+   * LeafReaderContext#docBase} range, then returns a {@link Scorer} per segment that iterates those
+   * docs in ascending doc-ID order and replays their pre-computed scores.
    */
   private static Query docAndScoreQuery(ScoreDoc[] scoreDocs) {
     return new Query() {

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java

@@ -52,7 +52,12 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) {
    */
   public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("101"), acceleratedHNSWParams);
+  }
+
+  private Lucene101AcceleratedHNSWCodec(
+      String name, Codec delegate, AcceleratedHNSWParams acceleratedHNSWParams) {
+    super(name, delegate);
     initializeFormat(acceleratedHNSWParams);
   }
 

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java

@@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) {
 
   public LuceneAcceleratedHNSWBinaryQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    super(NAME, LuceneProvider.getCodec("101"));
     initializeFormat(acceleratedHNSWParams);
   }
 

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java

@@ -8,12 +8,13 @@
 
 import com.nvidia.cuvs.LibraryException;
 import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.util.concurrent.Callable;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.apache.lucene.codecs.KnnVectorsFormat;
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.KnnVectorsWriter;
-import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
@@ -27,22 +28,86 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector
 
   private static final Logger log =
       Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class.getName());
-  private static final LuceneProvider LUCENE102_PROVIDER;
-  private static final LuceneProvider LUCENE99_PROVIDER;
-  private static final FlatVectorsFormat FLAT_VECTORS_FORMAT;
   private static final int MAX_DIMENSIONS = 4096;
+  private static volatile FlatVectorsFormat cachedFlatVectorsFormat;
 
   private final AcceleratedHNSWParams acceleratedHNSWParams;
+  private volatile KnnVectorsFormat cachedFallbackFormat;
 
-  static {
+  private static LuceneProvider getLucene99Provider() throws IOException {
     try {
-      LUCENE99_PROVIDER = LuceneProvider.getInstance("99");
-      LUCENE102_PROVIDER = LuceneProvider.getInstance("102");
-      FLAT_VECTORS_FORMAT =
-          LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE);
+      return LuceneProvider.getInstance(LuceneProvider.LUCENE_99_FORMAT_VERSION);
+    } catch (ClassNotFoundException e) {
+      throw new IOException("Lucene99 vector formats are not available in this runtime", e);
+    }
+  }
+
+  private static RuntimeException handleConstructionFailure(String formatName, Throwable failure)
+      throws IOException {
+    if (failure instanceof IOException
+        || failure instanceof RuntimeException
+        || failure instanceof Error) {
+      return Utils.handleThrowable(failure);
+    }
+    return new IllegalStateException("Unable to construct " + formatName, failure);
+  }
+
+  static <T> T constructLucene102Format(String formatName, Callable<T> constructor)
+      throws IOException {
+    try {
+      return constructor.call();
+    } catch (ClassNotFoundException e) {
+      throw new UnsupportedOperationException(
+          formatName + " is not available in this Lucene runtime", e);
+    } catch (InvocationTargetException e) {
+      throw handleConstructionFailure(formatName, e.getTargetException());
+    } catch (ReflectiveOperationException e) {
+      throw new IllegalStateException("Unable to construct " + formatName, e);
+    } catch (IOException | RuntimeException | Error e) {
+      throw Utils.handleThrowable(e);
     } catch (Exception e) {
-      throw new ExceptionInInitializerError(e.getMessage());
+      throw new IllegalStateException("Unable to construct " + formatName, e);
+    }
+  }
+
+  private static FlatVectorsFormat getOrCreateFlatVectorsFormat() throws IOException {
+    FlatVectorsFormat format = cachedFlatVectorsFormat;
+    if (format == null) {
+      synchronized (LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class) {
+        format = cachedFlatVectorsFormat;
+        if (format == null) {
+          format =
+              constructLucene102Format(
+                  "Lucene102BinaryQuantizedVectorsFormat",
+                  () ->
+                      LuceneProvider.getInstance(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION)
+                          .getLuceneBinaryQuantizedVectorsFormatInstance());
+          cachedFlatVectorsFormat = format;
+        }
+      }
     }
+    return format;
+  }
+
+  private KnnVectorsFormat getOrCreateFallbackFormat() throws IOException {
+    KnnVectorsFormat format = cachedFallbackFormat;
+    if (format == null) {
+      synchronized (this) {
+        format = cachedFallbackFormat;
+        if (format == null) {
+          format =
+              constructLucene102Format(
+                  "Lucene102HnswBinaryQuantizedVectorsFormat",
+                  () ->
+                      LuceneProvider.getInstance(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION)
+                          .getLuceneHnswBinaryQuantizedKnnVectorsFormatInstance(
+                              acceleratedHNSWParams.getMaxConn(),
+                              acceleratedHNSWParams.getBeamWidth()));
+          cachedFallbackFormat = format;
+        }
+      }
+    }
+    return format;
   }
 
   /**
@@ -70,27 +135,20 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat(
    */
   @Override
   public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-    var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state);
     if (isSupported()) {
+      var flatWriter = getOrCreateFlatVectorsFormat().fieldsWriter(state);
       log.log(
           Level.FINE,
           "cuVS is supported so using the Lucene99AcceleratedHNSWBinaryQuantizedVectorsWriter");
       return new LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter(
           state, acceleratedHNSWParams, flatWriter);
     } else {
-      try {
-        // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat format
-        log.log(
-            Level.WARNING,
-            "GPU based indexing not supported, falling back to using the"
-                + " Lucene102HnswBinaryQuantizedVectorsFormat");
-        KnnVectorsFormat fallbackFormat =
-            LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance(
-                acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth());
-        return fallbackFormat.fieldsWriter(state);
-      } catch (Exception e) {
-        throw Utils.handleThrowable(e);
-      }
+      // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat format
+      log.log(
+          Level.WARNING,
+          "GPU based indexing not supported, falling back to using the"
+              + " Lucene102HnswBinaryQuantizedVectorsFormat");
+      return getOrCreateFallbackFormat().fieldsWriter(state);
     }
   }
 
@@ -100,8 +158,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
   @Override
   public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
     try {
-      return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance(
-          state, FLAT_VECTORS_FORMAT.fieldsReader(state));
+      return getLucene99Provider()
+          .getLuceneHnswVectorsReaderInstance(
+              state, getOrCreateFlatVectorsFormat().fieldsReader(state));
     } catch (Exception e) {
       throw Utils.handleThrowable(e);
     }

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java

@@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) {
 
   public LuceneAcceleratedHNSWScalarQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    super(NAME, LuceneProvider.getCodec("101"));
     initializeFormat(acceleratedHNSWParams);
   }
 

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java

@@ -6,7 +6,6 @@
 
 import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported;
 
-import com.nvidia.cuvs.LibraryException;
 import java.io.IOException;
 import java.util.logging.Logger;
 import org.apache.lucene.codecs.KnnVectorsFormat;
@@ -25,26 +24,60 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsFormat extends KnnVector
 
   private static final Logger log =
       Logger.getLogger(LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.class.getName());
-  private static final LuceneProvider LUCENE_PROVIDER;
-  private static final FlatVectorsFormat FLAT_VECTORS_FORMAT;
   private static final int MAX_DIMENSIONS = 4096;
+  private static volatile FlatVectorsFormat cachedFlatVectorsFormat;
 
   private final AcceleratedHNSWParams acceleratedHNSWParams;
+  private volatile KnnVectorsFormat cachedFallbackFormat;
 
-  static {
+  private static LuceneProvider getLuceneProvider() throws IOException {
     try {
-      LUCENE_PROVIDER = LuceneProvider.getInstance("99");
-      FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneScalarQuantizedVectorsFormatInstance();
-    } catch (Exception e) {
-      throw new ExceptionInInitializerError(e.getMessage());
+      return LuceneProvider.getInstance(LuceneProvider.LUCENE_99_FORMAT_VERSION);
+    } catch (ClassNotFoundException e) {
+      throw new IOException("Lucene99 vector formats are not available in this runtime", e);
     }
   }
 
-  /**
-   * Initializes {@link LuceneAcceleratedHNSWScalarQuantizedVectorsFormat} with default values.
-   *
-   * @throws LibraryException if the native library fails to load
-   */
+  private static FlatVectorsFormat getOrCreateFlatVectorsFormat() throws IOException {
+    FlatVectorsFormat format = cachedFlatVectorsFormat;
+    if (format == null) {
+      synchronized (LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.class) {
+        format = cachedFlatVectorsFormat;
+        if (format == null) {
+          try {
+            format = getLuceneProvider().getLuceneScalarQuantizedVectorsFormatInstance();
+            cachedFlatVectorsFormat = format;
+          } catch (Exception e) {
+            throw Utils.handleThrowable(e);
+          }
+        }
+      }
+    }
+    return format;
+  }
+
+  private KnnVectorsFormat getOrCreateFallbackFormat() throws IOException {
+    KnnVectorsFormat format = cachedFallbackFormat;
+    if (format == null) {
+      synchronized (this) {
+        format = cachedFallbackFormat;
+        if (format == null) {
+          try {
+            format =
+                getLuceneProvider()
+                    .getLuceneHnswScalarQuantizedKnnVectorsFormatInstance(
+                        acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth());
+            cachedFallbackFormat = format;
+          } catch (Exception e) {
+            throw Utils.handleThrowable(e);
+          }
+        }
+      }
+    }
+    return format;
+  }
+
+  /** Initializes {@link LuceneAcceleratedHNSWScalarQuantizedVectorsFormat} with default values. */
   public LuceneAcceleratedHNSWScalarQuantizedVectorsFormat() {
     this(new AcceleratedHNSWParams.Builder().build());
   }
@@ -65,8 +98,8 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(
    */
   @Override
   public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-    var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state);
     if (isSupported()) {
+      var flatWriter = getOrCreateFlatVectorsFormat().fieldsWriter(state);
       log.info("cuVS is supported so using the Lucene99AcceleratedHNSWQuantizedVectorsWriter");
       return new LuceneAcceleratedHNSWScalarQuantizedVectorsWriter(
           state, acceleratedHNSWParams, flatWriter);
@@ -76,10 +109,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         log.warning(
             "GPU based indexing not supported, falling back to using the"
                 + " Lucene99HnswScalarQuantizedVectorsFormat");
-        KnnVectorsFormat fallbackFormat =
-            LUCENE_PROVIDER.getLuceneHnswScalarQuantizedVectorsFormatInstance(
-                acceleratedHNSWParams.getBeamWidth(), acceleratedHNSWParams.getMaxConn());
-        return fallbackFormat.fieldsWriter(state);
+        return getOrCreateFallbackFormat().fieldsWriter(state);
       } catch (Exception e) {
         throw Utils.handleThrowable(e);
       }
@@ -92,8 +122,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
   @Override
   public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
     try {
-      return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance(
-          state, FLAT_VECTORS_FORMAT.fieldsReader(state));
+      return getLuceneProvider()
+          .getLuceneHnswVectorsReaderInstance(
+              state, getOrCreateFlatVectorsFormat().fieldsReader(state));
     } catch (Exception e) {
       throw Utils.handleThrowable(e);
     }