[REVIEW] Lucene upgrade to 10.4 in cuvs-lucene

examples/java/cuvs-lucene/pom.xml

@@ -79,18 +79,18 @@
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-core</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
     </dependency>
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-codecs</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
         <scope>test</scope>
     </dependency>
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-backward-codecs</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
     </dependency>
     <dependency>
         <groupId>commons-io</groupId>

java/cuvs-lucene/README.md

@@ -12,11 +12,11 @@ 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_4_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.
 
 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. The on-disk format is Lucene99 HNSW, so indexes built on the GPU can be read by stock Lucene 10.4+.
   - `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

java/cuvs-lucene/pom.xml

@@ -110,28 +110,27 @@ SPDX-License-Identifier: Apache-2.0
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-core</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
     </dependency>
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-codecs</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
         <scope>test</scope>
     </dependency>
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-backward-codecs</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
     </dependency>
     <dependency>
         <groupId>org.apache.lucene</groupId>
-        <artifactId>lucene-misc</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
     </dependency>
     <dependency>
         <groupId>org.apache.lucene</groupId>
         <artifactId>lucene-test-framework</artifactId>
-        <version>10.2.0</version>
+        <version>10.4.0</version>
         <scope>test</scope>
     </dependency>
     <dependency>

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

@@ -235,6 +235,14 @@ private static CuVSMatrix buildCagraGraphForSubset(
     return CuVSMatrix.ofArray(remappedAdjacency);
   }
 
+  private static int[] getSortedNodes(NodesIterator nodesOnLevel) {
+    int[] nodes = new int[nodesOnLevel.size()];
+    int consumed = nodesOnLevel.consume(nodes);
+    assert consumed == nodesOnLevel.size();
+    Arrays.sort(nodes);
+    return nodes;
+  }
+
   /**
    * Returns a 2D array of offsets (information written while writing the meta info)
    *
@@ -260,7 +268,7 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde
     // rather than per level/per task below.
     int maxConn = graph.maxConn();
 
-    int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0));
+    int[] level0Nodes = getSortedNodes(graph.getNodesOnLevel(0));
     offsets[0] = new int[level0Nodes.length];
     if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) {
       writeLevel0Parallel(
@@ -270,7 +278,7 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde
     }
 
     for (int level = 1; level < numLevels; level++) {
-      int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level));
+      int[] sortedNodes = getSortedNodes(graph.getNodesOnLevel(level));
       offsets[level] = new int[sortedNodes.length];
       writeLevelSerial(
           graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0, maxConn);
@@ -366,8 +374,9 @@ private static void writeLevel0Parallel(
   }
 
   /**
-   * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt
-   * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical.
+   * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size +
+   * GroupVInts deltas) to {@code out}. Shared by the serial and parallel paths so encoding is
+   * identical.
    */
   private static void encodeNode(
       NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0)
@@ -388,9 +397,7 @@ private static void encodeNode(
       }
     }
     out.writeVInt(actualSize);
-    for (int i = 0; i < actualSize; i++) {
-      out.writeVInt(scratch[i]);
-    }
+    out.writeGroupVInts(scratch, actualSize);
   }
 
   /**

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

@@ -31,7 +31,7 @@ public class CuVS2510GPUSearchCodec extends FilterCodec {
   public CuVS2510GPUSearchCodec() throws Exception {
     this(
         NAME,
-        LuceneProvider.getCodec("101"),
+        LuceneProvider.getCodec("104"),
         new GPUSearchParams.Builder().build(),
         FilterBitsetCacheConfig.DEFAULT);
   }
@@ -55,7 +55,7 @@ public CuVS2510GPUSearchCodec(String name, Codec delegate) {
    * @throws Exception Exception raised when initializing the codec
    */
   public CuVS2510GPUSearchCodec(GPUSearchParams params) throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"), params, FilterBitsetCacheConfig.DEFAULT);
+    this(NAME, LuceneProvider.getCodec("104"), params, FilterBitsetCacheConfig.DEFAULT);
   }
 
   /**
@@ -67,7 +67,7 @@ public CuVS2510GPUSearchCodec(GPUSearchParams params) throws Exception {
    */
   public CuVS2510GPUSearchCodec(GPUSearchParams params, FilterBitsetCacheConfig filterCacheConfig)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"), params, filterCacheConfig);
+    this(NAME, LuceneProvider.getCodec("104"), params, filterCacheConfig);
   }
 
   /**

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

@@ -40,13 +40,14 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.AcceptDocs;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IOContext.Context;
 import org.apache.lucene.store.IndexInput;
-import org.apache.lucene.store.ReadAdvice;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.hnsw.IntToIntFunction;
@@ -127,8 +128,7 @@ public CuVS2510GPUVectorsReader(SegmentReadState state, FlatVectorsReader flatRe
       } finally {
         CodecUtil.checkFooter(meta, priorException);
       }
-      var ioContext = state.context.withReadAdvice(ReadAdvice.SEQUENTIAL);
-      cuvsIndexInput = openCuVSInput(state, versionMeta, ioContext);
+      cuvsIndexInput = openCuVSInput(state, versionMeta, state.context);
       /*
        * Only load indexes on the GPU when this reader is opening for searches.
        * Do not load indexes on the GPU when this reader is opening during merge calls.
@@ -450,11 +450,59 @@ private static FloatToFloatFunction getScoreNormalizationFunc(VectorSimilarityFu
     return score -> (1f / (1f + score));
   }
 
+  /** Maps AcceptDocs to vector ordinals. bits() may be null in Lucene 10.4. */
+  private static Bits computeAcceptedOrds(FloatVectorValues rawValues, AcceptDocs acceptDocs)
+      throws IOException {
+    if (acceptDocs == null) {
+      return null;
+    }
+    Bits live = acceptDocs.bits();
+    if (live != null) {
+      Bits mapped = rawValues.getAcceptOrds(live);
+      if (mapped != null) {
+        return mapped;
+      }
+    }
+    final Bits docAccept;
+    if (live != null) {
+      docAccept = live;
+    } else {
+      BitSet docBits = new BitSet();
+      DocIdSetIterator disi = acceptDocs.iterator();
+      for (int doc = disi.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = disi.nextDoc()) {
+        docBits.set(doc);
+      }
+      docAccept =
+          new Bits() {
+            @Override
+            public boolean get(int docId) {
+              return docBits.get(docId);
+            }
+
+            @Override
+            public int length() {
+              return docBits.length();
+            }
+          };
+    }
+    return new Bits() {
+      @Override
+      public boolean get(int ord) {
+        return docAccept.get(rawValues.ordToDoc(ord));
+      }
+
+      @Override
+      public int length() {
+        return rawValues.size();
+      }
+    };
+  }
+
   /**
    * Returns the k nearest neighbor documents using cuVS's CAGRA or brute force algorithm for this field, to the given vector.
    */
   @Override
-  public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     var fieldEntry = getFieldEntry(field, VectorEncoding.FLOAT32);
     if (fieldEntry.count() == 0 || knnCollector.k() == 0) {
@@ -468,12 +516,13 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
     }
 
     final FloatVectorValues rawValues = flatVectorsReader.getFloatVectorValues(field);
-    final Bits acceptedOrds = rawValues.getAcceptOrds(acceptDocs);
+    final Bits acceptedOrds = computeAcceptedOrds(rawValues, acceptDocs);
     BitSet[] mask = null;
     int maskLength = 0;
     int topK = knnCollector.k();
 
     if (acceptDocs != null) {
+      assert acceptedOrds != null;
       mask = new BitSet[1]; // As there is only one query "target"
       mask[0] = new BitSet(acceptedOrds.length());
       /*
@@ -597,7 +646,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
    * This is not supported.
    */
   @Override
-  public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
+  public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs)
       throws IOException {
     throw new UnsupportedOperationException("Byte vectors are not currently supported");
   }

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

@@ -28,6 +28,7 @@
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.IndexSearcher;
@@ -268,7 +269,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
   @Override
   protected TopDocs approximateSearch(
       LeafReaderContext context,
-      Bits acceptDocs,
+      AcceptDocs acceptDocs,
       int visitedLimit,
       KnnCollectorManager knnCollectorManager)
       throws IOException {

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

@@ -30,7 +30,7 @@ public class Lucene101AcceleratedHNSWCodec extends FilterCodec {
    * @throws Exception
    */
   public Lucene101AcceleratedHNSWCodec() throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
   }
 
   /**
@@ -52,7 +52,7 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) {
    */
   public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
     initializeFormat(acceleratedHNSWParams, 0);
   }
 
@@ -70,7 +70,7 @@ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams
    */
   Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
     initializeFormat(acceleratedHNSWParams, numInputVectors);
   }
 

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

@@ -25,7 +25,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedCodec extends FilterCodec {
   private KnnVectorsFormat format;
 
   public LuceneAcceleratedHNSWBinaryQuantizedCodec() throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
   }
 
   public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) {
@@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) {
 
   public LuceneAcceleratedHNSWBinaryQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
     initializeFormat(acceleratedHNSWParams);
   }
 

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

@@ -27,19 +27,17 @@ 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 LuceneProvider LUCENE_PROVIDER;
   private static final FlatVectorsFormat FLAT_VECTORS_FORMAT;
   private static final int MAX_DIMENSIONS = 4096;
 
   private final AcceleratedHNSWParams acceleratedHNSWParams;
 
   static {
     try {
-      LUCENE99_PROVIDER = LuceneProvider.getInstance("99");
-      LUCENE102_PROVIDER = LuceneProvider.getInstance("102");
+      LUCENE_PROVIDER = LuceneProvider.getInstance("104");
       FLAT_VECTORS_FORMAT =
-          LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE);
+          LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE);
     } catch (Exception e) {
       throw new ExceptionInInitializerError(e.getMessage());
     }
@@ -79,13 +77,13 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
           state, acceleratedHNSWParams, flatWriter);
     } else {
       try {
-        // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat format
+        // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat
         log.log(
             Level.WARNING,
             "GPU based indexing not supported, falling back to using the"
                 + " Lucene102HnswBinaryQuantizedVectorsFormat");
         KnnVectorsFormat fallbackFormat =
-            LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance(
+            LUCENE_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance(
                 acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth());
         return fallbackFormat.fieldsWriter(state);
       } catch (Exception e) {
@@ -100,7 +98,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
   @Override
   public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
     try {
-      return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance(
+      return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance(
           state, FLAT_VECTORS_FORMAT.fieldsReader(state));
     } catch (Exception e) {
       throw Utils.handleThrowable(e);

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

@@ -25,7 +25,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedCodec extends FilterCodec {
   private KnnVectorsFormat format;
 
   public LuceneAcceleratedHNSWScalarQuantizedCodec() throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
   }
 
   public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) {
@@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) {
 
   public LuceneAcceleratedHNSWScalarQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams)
       throws Exception {
-    this(NAME, LuceneProvider.getCodec("101"));
+    this(NAME, LuceneProvider.getCodec("104"));
     initializeFormat(acceleratedHNSWParams);
   }