diff --git a/examples/java/cuvs-lucene/pom.xml b/examples/java/cuvs-lucene/pom.xml index fa3207f1a2..9ca464900b 100644 --- a/examples/java/cuvs-lucene/pom.xml +++ b/examples/java/cuvs-lucene/pom.xml @@ -79,18 +79,18 @@ org.apache.lucene lucene-core - 10.2.0 + 10.4.0 org.apache.lucene lucene-codecs - 10.2.0 + 10.4.0 test org.apache.lucene lucene-backward-codecs - 10.2.0 + 10.4.0 commons-io diff --git a/java/cuvs-lucene/README.md b/java/cuvs-lucene/README.md index ae7885d142..884a65fb35 100644 --- a/java/cuvs-lucene/README.md +++ b/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 diff --git a/java/cuvs-lucene/pom.xml b/java/cuvs-lucene/pom.xml index bb797b6f8a..77d61b8d73 100644 --- a/java/cuvs-lucene/pom.xml +++ b/java/cuvs-lucene/pom.xml @@ -110,28 +110,27 @@ SPDX-License-Identifier: Apache-2.0 org.apache.lucene lucene-core - 10.2.0 + 10.4.0 org.apache.lucene lucene-codecs - 10.2.0 + 10.4.0 test org.apache.lucene lucene-backward-codecs - 10.2.0 + 10.4.0 org.apache.lucene - lucene-misc - 10.2.0 + 10.4.0 org.apache.lucene lucene-test-framework - 10.2.0 + 10.4.0 test diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index fdd7dd441f..86e28dcc3a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/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); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index ad286904cf..e1e3e17db6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/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); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 906a310ede..4f5dedc597 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/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"); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index f796d88799..c2b00095cb 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/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 { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index f95487e953..79992d3c03 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/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); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index 0b1653bc14..00aaa91dc9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/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); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 40a818683d..80189f7c34 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -27,8 +27,7 @@ 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; @@ -36,10 +35,9 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector 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); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java index 0705ed0a55..2587c7ce93 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java +++ b/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); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java index ead6daeaad..6d77354c01 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -33,7 +33,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsFormat extends KnnVector static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance("104"); FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneScalarQuantizedVectorsFormatInstance(); } catch (Exception e) { throw new ExceptionInInitializerError(e.getMessage()); @@ -72,10 +72,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, acceleratedHNSWParams, flatWriter); } else { try { - // Fallback to Lucene's Lucene99HnswScalarQuantizedVectorsFormat + // Fallback to Lucene's Lucene104HnswScalarQuantizedVectorsFormat log.warning( "GPU based indexing not supported, falling back to using the" - + " Lucene99HnswScalarQuantizedVectorsFormat"); + + " Lucene104HnswScalarQuantizedVectorsFormat"); KnnVectorsFormat fallbackFormat = LUCENE_PROVIDER.getLuceneHnswScalarQuantizedVectorsFormatInstance( acceleratedHNSWParams.getBeamWidth(), acceleratedHNSWParams.getMaxConn()); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 9caa2a6aa3..47483d0b97 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -8,10 +8,14 @@ import java.lang.invoke.VarHandle; import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; @@ -79,7 +83,7 @@ public class LuceneProvider { private static String luceneCodec = BASE + codecs + "LuceneCodec"; private static String luceneCodecFallback = BASE + fallbackCodecs + "LuceneCodec"; - private static LuceneProvider instance; + private static final Map INSTANCES = new HashMap<>(); private static MethodHandles.Lookup lookup = MethodHandles.lookup(); @@ -92,30 +96,35 @@ public class LuceneProvider { private Class scalarQuantizedVectorsFormat; private Class hnswScalarQuantizedVectorsFormat; - public static LuceneProvider getInstance(String version) throws ClassNotFoundException { - if (instance == null) { - instance = new LuceneProvider(version); + public static synchronized LuceneProvider getInstance(String version) + throws ClassNotFoundException { + LuceneProvider provider = INSTANCES.get(version); + if (provider == null) { + provider = new LuceneProvider(version); + INSTANCES.put(version, provider); } - return instance; + return provider; } private LuceneProvider(String version) throws ClassNotFoundException { + // Lucene 10.4 still ships float HNSW and flat vectors as lucene99. + String hnswVersion = "104".equals(version) ? "99" : version; flatVectorsFormat = loadClass( - setVersion(luceneFlatVectorsFormat, version), - setVersion(luceneFlatVectorsFormatFallback, version)); + setVersion(luceneFlatVectorsFormat, hnswVersion), + setVersion(luceneFlatVectorsFormatFallback, hnswVersion)); hnswVectorsFormat = loadClass( - setVersion(luceneHnswVectorsFormat, version), - setVersion(luceneHnswVectorsFormatFallback, version)); + setVersion(luceneHnswVectorsFormat, hnswVersion), + setVersion(luceneHnswVectorsFormatFallback, hnswVersion)); hnswVectorsReader = loadClass( - setVersion(luceneHnswVectorsReader, version), - setVersion(luceneHnswVectorsReaderFallback, version)); + setVersion(luceneHnswVectorsReader, hnswVersion), + setVersion(luceneHnswVectorsReaderFallback, hnswVersion)); hnswVectorsWriter = loadClass( - setVersion(luceneHnswVectorsWriter, version), - setVersion(luceneHnswVectorsWriterFallback, version)); + setVersion(luceneHnswVectorsWriter, hnswVersion), + setVersion(luceneHnswVectorsWriterFallback, hnswVersion)); scalarQuantizedVectorsFormat = loadClass( setVersion(luceneScalarQuantizedVectorsFormat, version), @@ -126,16 +135,16 @@ private LuceneProvider(String version) throws ClassNotFoundException { setVersion(luceneHnswScalarQuantizedVectorsFormat, version), setVersion(luceneHnswScalarQuantizedVectorsFormatFallback, version)); - // TODO: Find a better way if possible, but as a separate initiative. - if ("102".equals(version)) { + if ("104".equals(version)) { + // Binary quantized HNSW is still the lucene102 format (now in backward-codecs). binaryQuantizedVectorsFormat = loadClass( - setVersion(luceneBinaryQuantizedVectorsFormat, version), - setVersion(luceneBinaryQuantizedVectorsFormatFallback, version)); + setVersion(luceneBinaryQuantizedVectorsFormat, "102"), + setVersion(luceneBinaryQuantizedVectorsFormatFallback, "102")); hnswBinaryQuantizedVectorsFormat = loadClass( - setVersion(luceneHnswBinaryQuantizedVectorsFormat, version), - setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, version)); + setVersion(luceneHnswBinaryQuantizedVectorsFormat, "102"), + setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, "102")); } } @@ -258,17 +267,17 @@ public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws } } - public FlatVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( + public KnnVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( int maxConn, int beamWidth) throws Exception { try { - Constructor luceneHnswBinaryQuantizedVectorsFormatConstructor = - hnswBinaryQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); - return (FlatVectorsFormat) - luceneHnswBinaryQuantizedVectorsFormatConstructor.newInstance(maxConn, beamWidth); + Constructor ctor = + hnswBinaryQuantizedVectorsFormat.getConstructor( + Integer.TYPE, Integer.TYPE, Integer.TYPE, ExecutorService.class); + return (KnnVectorsFormat) ctor.newInstance(maxConn, beamWidth, 1, null); } catch (Exception e) { log.log( Level.SEVERE, - "Unable to initialize LuceneBinaryQuantizedVectorsFormat: " + e.getMessage()); + "Unable to initialize LuceneHnswBinaryQuantizedVectorsFormat: " + e.getMessage()); throw e; } } @@ -286,12 +295,12 @@ public FlatVectorsFormat getLuceneScalarQuantizedVectorsFormatInstance() throws } } - public FlatVectorsFormat getLuceneHnswScalarQuantizedVectorsFormatInstance( + public KnnVectorsFormat getLuceneHnswScalarQuantizedVectorsFormatInstance( int beamWidth, int maxConn) throws Exception { try { Constructor luceneHnswScalarQuantizedVectorsFormatConstructor = hnswScalarQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); - return (FlatVectorsFormat) + return (KnnVectorsFormat) luceneHnswScalarQuantizedVectorsFormatConstructor.newInstance(beamWidth, maxConn); } catch (Exception e) { log.log( diff --git a/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 6625ac72a4..86e77fcacb 100644 --- a/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/java/cuvs-lucene/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat -org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index d638180b06..1281225386 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -35,7 +35,7 @@ public void testNonexistentCodec() throws Exception { public void testExistingComponents() throws Exception { LuceneProvider provider = LuceneProvider.getInstance("99"); assertTrue(provider.getLuceneFlatVectorsFormatInstance(null) instanceof FlatVectorsFormat); - assertEquals(provider.getStaticIntParam("VERSION_CURRENT"), 0); + assertEquals(1, provider.getStaticIntParam("VERSION_CURRENT")); assertNotEquals(provider.getSimilarityFunctions().size(), 0); } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 2b960bf3cd..347b9cc4c4 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -7,6 +7,7 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import java.io.IOException; import java.util.List; import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; @@ -19,6 +20,7 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.search.AcceptDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; @@ -39,6 +41,16 @@ protected Codec getCodec() { return TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); } + @Override + protected boolean supportsFloatVectorFallback() { + return false; + } + + @Override + protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOException { + // Off-heap size checks assume Lucene's CPU reader; skip for the GPU format. + } + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); @@ -111,7 +123,9 @@ public void testTwoVectorFieldsPerDoc() throws Exception { assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); // opportunistically check boundary condition - search with a 0 topK - var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); + var topDocs = + r.searchNearestVectors( + "f1", randomVector(384), 0, AcceptDocs.fromLiveDocs(null, r.maxDoc()), 10); assertEquals(0, topDocs.scoreDocs.length); assertEquals(0, topDocs.totalHits.value()); } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 9bb140228b..f4f1e55839 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -19,6 +19,7 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.search.AcceptDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; @@ -39,6 +40,11 @@ protected Codec getCodec() { return TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); } + @Override + protected boolean supportsFloatVectorFallback() { + return false; + } + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); @@ -111,7 +117,9 @@ public void testTwoVectorFieldsPerDoc() throws Exception { assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); // opportunistically check boundary condition - search with a 0 topK - var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); + var topDocs = + r.searchNearestVectors( + "f1", randomVector(384), 0, AcceptDocs.fromLiveDocs(null, r.maxDoc()), 10); assertEquals(0, topDocs.scoreDocs.length); assertEquals(0, topDocs.totalHits.value()); } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java index df2ba21189..6794dcf618 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java @@ -15,6 +15,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.LeafReader; +import org.apache.lucene.search.AcceptDocs; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TopKnnCollector; import org.apache.lucene.store.ByteBuffersDirectory; @@ -77,7 +78,8 @@ public void testSelectiveFilterDoesNotLeakTrailingOrdinals() throws Exception { } for (float[] q : queries) { TopKnnCollector collector = new TopKnnCollector(topK, Integer.MAX_VALUE); - leaf.searchNearestVectors(VECTOR_FIELD, q, collector, acceptDocs); + leaf.searchNearestVectors( + VECTOR_FIELD, q, collector, AcceptDocs.fromLiveDocs(acceptDocs, leaf.maxDoc())); for (ScoreDoc hit : collector.topDocs().scoreDocs) { assertTrue( "per-segment search returned doc " + hit.doc + " outside filter category " + c, diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 5140782fc7..9a5c5ef1e3 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -26,6 +26,7 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.search.AcceptDocs; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; @@ -65,6 +66,11 @@ protected Codec getCodec() { return TestUtil.alwaysKnnVectorsFormat(knnVectorsFormat); } + @Override + protected boolean supportsFloatVectorFallback() { + return false; + } + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { final int R = 2, D = 128; float[][] f = new float[R][D]; @@ -174,7 +180,8 @@ public void testCosineSimilarity() throws Exception { assertEquals(R, values.size()); float[] queryVector = randomVector(D); - var topDocs = r.searchNearestVectors(F, queryVector, 2, null, 10); + AcceptDocs acceptAll = AcceptDocs.fromLiveDocs(null, r.maxDoc()); + var topDocs = r.searchNearestVectors(F, queryVector, 2, acceptAll, 10); assertTrue("Should return at least one result", topDocs.scoreDocs.length > 0); assertTrue("Scores should be non-negative", topDocs.scoreDocs[0].score >= 0); } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java index d95c60ae6b..97ac6b5f92 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java @@ -21,7 +21,6 @@ import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; import org.junit.Test; /** @@ -112,7 +111,7 @@ private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { assertEquals(a.numLevels(), b.numLevels()); for (int level = 0; level < a.numLevels(); level++) { - int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); + int[] nodes = sortedNodes(a.getNodesOnLevel(level)); for (int node : nodes) { assertArrayEquals( "node " + node + " at level " + level + " has different neighbors", @@ -122,6 +121,14 @@ private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception } } + private static int[] sortedNodes(HnswGraph.NodesIterator nodesOnLevel) { + int[] nodes = new int[nodesOnLevel.size()]; + int consumed = nodesOnLevel.consume(nodes); + assertEquals(nodesOnLevel.size(), consumed); + Arrays.sort(nodes); + return nodes; + } + private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { graph.seek(level, node); List arcs = new ArrayList<>();