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

java/cuvs-lucene/README.md

@@ -12,24 +12,24 @@ 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
 
 ### Prerequisites
 
 - A machine with an NVIDIA GPU
-- [CUDA 12.0+](https://developer.nvidia.com/cuda-toolkit-archive)
+- [CUDA 12.2+](https://developer.nvidia.com/cuda-toolkit-archive)
 - [JDK 22](https://jdk.java.net/archive/)
 - [Maven 3.9.6+](https://maven.apache.org/download.cgi)
-- A matching version of the [cuVS libraries](https://docs.rapids.ai/api/cuvs/stable/build/#build-from-source). For Maven usage, install the cuVS tarball and add it to your system library load path. See the cuVS [tarball install instructions](https://docs.rapids.ai/api/cuvs/stable/build/#download-extract).
+- A matching version of the [cuVS libraries](https://docs.rapids.ai/api/cuvs/stable/build/#build-from-source). For Maven usage, install the cuVS tarball and add it to your system library load path. See the cuVS [tarball install instructions](https://docs.nvidia.com/cuvs/installation/c#tarball).
 
 ### Maven
 
@@ -56,13 +56,13 @@ The resulting artifacts are written to `target/`.
 
 To run the tests, add `--run-java-tests` to any of the commands above. Be sure to set (manually, if needed)
 your `LD_LIBRARY_PATH` to include the directory with the appropriate (matching) version of `libcuvs.so`, as
-described in the cuVS [tarball install instructions](https://docs.rapids.ai/api/cuvs/stable/build/#download-extract).
+described in the cuVS [tarball install instructions](https://docs.nvidia.com/cuvs/installation/c#tarball).
 
 ## Getting Started
 
 The example below plugs the GPU-accelerated HNSW codec into a standard Lucene `IndexWriter`. Once the codec is set on the `IndexWriterConfig`, indexing proceeds exactly as it would with the default Lucene codec, and search uses the stock `KnnFloatVectorQuery`.
 
-Before running it, make sure cuVS is installed and available on your system library load path. The cuVS [tarball install instructions](https://docs.rapids.ai/api/cuvs/stable/build/#download-extract) show how to set this up.
+Before running it, make sure cuVS is installed and available on your system library load path. The cuVS [tarball install instructions](https://docs.nvidia.com/cuvs/installation/c#tarball) show how to set this up.
 
 ### RMM async allocation for GPU search
 
@@ -119,15 +119,71 @@ public class HelloCuvsLucene {
 }
 ```
 
-The artifacts would be built and available in the target / folder.
+The artifacts are built in the `target/` directory.
 
-### Running Tests
+Run the example with:
 
 ```sh
 mvn -q compile org.codehaus.mojo:exec-maven-plugin:3.5.1:java \
   -Dexec.mainClass=com.nvidia.cuvs.lucene.examples.HelloCuvsLucene
 ```
 
+### Using with PyLucene
+
+The complete codec set requires PyLucene generated against Lucene 10.2.0. The
+official PyLucene 10.0.0 distribution is not a compatible full-feature runtime:
+in particular, GPU search uses Lucene 10.2 APIs and binary quantization uses
+Lucene102 vector formats.
+
+Apache does not publish a PyLucene 10.2.0 release. Linux development and
+testing therefore require a custom PyLucene wrapper build generated against
+the Lucene 10.2.0 sources. Prepare and activate that matching external
+environment before using PyLucene or running pytest; Maven can build the jar
+independently and does not produce the PyLucene runtime.
+
+Build the standard thin `cuvs-lucene` jar:
+
+```sh
+mvn clean package -DskipTests
+```
+
+Add that jar and the matching base `cuvs-java` jar to the classpath passed to
+`lucene.initVM(...)`. PyLucene can then load the codec through Lucene's service
+provider lookup:
+
+```python
+from org.apache.lucene.codecs import Codec
+
+codec = Codec.forName("Lucene101AcceleratedHNSWCodec")
+```
+
+Use the returned codec with `IndexWriterConfig.setCodec(codec)`. Initialize the
+JVM only after the custom 10.2 environment and every application jar are on its
+classpath; PyLucene cannot replace that classpath after `lucene.initVM(...)`.
+
+### Running Tests
+
+Run the Java tests with `mvn clean test`. Once the custom PyLucene 10.2
+environment is activated and the cuVS classpath and native-library environment
+are available, run the full parametrized CPU/GPU end-to-end suite directly
+with pytest:
+
+```sh
+python3 -m pytest -q -s src/test/python/test_pylucene_end_to_end.py
+```
+
+The cases live in `src/test/python/test_pylucene_end_to_end.py`; reusable
+runtime helpers are in `pylucene_test_support.py`, and the Java test bridge is
+compiled to `target/test-classes`. Set `CUVS_LUCENE_JAR`,
+`CUVS_LUCENE_CUVS_JAVA_JAR`, or `CUVS_LUCENE_PYLUCENE_TEST_CLASSES` only when
+their standard Maven locations are not appropriate. The helper verifies that
+PyLucene's Lucene version exactly matches the `lucene-core` version in this
+checkout's POM before starting the JVM, so a mismatched wrapper fails with an
+actionable error instead of a later linkage failure.
+
+The pytest IDs identify CPU HNSW, CAGRA-built HNSW, and CAGRA-search cases;
+use pytest's `-k` option for a focused run.
+
 For more examples, including one that indexes and searches entirely on the GPU using `CuVS2510GPUSearchCodec`, please refer to the [`examples/`](examples) directory.
 
 ## Contributing

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

@@ -447,9 +447,8 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
       }
       topK = Math.min(knnCollector.k() + 10, mask[0].cardinality());
       // numDocs must be the total vector count so cuVS sizes the prefilter to cover every ordinal.
-      // BitSet.length() is (highest set bit + 1), which under a selective filter is smaller than
-      // the
-      // vector count, leaving the trailing ordinals outside the filter and thus default-accepted.
+      // A selective BitSet may end before the vector count, which would leave trailing ordinals
+      // outside the filter and therefore accepted by default.
       maskLength = acceptedOrds.length();
     }
 

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

@@ -249,7 +249,7 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab
                 .withIndexParams(params)
                 .build();
         var deviceVectors = dataset.toDevice(getCuVSResourcesInstance());
-        var indexDataset = index.makePaddedDataset(deviceVectors)) {
+        var indexDataset = index.makePaddedDatasetView(deviceVectors)) {
       index.updateDataset(indexDataset);
       index.serialize(os);
     }

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
@@ -161,11 +161,9 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
       gpuReaders.add(gpuReader);
     }
 
-    // Build one filter handle per segment encoding (filter ∩ that segment's liveDocs) whenever any
-    // filtering is required — either an explicit Lucene filter, or live-document deletes in at
-    // least
-    // one segment. Each segment's handle becomes that partition's filter; a segment with neither an
-    // explicit filter nor deletes gets a null entry (unfiltered for that partition).
+    // Build one filter handle per segment whenever an explicit filter or any segment deletion
+    // requires filtering. Each handle encodes (filter ∩ segment liveDocs); an unaffected segment
+    // gets a null entry and remains unfiltered for that partition.
     boolean hasExplicitFilter = (filter != null);
     boolean hasDeletes = false;
     for (LeafReaderContext ctx : leaves) {
@@ -387,9 +385,7 @@ private List<FilterBitsetHandle> buildPerSegmentFilterHandles(
   private FilterBitsetHandle buildSegmentFilterHandle(
       Weight filterWeight, LeafReaderContext ctx, FloatVectorValues fvv) throws IOException {
     Bits liveDocs = ctx.reader().getLiveDocs();
-    // When filterWeight is null, accept all live documents (acceptDocs == liveDocs, which may
-    // itself
-    // be null to mean "all docs accepted" in this segment).
+    // Without an explicit filter, accept liveDocs directly; null means every document is live.
     Bits acceptDocs = (filterWeight != null) ? evalFilter(filterWeight, ctx, liveDocs) : liveDocs;
     Bits acceptedOrds = fvv.getAcceptOrds(acceptDocs);
     int numOrds = fvv.size();
@@ -481,10 +477,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/Lucene99AcceleratedHNSWVectorsFormat.java

@@ -81,7 +81,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
     } else {
       log.log(
           Level.WARNING,
-          "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter");
+          "GPU based indexing not supported, falling back to using the"
+              + " Lucene99HnswVectorsWriter");
       try {
         return LUCENE_PROVIDER.getLuceneHnswVectorsWriterInstance(
             state,

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

@@ -13,7 +13,6 @@
 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,24 +26,61 @@ 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 final LuceneProvider LUCENE_99_PROVIDER =
+      getLuceneProvider(LuceneProvider.LUCENE_99_FORMAT_VERSION);
+  private static volatile FlatVectorsFormat cachedFlatVectorsFormat;
 
   private final AcceleratedHNSWParams acceleratedHNSWParams;
+  private volatile KnnVectorsFormat cachedFallbackFormat;
 
-  static {
+  private static LuceneProvider getLuceneProvider(String version) {
     try {
-      LUCENE99_PROVIDER = LuceneProvider.getInstance("99");
-      LUCENE102_PROVIDER = LuceneProvider.getInstance("102");
-      FLAT_VECTORS_FORMAT =
-          LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE);
+      return LuceneProvider.getInstance(version);
     } catch (Exception e) {
-      throw new ExceptionInInitializerError(e.getMessage());
+      throw new UnsupportedOperationException(
+          "Lucene" + version + " vector formats are not available in this runtime", e);
     }
   }
 
+  private static FlatVectorsFormat getOrCreateFlatVectorsFormat() {
+    FlatVectorsFormat format = cachedFlatVectorsFormat;
+    if (format == null) {
+      synchronized (LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class) {
+        format = cachedFlatVectorsFormat;
+        if (format == null) {
+          try {
+            format =
+                getLuceneProvider(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION)
+                    .getLuceneBinaryQuantizedVectorsFormatInstance();
+            cachedFlatVectorsFormat = format;
+          } catch (Exception e) {
+            throw new UnsupportedOperationException(
+                "Binary quantized vectors require Lucene102 vector formats", e);
+          }
+        }
+      }
+    }
+    return format;
+  }
+
+  private KnnVectorsFormat getOrCreateFallbackFormat() throws Exception {
+    KnnVectorsFormat format = cachedFallbackFormat;
+    if (format == null) {
+      synchronized (this) {
+        format = cachedFallbackFormat;
+        if (format == null) {
+          format =
+              getLuceneProvider(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION)
+                  .getLuceneHnswBinaryQuantizedKnnVectorsFormatInstance(
+                      acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth());
+          cachedFallbackFormat = format;
+        }
+      }
+    }
+    return format;
+  }
+
   /**
    * Initializes {@link LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat} with default values.
    *
@@ -70,8 +106,8 @@ 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");
@@ -84,12 +120,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
             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);
+        return getOrCreateFallbackFormat().fieldsWriter(state);
       } catch (Exception e) {
-        throw new RuntimeException(e.getMessage());
+        throw new IOException("Unable to initialize the binary quantized fallback writer", e);
       }
     }
   }
@@ -100,10 +133,10 @@ 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 LUCENE_99_PROVIDER.getLuceneHnswVectorsReaderInstance(
+          state, getOrCreateFlatVectorsFormat().fieldsReader(state));
     } catch (Exception e) {
-      throw new RuntimeException(e.getMessage());
+      throw new IOException("Unable to initialize the binary quantized vectors reader", 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);
   }