diff --git a/examples/pom.xml b/examples/pom.xml index 768a9a87..a2af7e60 100644 --- a/examples/pom.xml +++ b/examples/pom.xml @@ -30,18 +30,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/pom.xml b/pom.xml index 00114972..bc66c14e 100644 --- a/pom.xml +++ b/pom.xml @@ -47,23 +47,23 @@ 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-test-framework - 10.2.0 + 10.4.0 test diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 9e663d2c..ab387a6a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -10,6 +10,7 @@ import java.util.Objects; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.function.Supplier; public class AcceleratedHNSWParams { @@ -17,33 +18,40 @@ public class AcceleratedHNSWParams { * TODO: Update boundaries for all parameters when a consensus is reached. * Issue: https://github.com/rapidsai/cuvs-lucene/issues/99 */ - private static final int MIN_WRITER_THREADS = 1; - private static final int MAX_WRITER_THREADS = 512; - private static final int MIN_INT_GRAPH_DEG = 2; - private static final int MAX_INT_GRAPH_DEG = 512; - private static final int MIN_GRAPH_DEG = 1; - private static final int MAX_GRAPH_DEG = 512; - private static final int MIN_HNSW_LAYERS = 1; - private static final int MAX_HNSW_LAYERS = 3; - private static final int MIN_MAX_CONN = 1; - private static final int MAX_MAX_CONN = 512; - private static final int MIN_BEAM_WIDTH = 1; - private static final int MAX_BEAM_WIDTH = 512; - private static final int MIN_NUM_MERGE_WORKERS = 1; - private static final int MAX_NUM_MERGE_WORKERS = 512; - - private static final int DEFAULT_WRITER_THREADS = 1; - private static final int DEFAULT_INT_GRAPH_DEGREE = 128; - private static final int DEFAULT_GRAPH_DEGREE = 64; - private static final int DEFAULT_HNSW_LAYERS = 1; - private static final int DEFAULT_MAX_CONN = 32; - private static final int DEFAULT_BEAM_WIDTH = 32; - private static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = + public static final int MIN_WRITER_THREADS = 1; + public static final int MAX_WRITER_THREADS = 512; + public static final int MIN_INT_GRAPH_DEG = 2; + public static final int MAX_INT_GRAPH_DEG = 512; + public static final int MIN_GRAPH_DEG = 1; + public static final int MAX_GRAPH_DEG = 512; + public static final int MIN_HNSW_LAYERS = 1; + public static final int MAX_HNSW_LAYERS = 3; + public static final int MIN_MAX_CONN = 1; + public static final int MAX_MAX_CONN = 512; + public static final int MIN_BEAM_WIDTH = 1; + public static final int MAX_BEAM_WIDTH = 512; + public static final int MIN_NUM_MERGE_WORKERS = 1; + public static final int MAX_NUM_MERGE_WORKERS = 512; + + public static final int DEFAULT_WRITER_THREADS = 1; + public static final int DEFAULT_INT_GRAPH_DEGREE = 128; + public static final int DEFAULT_GRAPH_DEGREE = 64; + public static final int DEFAULT_HNSW_LAYERS = 1; + public static final int DEFAULT_MAX_CONN = 32; + public static final int DEFAULT_BEAM_WIDTH = 32; + public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = CagraGraphBuildAlgo.NN_DESCENT; - private static final CuVSIvfPqParams DEFAULT_IVF_PQ_PARAMS = - new CuVSIvfPqParams.Builder().build(); - private static final int DEFAULT_NUM_MERGE_WORKERS = 1; - private static final ExecutorService DEFAULT_MERGE_EXE_SRVC = Executors.newFixedThreadPool(1); + public static final int DEFAULT_NUM_MERGE_WORKERS = 1; + + public static final Supplier DEFAULT_IVF_PQ_PARAMS = + () -> { + return new CuVSIvfPqParams.Builder().build(); + }; + + public static final Supplier DEFAULT_MERGE_EXE_SRVC = + () -> { + return Executors.newFixedThreadPool(DEFAULT_NUM_MERGE_WORKERS); + }; private final int writerThreads; private final int intermediateGraphDegree; @@ -57,7 +65,7 @@ public class AcceleratedHNSWParams { private final ExecutorService mergeExec; /** - * Constructs an instance of {@link GPUSearchParams} with specific parameter values. + * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. * * @param writerThreads Number of cuVS writer threads to use. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. @@ -67,7 +75,7 @@ public class AcceleratedHNSWParams { * @param maxConn The max connection parameter used when building HNSW index with the fallback mechanism. * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism. * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. - * @param cagraGraphBuildAlgo An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. + * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. * @param numMergeWorkers The number of merge workers to use with the fallback mechanism. * @param mergeExec The instance of {@link ExecutorService} to use with the fallback mechanism. */ @@ -177,7 +185,7 @@ public int getNumMergeWorkers() { } /** - * Get the instance of the {@link ExecutorService} to be used in the fallback mechanism * + * Get the instance of the {@link ExecutorService} to be used in the fallback mechanism * * @return the instance of the {@link ExecutorService} */ @@ -220,14 +228,14 @@ public static class Builder { private int maxConn = DEFAULT_MAX_CONN; private int beamWidth = DEFAULT_BEAM_WIDTH; private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO; - private CuVSIvfPqParams cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS; private int numMergeWorkers = DEFAULT_NUM_MERGE_WORKERS; - private ExecutorService mergeExec = DEFAULT_MERGE_EXE_SRVC; + private CuVSIvfPqParams cuVSIvfPqParams = null; + private ExecutorService mergeExec = null; /** * Set the number of cuVS writer threads while building the index * Valid range - Minimum: {@value MIN_WRITER_THREADS}, Maximum: {@value MAX_WRITER_THREADS} - * Default value - 64 + * Default value - {@value DEFAULT_WRITER_THREADS} * * @param writerThreads * @return instance of {@link Builder} @@ -240,7 +248,7 @@ public Builder withWriterThreads(int writerThreads) { /** * Set the intermediate graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} - * Default value - 128 + * Default value - {@value DEFAULT_INT_GRAPH_DEGREE} * * @param intermediateGraphDegree * @return instance of {@link Builder} @@ -253,7 +261,7 @@ public Builder withIntermediateGraphDegree(int intermediateGraphDegree) { /** * Set the graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_GRAPH_DEG}, Maximum: {@value MAX_GRAPH_DEG} - * Default value - 64 + * Default value - {@value DEFAULT_GRAPH_DEGREE} * * @param graphDegree * @return instance of {@link Builder} @@ -266,7 +274,7 @@ public Builder withGraphDegree(int graphDegree) { /** * Set the number of HNSW layers to construct while building the HNSW index * Valid range - Minimum: {@value MIN_HNSW_LAYERS}, Maximum: {@value MAX_HNSW_LAYERS} - * Default value - 2 + * Default value - {@value DEFAULT_HNSW_LAYERS} * * @param hnswLayers the number of HNSW layers * @return instance of {@link Builder} @@ -279,7 +287,7 @@ public Builder withHNSWLayer(int hnswLayers) { /** * Set the max connections parameter while building HNSW index with fallback mechanism * Valid range - Minimum: {@value MIN_MAX_CONN}, Maximum: {@value MAX_MAX_CONN} - * Default value - 8 + * Default value - {@value DEFAULT_MAX_CONN} * * @param maxConn the max connections parameter * @return instance of {@link Builder} @@ -292,7 +300,7 @@ public Builder withMaxConn(int maxConn) { /** * Set the beam width parameter while building HNSW index with fallback mechanism * Valid range - Minimum: {@value MIN_BEAM_WIDTH}, Maximum: {@value MAX_BEAM_WIDTH} - * Default value - 16 + * Default value - {@value DEFAULT_BEAM_WIDTH} * * @param beamWidth the beam width parameter * @return instance of {@link Builder} @@ -304,7 +312,7 @@ public Builder withBeamWidth(int beamWidth) { /** * Set the CAGRA graph build algorithm to use - * Default NN_DESCENT + * Default value - NN_DESCENT * * @param cagraGraphBuildAlgo * @return instance of {@link Builder} @@ -327,7 +335,7 @@ public Builder withCuVSIvfPqParams(CuVSIvfPqParams cuVSIvfPqParams) { /** * Set the number of merge workers to be used with the fallback mechanism - * Default value - 1 + * Default value - {@value DEFAULT_NUM_MERGE_WORKERS} * * @param numMergeWorkers number of merge workers to set * @return instance of {@link Builder} @@ -407,9 +415,6 @@ private void validate() throws IllegalArgumentException { if (Objects.isNull(cagraGraphBuildAlgo)) { throw new IllegalArgumentException("cagraGraphBuildAlgo cannot be null."); } - if (Objects.isNull(cuVSIvfPqParams)) { - throw new IllegalArgumentException("cuVSIvfPqParams cannot be null."); - } if (numMergeWorkers < MIN_NUM_MERGE_WORKERS || numMergeWorkers > MAX_NUM_MERGE_WORKERS) { throw new IllegalArgumentException( "numMergeWorkers not in valid range. Valid range: [" @@ -418,9 +423,6 @@ private void validate() throws IllegalArgumentException { + MAX_NUM_MERGE_WORKERS + "]"); } - if (Objects.isNull(mergeExec)) { - throw new IllegalArgumentException("mergeExec cannot be null."); - } } /** @@ -429,6 +431,12 @@ private void validate() throws IllegalArgumentException { * @return instance of {@link AcceleratedHNSWParams} */ public AcceleratedHNSWParams build() { + if (Objects.isNull(cuVSIvfPqParams)) { + cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS.get(); + } + if (Objects.isNull(mergeExec)) { + mergeExec = DEFAULT_MERGE_EXE_SRVC.get(); + } validate(); return new AcceleratedHNSWParams( writerThreads, diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 457ad22d..c4efdd9d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -34,7 +34,8 @@ public class AcceleratedHNSWUtils { public enum QuantizationType { BINARY, - SCALAR + SCALAR, + NONE } private static final LuceneProvider LUCENE_PROVIDER; @@ -42,7 +43,7 @@ public enum QuantizationType { static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); } catch (Exception e) { throw new ExceptionInInitializerError(e.getMessage()); @@ -84,7 +85,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( int size, int dimensions, CuVSMatrix adjacencyListMatrix, - List vectors, + List vectors, int hnswLayers, int graphDegree, CagraIndexParams params, @@ -132,17 +133,32 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( layerNodes.add(selectedNodes); - // Extract vectors for selected nodes - int bytesPerVector = (dimensions + 7) / 8; - byte[][] selectedVectors = new byte[nextLayerSize][]; - for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = vectors.get(selectedNodes[i]); - } + if (quantization == QuantizationType.NONE) { + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); + } + + // Build CAGRA graph for this layer + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + + } else { + + // Extract vectors for selected nodes + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + } - // Build CAGRA graph for this layer - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); + // Build CAGRA graph for this layer + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); + } // Update for next iteration currentLayerSize = nextLayerSize; @@ -160,7 +176,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( * Builds a CAGRA graph for a subset of binary quantized vectors */ private static CuVSMatrix buildCagraGraphForSubset( - byte[][] vectors, + Object vectors, int[] selectedNodes, int bytesPerVector, CagraIndexParams params, @@ -172,9 +188,12 @@ private static CuVSMatrix buildCagraGraphForSubset( if (quantization == QuantizationType.BINARY) { subsetDataset = - createByteMatrixFromArray(vectors, bytesPerVector, getCuVSResourcesInstance()); + createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); + } else if (quantization == QuantizationType.SCALAR) { + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); } else { - subsetDataset = createByteMatrixFromArray(vectors, dimensions, getCuVSResourcesInstance()); + subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } // Build CAGRA index for the subset @@ -211,8 +230,18 @@ 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) + * Returns a 2D array of offsets (information written while writing the meta info). Neighbor ids + * are written with {@link IndexOutput#writeGroupVInts(int[], int)} to match Apache Lucene 10.4 + * {@code Lucene99HnswVectorsFormat} (GroupVarInt encoding). * * @param graph instance of GPUBuiltHnswGraph * @param vectorIndex instance of IndexOutput @@ -226,7 +255,7 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde int[][] offsets = new int[graph.numLevels()][]; int[] scratch = new int[graph.maxConn() * 2]; for (int level = 0; level < graph.numLevels(); level++) { - int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); + int[] sortedNodes = getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; int nodeOffsetId = 0; @@ -259,10 +288,7 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde } // Write the size after duplicates are removed vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); - } + vectorIndex.writeGroupVInts(scratch, actualSize); offsets[level][nodeOffsetId++] = Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 2de761a7..c3befa23 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -29,7 +29,7 @@ public class CuVS2510GPUSearchCodec extends FilterCodec { * @throws Exception */ public CuVS2510GPUSearchCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getLucene104Codec()); initializeFormat(new GPUSearchParams.Builder().build()); } @@ -53,7 +53,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")); + this(NAME, LuceneProvider.getLucene104Codec()); initializeFormat(params); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index ccc61eae..949f7938 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -39,7 +39,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); } catch (Exception e) { diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index c783660d..4357193e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/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; @@ -69,7 +70,7 @@ public class CuVS2510GPUVectorsReader extends KnnVectorsReader { static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); } catch (Exception e) { throw new ExceptionInInitializerError(e.getMessage()); @@ -111,8 +112,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. @@ -392,11 +392,63 @@ private static FloatToFloatFunction getScoreNormalizationFunc(VectorSimilarityFu return score -> (1f / (1f + score)); } + /** + * Maps live-document acceptance to vector ordinals. {@link AcceptDocs#bits()} may be null in + * Lucene 10.4+ while {@link FloatVectorValues#getAcceptOrds} returns null when passed null bits; + * this method fills that gap for GPU prefilter construction. + */ + 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) { @@ -410,12 +462,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()); /* @@ -528,7 +581,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/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 2bb6732c..7bfb4ae2 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -73,7 +73,7 @@ public class CuVS2510GPUVectorsWriter extends KnnVectorsWriter { static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); } catch (Exception e) { throw new ExceptionInInitializerError(e.getMessage()); diff --git a/src/main/java/com/nvidia/cuvs/lucene/QuantizedFieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java similarity index 84% rename from src/main/java/com/nvidia/cuvs/lucene/QuantizedFieldWriter.java rename to src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 8a6bbfed..627d931d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/QuantizedFieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -17,10 +17,10 @@ import org.apache.lucene.index.FieldInfo; import org.apache.lucene.util.RamUsageEstimator; -public class QuantizedFieldWriter extends KnnFieldVectorsWriter { +public class FieldWriter extends KnnFieldVectorsWriter { private static final long SHALLOW_SIZE = - RamUsageEstimator.shallowSizeOfInstance(QuantizedFieldWriter.class); + RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); private final FieldInfo fieldInfo; private final FlatFieldVectorsWriter flatFieldVectorsWriter; @@ -28,7 +28,7 @@ public class QuantizedFieldWriter extends KnnFieldVectorsWriter { private QuantizationType quantizationType; @SuppressWarnings("unchecked") - public QuantizedFieldWriter( + public FieldWriter( QuantizationType quantizationType, FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { @@ -49,14 +49,20 @@ public void addValue(int docID, Object vectorValue) throws IOException { flatFieldVectorsWriter.addValue(docID, (float[]) vectorValue); } - List getVectors() { + List getByteVectors() { if (quantizationType == QuantizationType.BINARY) { return quantizeFloatVectorsToBinary(flatFieldVectorsWriter.getVectors()); - } else { + } else if (quantizationType == QuantizationType.SCALAR) { return quantizeFloatVectorsToScalar(flatFieldVectorsWriter.getVectors()); + } else { + throw new UnsupportedOperationException("Not applicable for QuantizationType.NONE"); } } + List getFloatVectors() { + return flatFieldVectorsWriter.getVectors(); + } + FieldInfo fieldInfo() { return fieldInfo; } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 6a9daae7..f06c8224 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -7,11 +7,11 @@ import java.io.IOException; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.AcceptDocs; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.knn.KnnCollectorManager; -import org.apache.lucene.util.Bits; /** * Extends upon KnnFloatVectorQuery for GPU-only search. @@ -43,7 +43,7 @@ public GPUKnnFloatVectorQuery( @Override protected TopDocs approximateSearch( LeafReaderContext context, - Bits acceptDocs, + AcceptDocs acceptDocs, int visitedLimit, KnnCollectorManager knnCollectorManager) throws IOException { diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index 86004141..d253555a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -12,7 +12,9 @@ import org.apache.lucene.codecs.KnnVectorsFormat; /** - * CuVS based codec for GPU based vector search + * A codec that enables GPU-based accelerated HNSW capability and can be used + * to accelerated indexing using GPUs and search using CPUs. Fallbacks to CPU + * based indexing when used on a machine without a GPU and/or cuVS. * * @since 25.10 */ @@ -22,25 +24,51 @@ public class Lucene101AcceleratedHNSWCodec extends FilterCodec { private static final String NAME = "Lucene101AcceleratedHNSWCodec"; private KnnVectorsFormat format; + /** + * Default constructor for {@link Lucene101AcceleratedHNSWCodec}. + * + * @throws Exception + */ public Lucene101AcceleratedHNSWCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getLucene104Codec()); } + /** + * Constructor for {@link Lucene101AcceleratedHNSWCodec}. + * + * @param name the codec's name + * @param delegate the delegate codec to filter + */ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { super(name, delegate); initializeFormatDefaultValues(); } + /** + * Constructor for {@link Lucene101AcceleratedHNSWCodec}. + * + * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} + * @throws Exception exception + */ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getLucene104Codec()); initializeFormat(acceleratedHNSWParams); } + /** + * Initialize an instance of {@link Lucene99AcceleratedHNSWVectorsFormat} + * with an instance of {@link AcceleratedHNSWParams} with default parameter values. + */ private void initializeFormatDefaultValues() { initializeFormat(new AcceleratedHNSWParams.Builder().build()); } + /** + * Initialize an instance of {@link Lucene99AcceleratedHNSWVectorsFormat}. + * + * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} to use + */ private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { try { format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams); @@ -53,7 +81,7 @@ private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { } /** - * Get the configured {@link KnnVectorsFormat} + * Get the configured {@link KnnVectorsFormat}. * * @return the instance of the {@link KnnVectorsFormat} */ @@ -63,7 +91,7 @@ public KnnVectorsFormat knnVectorsFormat() { } /** - * Set the {@link KnnVectorsFormat} + * Set the {@link KnnVectorsFormat}. * * @param format the {@link KnnVectorsFormat} to set */ diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 233313ec..49c7246b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -17,6 +17,7 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.TaskExecutor; /** * cuVS based KnnVectorsFormat for indexing on GPU and searching on the CPU. @@ -28,9 +29,8 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final Logger log = Logger.getLogger(Lucene99AcceleratedHNSWVectorsFormat.class.getName()); private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; - private static final Integer NUM_MERGE_WORKERS; - private static final int maxDimensions = 4096; - private AcceleratedHNSWParams acceleratedHNSWParams; + private static final int MAX_DIMENSIONS = 4096; + private final AcceleratedHNSWParams acceleratedHNSWParams; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; @@ -40,8 +40,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); - NUM_MERGE_WORKERS = LUCENE_PROVIDER.getStaticIntParam("DEFAULT_BEAM_WIDTH"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); } catch (Exception e) { @@ -50,7 +49,8 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { } /** - * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} with default parameter values. + * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} with an instance + * of {@link AcceleratedHNSWParams} with default parameter values. * * @throws LibraryException if the native library fails to load */ @@ -59,7 +59,8 @@ public Lucene99AcceleratedHNSWVectorsFormat() { } /** - * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} with the given threads, graph degree, etc. + * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} with an instance + * of {@link AcceleratedHNSWParams}. * * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} */ @@ -81,17 +82,15 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException log.log( Level.WARNING, "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); - // TODO: Make num merge workers configurable. try { return LUCENE_PROVIDER.getLuceneHnswVectorsWriterInstance( state, acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth(), flatWriter, - NUM_MERGE_WORKERS, - null); + acceleratedHNSWParams.getNumMergeWorkers(), + new TaskExecutor(acceleratedHNSWParams.getMergeExec())); } catch (Exception e) { - // maybe there is a better suited option to throwing RuntimeException? Need to explore. throw new RuntimeException(e.getMessage()); } } @@ -106,7 +105,6 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance( state, FLAT_VECTORS_FORMAT.fieldsReader(state)); } catch (Exception e) { - // maybe there is a better suited option to throwing RuntimeException? Need to explore. throw new RuntimeException(e.getMessage()); } } @@ -116,6 +114,6 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException */ @Override public int getMaxDimensions(String fieldName) { - return maxDimensions; + return MAX_DIMENSIONS; } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 5e74eba9..6653ec97 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -5,61 +5,48 @@ package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.cagraIndexParams; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; import static org.apache.lucene.index.VectorEncoding.FLOAT32; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.RowView; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import java.util.Objects; -import java.util.Random; -import java.util.SortedSet; -import java.util.TreeSet; -import java.util.logging.Logger; -import java.util.stream.IntStream; import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.codecs.KnnFieldVectorsWriter; -import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; -import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.FieldInfos; -import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; -import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.MergeState; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.Sorter; import org.apache.lucene.index.Sorter.DocMap; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.internal.hppc.IntObjectHashMap; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; -import org.apache.lucene.util.hnsw.NeighborArray; -import org.apache.lucene.util.packed.DirectMonotonicWriter; /** * This class extends upon the KnnVectorsWriter to - * enable the creation of GPU-based accelerated vector search indexes. + * enable the creation of GPU-based accelerated HNSW based vector search. * * @since 25.10 */ @@ -67,31 +54,24 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(Lucene99AcceleratedHNSWVectorsWriter.class); - - @SuppressWarnings("unused") - private static final Logger log = - Logger.getLogger(Lucene99AcceleratedHNSWVectorsWriter.class.getName()); - - private static final String CUVS_COMPONENT = "CUVS"; + private static final String COMPONENT = "Lucene99AcceleratedHNSWVectorsWriter"; private static final LuceneProvider LUCENE_PROVIDER; private static final Integer VERSION_CURRENT; - private static final List VECTOR_SIMILARITY_FUNCTIONS; - private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors - private final List fields = new ArrayList<>(); + private final AcceleratedHNSWParams acceleratedHNSWParams; + private final FlatVectorsWriter flatVectorsWriter; + private final List fields = new ArrayList<>(); private final InfoStream infoStream; - private IndexOutput cuvsIndex = null; - private IndexOutput hnswMeta = null, hnswVectorIndex = null; - private boolean finished; + private IndexOutput hnswMeta = null; + private IndexOutput hnswVectorIndex = null; private String vemFileName; private String vexFileName; - private final AcceleratedHNSWParams acceleratedHNSWParams; + private boolean finished; static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); - VECTOR_SIMILARITY_FUNCTIONS = LUCENE_PROVIDER.getSimilarityFunctions(); } catch (Exception e) { throw new ExceptionInInitializerError(e.getMessage()); } @@ -114,20 +94,15 @@ public Lucene99AcceleratedHNSWVectorsWriter( this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; this.acceleratedHNSWParams = acceleratedHNSWParams; - vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); - vexFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); - boolean success = false; try { - hnswMeta = state.directory.createOutput(vemFileName, state.context); hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); - CodecUtil.writeIndexHeader( hnswMeta, HNSW_META_CODEC_NAME, @@ -140,8 +115,8 @@ public Lucene99AcceleratedHNSWVectorsWriter( VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); - success = true; + printInfoStream(infoStream, COMPONENT, "Lucene99AcceleratedHNSWVectorsWriter is initialized"); } finally { if (success == false) { IOUtils.closeWhileHandlingException(this); @@ -156,55 +131,34 @@ public Lucene99AcceleratedHNSWVectorsWriter( public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { var encoding = fieldInfo.getVectorEncoding(); if (encoding != FLOAT32) { - throw new IllegalArgumentException("expected float32, got:" + encoding); + throw new IllegalArgumentException("Expected float32, got:" + encoding); } var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); - @SuppressWarnings("unchecked") - var flatWriter = (FlatFieldVectorsWriter) writer; - var cuvsFieldWriter = new GPUFieldWriter(fieldInfo, flatWriter); + var cuvsFieldWriter = new FieldWriter(QuantizationType.NONE, fieldInfo, writer); fields.add(cuvsFieldWriter); return writer; } /** - * A utility method to print info/debugging messages using InfoStream. - * - * @param msg the debugging message to print - */ - private void info(String msg) { - if (infoStream.isEnabled(CUVS_COMPONENT)) { - infoStream.message(CUVS_COMPONENT, msg); - } - } - - /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index + * Builds the intermediate CAGRA index and builds and writes the HNSW index. * * @param fieldInfo instance of FieldInfo that has the field description * @param vectors vectors to index * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo); + writeEmpty(fieldInfo, hnswMeta); + return; + } + if (vectors.size() < 2) { + writeSingleVectorGraph(fieldInfo, vectors); return; } - try { CuVSMatrix dataset = Utils.createFloatMatrix( vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); - - if (dataset.size() < 2) { - // Handle single vector case by creating a dummy HNSW graph - // TODO: Remove this workaround once https://github.com/rapidsai/cuvs/pull/1256 is merged - writeSingleVectorGraph(fieldInfo, vectors); - return; - } - - long startTime = System.nanoTime(); - CagraIndexParams params = cagraIndexParams( acceleratedHNSWParams.getWriterThreads(), @@ -217,14 +171,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro .withDataset(dataset) .withIndexParams(params) .build(); - - // Get the adjacency list from CAGRA index CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int size = (int) dataset.size(); int dimensions = fieldInfo.getVectorDimension(); - - // Create multi-layer HNSW graph from CAGRA GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, @@ -232,16 +181,13 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro dimensions, adjacencyListMatrix, vectors, - acceleratedHNSWParams.getHnswLayers()); - + acceleratedHNSWParams.getHnswLayers(), + acceleratedHNSWParams.getGraphdegree(), + params, + QuantizationType.NONE); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - - // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -250,283 +196,14 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro vectorIndexLength, size, hnswGraph, - graphLevelNodeOffsets); - - long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - startTime); - info("HNSW graph created in " + elapsedMillis + "ms, with " + dataset.size() + " vectors"); - + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); cagraIndex.close(); - } catch (Throwable t) { Utils.handleThrowable(t); } } - /** - * Creates a multi-layer HNSW graph with dynamic number of layers. - * M = cagraGraphDegree/2 - * Each layer contains 1/M nodes from the previous layer - * Creates layers until the highest layer has ≤ M nodes - */ - private GPUBuiltHnswGraph createMultiLayerHnswGraph( - FieldInfo fieldInfo, - int size, - int dimensions, - CuVSMatrix adjacencyListMatrix, - List vectors, - int hnswLayers) - throws Throwable { - - // Calculate M as cagraGraphDegree/2 - int M = acceleratedHNSWParams.getGraphdegree() / 2; - - // Store all layers data - List layerNodes = new ArrayList<>(); - List layerAdjacencies = new ArrayList<>(); - - // Layer 0: Use full CAGRA adjacency list - layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list - layerAdjacencies.add(adjacencyListMatrix); - - int currentLayerSize = size; - int layerIndex = 1; - Random random = new Random(); - - while (layerIndex < hnswLayers && currentLayerSize > 1) { - // Calculate size for next layer (1/M of current layer) - int nextLayerSize = Math.max(2, currentLayerSize / M); - // Select nodes for this layer - SortedSet selectedNodesSet = new TreeSet<>(); - - if (layerIndex == 1) { - // Select from all nodes (Layer 0) - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(random.nextInt(size)); - } - } else { - // Select from previous layer nodes - int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); - while (selectedNodesSet.size() < nextLayerSize) { - int idx = random.nextInt(prevLayerNodes.length); - selectedNodesSet.add(prevLayerNodes[idx]); - } - } - - // Convert to sorted array - int[] selectedNodes = - selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - - layerNodes.add(selectedNodes); - - // Extract vectors for selected nodes - float[][] selectedVectors = new float[nextLayerSize][]; - for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = vectors.get(selectedNodes[i]); - } - - // Build CAGRA graph for this layer - layerAdjacencies.add(buildCagraGraphForSubset(selectedVectors, selectedNodes)); - - // Update for next iteration - currentLayerSize = nextLayerSize; - layerIndex++; - - // Use different seed for each layer - random = new Random(new Random().nextLong()); - } - - // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); - } - - /** - * Builds a CAGRA graph for a subset of vectors - */ - private CuVSMatrix buildCagraGraphForSubset(float[][] vectors, int[] selectedNodes) - throws Throwable { - // Create CuVSMatrix from the subset vectors - CuVSMatrix subsetDataset = CuVSMatrix.ofArray(vectors); - - // Build CAGRA index for the subset - CagraIndexParams params = - cagraIndexParams( - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getCagraGraphBuildAlgo(), - acceleratedHNSWParams.getCuVSIvfPqParams()); - - CagraIndex subsetIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(subsetDataset) - .withIndexParams(params) - .build(); - - // Get adjacency list from subset CAGRA index - CuVSMatrix cagraGraph = subsetIndex.getGraph(); - - long numNodes = cagraGraph.size(); - long degree = cagraGraph.columns(); - - // Create a re-mapped adjacency list - int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; - - for (int i = 0; i < numNodes; i++) { - RowView rv = cagraGraph.getRow(i); - for (int j = 0; j < degree && j < rv.size(); j++) { - int subsetIndex1 = rv.getAsInt(j); - // Map subset index to original node ID - if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { - remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; - } else { - // Invalid index, use self-reference - remappedAdjacency[i][j] = selectedNodes[i]; - } - } - } - - subsetIndex.close(); - return CuVSMatrix.ofArray(remappedAdjacency); - } - - /** - * Writes the meta information for the index. - * - * @param vectorIndex instance of IndexOutput - * @param meta instance of IndexOutput - * @param field instance of FieldInfo - * @param vectorIndexOffset vector index offset - * @param vectorIndexLength vector index length - * @param count the count of vectors - * @param graph instance of HnswGraph - * @param graphLevelNodeOffsets graph level node offsets - * @throws IOException I/O Exceptions - */ - private void writeMeta( - IndexOutput vectorIndex, - IndexOutput meta, - FieldInfo field, - long vectorIndexOffset, - long vectorIndexLength, - int count, - HnswGraph graph, - int[][] graphLevelNodeOffsets) - throws IOException { - - meta.writeInt(field.number); - meta.writeInt(field.getVectorEncoding().ordinal()); - meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); - meta.writeVLong(vectorIndexOffset); - meta.writeVLong(vectorIndexLength); - meta.writeVInt(field.getVectorDimension()); - meta.writeInt(count); - meta.writeVInt(acceleratedHNSWParams.getGraphdegree() / 2); // M = cagraGraphDegree/2 - - // write graph nodes on each level - if (graph == null) { - meta.writeVInt(0); - } else { - meta.writeVInt(graph.numLevels()); - long valueCount = 0; - for (int level = 0; level < graph.numLevels(); level++) { - NodesIterator nodesOnLevel = graph.getNodesOnLevel(level); - valueCount += nodesOnLevel.size(); - if (level > 0) { - int[] nol = new int[nodesOnLevel.size()]; - int numberConsumed = nodesOnLevel.consume(nol); - Arrays.sort(nol); - assert numberConsumed == nodesOnLevel.size(); - meta.writeVInt(nol.length); // number of nodes on a level - for (int i = nodesOnLevel.size() - 1; i > 0; --i) { - nol[i] -= nol[i - 1]; - } - for (int n : nol) { - meta.writeVInt(n); - } - } else { - assert nodesOnLevel.size() == count : "Level 0 expects to have all nodes"; - } - } - - long start = vectorIndex.getFilePointer(); - meta.writeLong(start); - meta.writeVInt(16); // DIRECT_MONOTONIC_BLOCK_SHIFT); - - final DirectMonotonicWriter memoryOffsetsWriter = - DirectMonotonicWriter.getInstance(meta, vectorIndex, valueCount, 16); - long cumulativeOffsetSum = 0; - for (int[] levelOffsets : graphLevelNodeOffsets) { - for (int v : levelOffsets) { - memoryOffsetsWriter.add(cumulativeOffsetSum); - cumulativeOffsetSum += v; - } - } - - memoryOffsetsWriter.finish(); - - meta.writeLong(vectorIndex.getFilePointer() - start); - } - } - - /** - * Returns a 2D array of offsets (information written while writing the meta info) - * - * @param graph instance of GPUBuiltHnswGraph - * @param vectorIndex instance of IndexOutput - * @return a 2D array of offsets - * @throws IOException I/O Exceptions - */ - private int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { - // write vectors' neighbors on each level into the vectorIndex file - int countOnLevel0 = graph.size(); - int[][] offsets = new int[graph.numLevels()][]; - int[] scratch = new int[graph.maxConn() * 2]; - for (int level = 0; level < graph.numLevels(); level++) { - int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); - offsets[level] = new int[sortedNodes.length]; - int nodeOffsetId = 0; - - for (int node : sortedNodes) { - // Get node neighbors - NeighborArray neighbors = graph.getNeighbors(level, node); - // Get the size of the neighbor array - int size = neighbors.size(); - // Write size in VInt as the neighbors list is typically small - long offsetStart = vectorIndex.getFilePointer(); - // Get neighbors - int[] nnodes = neighbors.nodes(); - // Sort them - Arrays.sort(nnodes, 0, size); - // Now that we have sorted, do delta encoding to minimize the required bits to store the - // information - int actualSize = 0; - if (size > 0) { - scratch[0] = nnodes[0]; - actualSize = 1; - } - // De-duplication - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - // Sorting step helps here - if (nnodes[i - 1] == nnodes[i]) { - continue; - } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; - } - // Write the size after duplicates are removed - vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); - } - offsets[level][nodeOffsetId++] = - Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); - } - } - // Return offsets (information written while writing the meta info) - return offsets; - } - /** * Build the indexes and writes it to the disk. */ @@ -548,8 +225,8 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * @param fieldData * @throws IOException */ - private void writeField(GPUFieldWriter fieldData) throws IOException { - writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); + private void writeField(FieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getFloatVectors()); } /** @@ -559,18 +236,15 @@ private void writeField(GPUFieldWriter fieldData) throws IOException { * @param sortMap instance of the DocMap * @throws IOException */ - private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) - throws IOException { - + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); - final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord + final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); - List sortedVectors = new ArrayList(); - for (int i = 0; i < fieldData.getVectors().size(); i++) { - sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); + List floatVectors = fieldData.getFloatVectors(); + for (int i = 0; i < floatVectors.size(); i++) { + sortedVectors.add(floatVectors.get(new2OldOrd[i])); } - writeFieldInternal(fieldData.fieldInfo(), sortedVectors); } @@ -583,23 +257,13 @@ private void writeSortingField(GPUFieldWriter fieldData, Sorter.DocMap sortMap) */ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) throws IOException { - // Workaround for CAGRA not supporting single vector indexes - // TODO: Remove this method once https://github.com/rapidsai/cuvs/pull/1256 is merged try { int size = 1; int dimensions = fieldInfo.getVectorDimension(); - - // Create a dummy HNSW graph for a single vector GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - - // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -608,127 +272,17 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) vectorIndexLength, size, hnswGraph, - graphLevelNodeOffsets); - - long elapsedMillis = Utils.nanosToMillis(System.nanoTime() - System.nanoTime()); - info("Single vector HNSW graph created in " + elapsedMillis + "ms, with " + size + " vector"); - - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - - /** - * Creates a dummy HNSW graph for a single vector. - * The graph will have 1 level with 1 node and no neighbors. - */ - private GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimensions) throws Throwable { - // Create adjacency list for single node with no neighbors - int[][] singleNodeAdjacency = new int[][] {{-1}}; // -1 indicates no neighbors - - // Create CuVSMatrix from the adjacency list - CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency); - - // Create layer data for single-level graph - List layerNodes = new ArrayList<>(); - List layerAdjacencies = new ArrayList<>(); - - // Layer 0: contains all nodes (just the single node) - layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list - layerAdjacencies.add(adjacencyMatrix); - - // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); - } - - /** - * Writes an empty meta information for the field. - * - * @param fieldInfo instance of FieldInfo - * @throws IOException I/O Exceptions - */ - private void writeEmpty(FieldInfo fieldInfo) throws IOException { - writeMeta(null, hnswMeta, fieldInfo, 0, 0, 0, null, null); - } - - static int distFuncToOrd(VectorSimilarityFunction func) { - for (int i = 0; i < VECTOR_SIMILARITY_FUNCTIONS.size(); i++) { - if (VECTOR_SIMILARITY_FUNCTIONS.get(i).equals(func)) { - return (byte) i; - } - } - throw new IllegalArgumentException("invalid distance function: " + func); - } - - /** - * Uses the CAGRA API to merge the CAGRA indexes. - * - * @param fieldInfo instance of FieldInfo - * @param mergeState instance of MergeState - * @throws IOException I/O Exceptions - */ - private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { - try { - - List cagraIndexes = new ArrayList<>(); - // We need this count so that the merged segment's meta information has the vector count. - int totalVectorCount = 0; - - for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { - KnnVectorsReader knnReader = mergeState.knnVectorsReaders[i]; - // Access the CAGRA index for this field from the reader - - if (knnReader != null) { - if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { - if (cvr != null) { - totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); - CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); - if (cagraIndex != null) { - cagraIndexes.add(cagraIndex); - } - } - } else { - // This should never happen - throw new RuntimeException( - "Reader is not of CuVSVectorsReader type. Instead it is: " + knnReader.getClass()); - } - } - } - assert cagraIndexes.size() > 1; - - CagraIndex mergedIndex = - CagraIndex.merge(cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()])); - writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); - info("Successfully merged " + cagraIndexes.size() + " CAGRA indexes using native merge API"); - + graphLevelNodeOffsets, + acceleratedHNSWParams.getGraphdegree()); } catch (Throwable t) { Utils.handleThrowable(t); } } /** - * Creates List from merged vectors - */ - private List createListFromMergedVectors(FloatVectorValues mergedVectorValues) - throws IOException { - List vectors = new ArrayList(); - KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - float[] vector = mergedVectorValues.vectorValue(iter.index()); - vectors.add(vector.clone()); - } - return vectors; - } - - /** - * Fallback method that rebuilds indexes from merged vectors. - * Used when native CAGRA merge() is not possible. Also used - * when non-CAGRA index types are used (for e.g. Brute Force index). + * Create combined data set for the merged segment and call writeFieldInternal. */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { - if (fieldInfo.getVectorEncoding() != FLOAT32) { - throw new AssertionError("Only Float32 supported"); - } try { List dataset = createListFromMergedVectors( @@ -739,74 +293,13 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws } } - /** - * Extracts the CAGRA index for a specific field from a CuVSVectorsReader. - */ - private CagraIndex getCagraIndexFromReader(CuVS2510GPUVectorsReader reader, String fieldName) { - try { - IntObjectHashMap cuvsIndices = reader.getCuvsIndexes(); - FieldInfos fieldInfos = reader.getFieldInfos(); - FieldInfo fieldInfo = fieldInfos.fieldInfo(fieldName); - - if (fieldInfo != null) { - GPUIndex cuvsIndex = cuvsIndices.get(fieldInfo.number); - if (cuvsIndex != null) { - return cuvsIndex.getCagraIndex(); - } - } - } catch (Exception e) { - e.printStackTrace(); - info("Failed to extract CAGRA index for field " + fieldName + ": " + e.getMessage()); - } - return null; - } - - /** - * Writes a pre-built merged CAGRA index to the output. - */ - private void writeMergedCagraIndex(FieldInfo fieldInfo, CagraIndex mergedIndex, int vectorCount) - throws IOException { - try { - var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); - - // Serialize the merged index - Path tmpFile = - Files.createTempFile(getCuVSResourcesInstance().tempDirectory(), "mergedindex", "cag"); - mergedIndex.serialize(cagraIndexOutputStream, tmpFile); - - // TODO: Path to writeFieldInternal missing. Fix this. - - // Clean up the merged index - mergedIndex.close(); - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - /** * Write field for merging. */ @Override public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { flatVectorsWriter.mergeOneField(fieldInfo, mergeState); - - // Since CAGRA merge does not support merging of indexes with purging of deletes, - // we fallback to vector-based re-indexing. Issue: - // https://github.com/rapidsai/cuvs/issues/1253 - boolean hasDeletions = - IntStream.range(0, mergeState.liveDocs.length) - .anyMatch( - i -> - mergeState.liveDocs[i] == null - || IntStream.range(0, mergeState.maxDocs[i]) - .anyMatch(j -> !mergeState.liveDocs[i].get(j))); - - if (mergeState.knnVectorsReaders.length > 1 && !hasDeletions) { - mergeCagraIndexes(fieldInfo, mergeState); - } else { - // CAGRA's merge API does not handle the trivial case of merging 1 index. - vectorBasedMerge(fieldInfo, mergeState); - } + vectorBasedMerge(fieldInfo, mergeState); } /** @@ -819,11 +312,6 @@ public void finish() throws IOException { } finished = true; flatVectorsWriter.finish(); - - if (cuvsIndex != null) { - CodecUtil.writeFooter(cuvsIndex); - } - if (hnswMeta != null) { // write end of fields marker hnswMeta.writeInt(-1); @@ -839,7 +327,8 @@ public void finish() throws IOException { */ @Override public void close() throws IOException { - IOUtils.close(cuvsIndex, hnswMeta, hnswVectorIndex, flatVectorsWriter); + printInfoStream(infoStream, COMPONENT, "Closing resources"); + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); closeCuVSResourcesInstance(); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index f2c1aa37..96b0ec23 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/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.getLucene104Codec()); } 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.getLucene104Codec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 0f8d9602..e7678b49 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/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,12 +35,11 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector static { try { - LUCENE99_PROVIDER = LuceneProvider.getInstance("99"); - LUCENE102_PROVIDER = LuceneProvider.getInstance("102"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_CODEC_LINE); FLAT_VECTORS_FORMAT = - LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); } catch (Exception e) { - throw new ExceptionInInitializerError(e.getMessage()); + throw new ExceptionInInitializerError(e); } } @@ -79,14 +77,13 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, acceleratedHNSWParams, flatWriter); } else { try { - // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat format + // Lucene 10.4: binary-quantized HNSW lives in backward_codecs.lucene102 log.log( Level.WARNING, "GPU based indexing not supported, falling back to using the" - + " Lucene102HnswBinaryQuantizedVectorsFormat"); + + " Lucene102HnswBinaryQuantizedVectorsFormat (backward-codecs)"); KnnVectorsFormat fallbackFormat = - LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance( - acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); + LUCENE_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance(); return fallbackFormat.fieldsWriter(state); } catch (Exception e) { throw new RuntimeException(e.getMessage()); @@ -100,7 +97,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 new RuntimeException(e.getMessage()); diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 1bdc62db..ade0e702 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -62,7 +62,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector private static final String COMPONENT = "Lucene99AcceleratedHNSWBinaryQuantizedVectorsWriter"; private final FlatVectorsWriter flatVectorsWriter; - private final List fields = new ArrayList<>(); + private final List fields = new ArrayList<>(); private final InfoStream infoStream; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; @@ -133,7 +133,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException throw new IllegalArgumentException("expected float32, got:" + encoding); } var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); - var cuvsFieldWriter = new QuantizedFieldWriter(QuantizationType.BINARY, fieldInfo, writer); + var cuvsFieldWriter = new FieldWriter(QuantizationType.BINARY, fieldInfo, writer); fields.add(cuvsFieldWriter); return writer; } @@ -238,8 +238,8 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * @param fieldData * @throws IOException */ - private void writeField(QuantizedFieldWriter fieldData) throws IOException { - writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); + private void writeField(FieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getByteVectors()); } /** @@ -249,15 +249,14 @@ private void writeField(QuantizedFieldWriter fieldData) throws IOException { * @param sortMap instance of the DocMap * @throws IOException */ - private void writeSortingField(QuantizedFieldWriter fieldData, Sorter.DocMap sortMap) - throws IOException { + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); List sortedVectors = new ArrayList(); - List vectors = fieldData.getVectors(); + List vectors = fieldData.getByteVectors(); for (int i = 0; i < vectors.size(); i++) { sortedVectors.add(vectors.get(new2OldOrd[i])); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java index 0c7736a0..1b62ce8c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java +++ b/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.getLucene104Codec()); } 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.getLucene104Codec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java index 534390b8..bedc7d5b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -33,10 +33,10 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsFormat extends KnnVector static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_CODEC_LINE); FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneScalarQuantizedVectorsFormatInstance(); } catch (Exception e) { - throw new ExceptionInInitializerError(e.getMessage()); + throw new ExceptionInInitializerError(e); } } @@ -72,10 +72,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, acceleratedHNSWParams, flatWriter); } else { try { - // Fallback to Lucene's Lucene99HnswScalarQuantizedVectorsFormat + // Lucene 10.4 core 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/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 11a085ee..0ae1c8d9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -63,7 +63,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector private static final Integer VERSION_CURRENT; private final FlatVectorsWriter flatVectorsWriter; - private final List fields = new ArrayList<>(); + private final List fields = new ArrayList<>(); private final InfoStream infoStream; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; @@ -73,7 +73,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector static { try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + LUCENE_PROVIDER = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); } catch (Exception e) { throw new ExceptionInInitializerError(e.getMessage()); @@ -142,7 +142,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException throw new IllegalArgumentException("expected float32, got:" + encoding); } var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); - var cuvsFieldWriter = new QuantizedFieldWriter(QuantizationType.SCALAR, fieldInfo, writer); + var cuvsFieldWriter = new FieldWriter(QuantizationType.SCALAR, fieldInfo, writer); fields.add(cuvsFieldWriter); return writer; } @@ -166,7 +166,7 @@ private static byte[] convertSignedToUnsigned(byte[] signedVector) { * @param vectors quantized vectors * @throws IOException */ - private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { if (vectors.size() == 0) { writeEmpty(fieldInfo, hnswMeta); return; @@ -177,8 +177,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw // Convert 7-bit signed bytes to 8-bit unsigned bytes for cuVS compatibility List unsignedVectors = new ArrayList<>(vectors.size()); - for (byte[] signedVector : vectors) { - unsignedVectors.add(convertSignedToUnsigned(signedVector)); + for (Object signedVector : vectors) { + unsignedVectors.add(convertSignedToUnsigned((byte[]) signedVector)); } // Create CuVSMatrix with BYTE data type (unsigned bytes) @@ -264,8 +264,8 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * @param fieldData * @throws IOException */ - private void writeField(QuantizedFieldWriter fieldData) throws IOException { - writeFieldInternal(fieldData.fieldInfo(), fieldData.getVectors()); + private void writeField(FieldWriter fieldData) throws IOException { + writeFieldInternal(fieldData.fieldInfo(), fieldData.getByteVectors()); } /** @@ -275,16 +275,16 @@ private void writeField(QuantizedFieldWriter fieldData) throws IOException { * @param sortMap instance of the DocMap * @throws IOException */ - private void writeSortingField(QuantizedFieldWriter fieldData, Sorter.DocMap sortMap) - throws IOException { + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); List sortedVectors = new ArrayList(); - for (int i = 0; i < fieldData.getVectors().size(); i++) { - sortedVectors.add(fieldData.getVectors().get(new2OldOrd[i])); + List byteVectors = fieldData.getByteVectors(); + for (int i = 0; i < byteVectors.size(); i++) { + sortedVectors.add(byteVectors.get(new2OldOrd[i])); } writeFieldInternal(fieldData.fieldInfo(), sortedVectors); diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 7635e323..39f40c4b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -9,9 +9,12 @@ import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; 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; @@ -24,7 +27,8 @@ import org.apache.lucene.search.TaskExecutor; /** - * Dynamically loads Lucene format, reader, and writer classes with a fallback mechanism. + * Reflectively loads Lucene codec and vector format classes. This module targets Apache + * Lucene 10.4.x only ({@link #LUCENE_CODEC_LINE} / {@link #LUCENE_FLOAT_HNSW_LINE}). * * @since 25.12 */ @@ -32,6 +36,16 @@ public class LuceneProvider { static final Logger log = Logger.getLogger(LuceneProvider.class.getName()); + /** Lucene {@code Codec} SPI line id for the Lucene 10.4 core delegate (e.g. {@code Lucene104Codec}). */ + public static final String LUCENE_CODEC_LINE = "104"; + + /** + * Lucene 10.4 keeps float HNSW and the default {@code FlatVectorsFormat} under {@code + * codecs.lucene99} (there is no {@code Lucene104FlatVectorsFormat} / {@code Lucene104HnswVectors*} + * in {@code lucene-core}). + */ + public static final String LUCENE_FLOAT_HNSW_LINE = "99"; + private static final String BASE = "org.apache.lucene."; private static String codecs = "codecs.lucene."; private static String fallbackCodecs = "backward_codecs.lucene."; @@ -56,16 +70,6 @@ public class LuceneProvider { private static String luceneHnswVectorsWriterFallback = BASE + fallbackCodecs + "LuceneHnswVectorsWriter"; - private static String luceneBinaryQuantizedVectorsFormat = - BASE + codecs + "LuceneBinaryQuantizedVectorsFormat"; - private static String luceneBinaryQuantizedVectorsFormatFallback = - BASE + fallbackCodecs + "LuceneBinaryQuantizedVectorsFormat"; - - private static String luceneHnswBinaryQuantizedVectorsFormat = - BASE + codecs + "LuceneHnswBinaryQuantizedVectorsFormat"; - private static String luceneHnswBinaryQuantizedVectorsFormatFallback = - BASE + fallbackCodecs + "LuceneHnswBinaryQuantizedVectorsFormat"; - private static String luceneScalarQuantizedVectorsFormat = BASE + codecs + "LuceneScalarQuantizedVectorsFormat"; private static String luceneScalarQuantizedVectorsFormatFallback = @@ -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 ConcurrentHashMap<>(); private static MethodHandles.Lookup lookup = MethodHandles.lookup(); @@ -93,29 +97,43 @@ public class LuceneProvider { private Class hnswScalarQuantizedVectorsFormat; public static LuceneProvider getInstance(String version) throws ClassNotFoundException { - if (instance == null) { - instance = new LuceneProvider(version); + try { + return INSTANCES.computeIfAbsent(version, LuceneProvider::newInstanceUnchecked); + } catch (RuntimeException e) { + if (e.getCause() instanceof ClassNotFoundException cnfe) { + throw cnfe; + } + throw e; + } + } + + private static LuceneProvider newInstanceUnchecked(String version) { + try { + return new LuceneProvider(version); + } catch (ClassNotFoundException e) { + throw new RuntimeException(e); } - return instance; } private LuceneProvider(String version) throws ClassNotFoundException { + final String floatHnswPackageLine = + LUCENE_CODEC_LINE.equals(version) ? LUCENE_FLOAT_HNSW_LINE : version; flatVectorsFormat = loadClass( - setVersion(luceneFlatVectorsFormat, version), - setVersion(luceneFlatVectorsFormatFallback, version)); + setVersion(luceneFlatVectorsFormat, floatHnswPackageLine), + setVersion(luceneFlatVectorsFormatFallback, floatHnswPackageLine)); hnswVectorsFormat = loadClass( - setVersion(luceneHnswVectorsFormat, version), - setVersion(luceneHnswVectorsFormatFallback, version)); + setVersion(luceneHnswVectorsFormat, floatHnswPackageLine), + setVersion(luceneHnswVectorsFormatFallback, floatHnswPackageLine)); hnswVectorsReader = loadClass( - setVersion(luceneHnswVectorsReader, version), - setVersion(luceneHnswVectorsReaderFallback, version)); + setVersion(luceneHnswVectorsReader, floatHnswPackageLine), + setVersion(luceneHnswVectorsReaderFallback, floatHnswPackageLine)); hnswVectorsWriter = loadClass( - setVersion(luceneHnswVectorsWriter, version), - setVersion(luceneHnswVectorsWriterFallback, version)); + setVersion(luceneHnswVectorsWriter, floatHnswPackageLine), + setVersion(luceneHnswVectorsWriterFallback, floatHnswPackageLine)); scalarQuantizedVectorsFormat = loadClass( setVersion(luceneScalarQuantizedVectorsFormat, version), @@ -126,16 +144,13 @@ 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)) { - binaryQuantizedVectorsFormat = - loadClass( - setVersion(luceneBinaryQuantizedVectorsFormat, version), - setVersion(luceneBinaryQuantizedVectorsFormatFallback, version)); - hnswBinaryQuantizedVectorsFormat = - loadClass( - setVersion(luceneHnswBinaryQuantizedVectorsFormat, version), - setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, version)); + if (LUCENE_CODEC_LINE.equals(version)) { + final String lucene102BinaryFlat = + BASE + "backward_codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat"; + final String lucene102HnswBinary = + BASE + "backward_codecs.lucene102.Lucene102HnswBinaryQuantizedVectorsFormat"; + binaryQuantizedVectorsFormat = loadClass(lucene102BinaryFlat, lucene102BinaryFlat); + hnswBinaryQuantizedVectorsFormat = loadClass(lucene102HnswBinary, lucene102HnswBinary); } } @@ -173,6 +188,18 @@ public static Codec getCodec(String version) return (Codec) codecClassConstructor.newInstance(); } + /** Lucene 10.4 default codec delegate ({@link #LUCENE_CODEC_LINE}). */ + public static Codec getLucene104Codec() + throws ClassNotFoundException, + NoSuchMethodException, + SecurityException, + InstantiationException, + IllegalAccessException, + IllegalArgumentException, + InvocationTargetException { + return getCodec(LUCENE_CODEC_LINE); + } + public FlatVectorsFormat getLuceneFlatVectorsFormatInstance(FlatVectorsScorer scorer) throws Exception { try { @@ -258,17 +285,15 @@ public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws } } - public FlatVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( - int maxConn, int beamWidth) throws Exception { + /** Lucene 10.4 {@code Lucene102HnswBinaryQuantizedVectorsFormat} (backward-codecs) no-arg ctor. */ + public KnnVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance() throws Exception { try { - Constructor luceneHnswBinaryQuantizedVectorsFormatConstructor = - hnswBinaryQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); - return (FlatVectorsFormat) - luceneHnswBinaryQuantizedVectorsFormatConstructor.newInstance(maxConn, beamWidth); + Constructor ctor = hnswBinaryQuantizedVectorsFormat.getConstructor(); + return (KnnVectorsFormat) ctor.newInstance(); } catch (Exception e) { log.log( Level.SEVERE, - "Unable to initialize LuceneBinaryQuantizedVectorsFormat: " + e.getMessage()); + "Unable to initialize LuceneHnswBinaryQuantizedVectorsFormat: " + e.getMessage()); throw e; } } @@ -286,12 +311,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/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 91d7f5a0..c55aa7c3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -4,13 +4,18 @@ */ package com.nvidia.cuvs.lucene; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; import java.io.IOException; import java.time.Duration; +import java.util.ArrayList; import java.util.List; import java.util.logging.Level; import java.util.logging.Logger; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.util.InfoStream; /** @@ -170,10 +175,30 @@ static void handleThrowableWithIgnore(Throwable t, String msg) throws IOExceptio handleThrowable(t); } + /** + * Creates a list of float vectors from the input + * + * @param mergedVectorValues instance of {@link FloatVectorValues} + * @return a list of float arrays + * @throws IOException I/O Exception + */ + static List createListFromMergedVectors(FloatVectorValues mergedVectorValues) + throws IOException { + List vectors = new ArrayList(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + float[] vector = mergedVectorValues.vectorValue(iter.index()); + vectors.add(vector.clone()); + } + return vectors; + } + /** * Utility to print info/debug messages via InfoStream. * - * @param msg + * @param infoStream the writer's infostream + * @param component the name of the index writer + * @param msg the log message to push via the InfoStream */ static void info(InfoStream infoStream, String component, String msg) { if (infoStream.isEnabled(component)) { diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 6625ac72..86e77fca 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/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/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index 1fc9b8f2..66a657e6 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -5,7 +5,31 @@ package com.nvidia.cuvs.lucene; -import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_BEAM_WIDTH; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_INT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_MAX_CONN; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_NUM_MERGE_WORKERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_BEAM_WIDTH; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_INT_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_MAX_CONN; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_NUM_MERGE_WORKERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_BEAM_WIDTH; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_INT_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_MAX_CONN; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_NUM_MERGE_WORKERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_WRITER_THREADS; +import static java.lang.Integer.MAX_VALUE; +import static java.lang.Integer.MIN_VALUE; + import java.util.Random; import java.util.logging.Logger; import org.apache.lucene.tests.util.LuceneTestCase; @@ -24,20 +48,22 @@ public class TestAcceleratedHNSWParams extends LuceneTestCase { @Test public void testAcceleratedHNSWParamsDefaultValues() { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - assertEquals(32, params.getBeamWidth()); - assertEquals(64, params.getGraphdegree()); - assertEquals(1, params.getHnswLayers()); - assertEquals(128, params.getIntermediateGraphDegree()); - assertEquals(32, params.getMaxConn()); - assertEquals(1, params.getWriterThreads()); - assertEquals(1, params.getNumMergeWorkers()); - assertEquals(CagraGraphBuildAlgo.NN_DESCENT, params.getCagraGraphBuildAlgo()); + assertEquals(DEFAULT_BEAM_WIDTH, params.getBeamWidth()); + assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); + assertEquals(DEFAULT_HNSW_LAYERS, params.getHnswLayers()); + assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); + assertEquals(DEFAULT_MAX_CONN, params.getMaxConn()); + assertEquals(DEFAULT_WRITER_THREADS, params.getWriterThreads()); + assertEquals(DEFAULT_NUM_MERGE_WORKERS, params.getNumMergeWorkers()); + assertEquals(DEFAULT_CAGRA_GRAPH_BUILD_ALGO, params.getCagraGraphBuildAlgo()); } @Test public void testAcceleratedHNSWParamsInvalidBeamWidth() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(17, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_BEAM_WIDTH), random.nextInt(MAX_BEAM_WIDTH + 1, MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withBeamWidth(v).build()); @@ -47,7 +73,9 @@ public void testAcceleratedHNSWParamsInvalidBeamWidth() { @Test public void testAcceleratedHNSWParamsInvalidGraphDegree() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(65, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_GRAPH_DEG), random.nextInt(MAX_GRAPH_DEG + 1, MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withGraphDegree(v).build()); @@ -57,7 +85,9 @@ public void testAcceleratedHNSWParamsInvalidGraphDegree() { @Test public void testAcceleratedHNSWParamsInvalidHNSWLayers() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(3, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_HNSW_LAYERS), random.nextInt(MAX_HNSW_LAYERS + 1, MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withHNSWLayer(v).build()); @@ -67,7 +97,10 @@ public void testAcceleratedHNSWParamsInvalidHNSWLayers() { @Test public void testAcceleratedHNSWParamsInvalidIntGraphDegree() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(129, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_INT_GRAPH_DEG), + random.nextInt(MAX_INT_GRAPH_DEG + 1, MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withIntermediateGraphDegree(v).build()); @@ -77,7 +110,9 @@ public void testAcceleratedHNSWParamsInvalidIntGraphDegree() { @Test public void testAcceleratedHNSWParamsInvalidMaxConn() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(9, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_MAX_CONN), random.nextInt(MAX_MAX_CONN + 1, MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withMaxConn(v).build()); @@ -87,7 +122,10 @@ public void testAcceleratedHNSWParamsInvalidMaxConn() { @Test public void testAcceleratedHNSWParamsInvalidWriterThreads() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(2, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_WRITER_THREADS), + random.nextInt(MAX_WRITER_THREADS + 1, Integer.MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withWriterThreads(v).build()); @@ -97,20 +135,16 @@ public void testAcceleratedHNSWParamsInvalidWriterThreads() { @Test public void testAcceleratedHNSWParamsInvalidNumMergeWorkers() { for (int v : - new int[] {random.nextInt(Integer.MIN_VALUE, 1), random.nextInt(33, Integer.MAX_VALUE)}) { + new int[] { + random.nextInt(MIN_VALUE, MIN_NUM_MERGE_WORKERS), + random.nextInt(MAX_NUM_MERGE_WORKERS + 1, MAX_VALUE) + }) { assertThrows( IllegalArgumentException.class, () -> new AcceleratedHNSWParams.Builder().withNumMergeWorkers(v).build()); } } - @Test - public void testAcceleratedHNSWParamsInvalidMergeExecutorService() { - assertThrows( - IllegalArgumentException.class, - () -> new AcceleratedHNSWParams.Builder().withMergeExecutorService(null).build()); - } - @Test public void testAcceleratedHNSWParamsInvalidCagraGraphBuildAlgo() { assertThrows( @@ -118,13 +152,6 @@ public void testAcceleratedHNSWParamsInvalidCagraGraphBuildAlgo() { () -> new AcceleratedHNSWParams.Builder().withCagraGraphBuildAlgo(null).build()); } - @Test - public void testAcceleratedHNSWParamsInvalidCuVSIvfPqParams() { - assertThrows( - IllegalArgumentException.class, - () -> new AcceleratedHNSWParams.Builder().withCuVSIvfPqParams(null).build()); - } - @BeforeClass public static void beforeClass() { random = random(); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e660..882cef70 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -21,8 +21,8 @@ public class TestBackCompat { @Test public void testFallback() throws Exception { - // Lucene99Codec exists in the org.apache.lucene.backward_codecs.lucene99 - Codec c = LuceneProvider.getCodec("99"); + // Lucene 10.4 backward-codecs: Lucene99Codec (SPI line id matches float/HNSW package line "99") + Codec c = LuceneProvider.getCodec(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); assertEquals(c.getName(), "Lucene99"); } @@ -33,9 +33,10 @@ public void testNonexistentCodec() throws Exception { @Test public void testExistingComponents() throws Exception { - LuceneProvider provider = LuceneProvider.getInstance("99"); + LuceneProvider provider = LuceneProvider.getInstance(LuceneProvider.LUCENE_FLOAT_HNSW_LINE); assertTrue(provider.getLuceneFlatVectorsFormatInstance(null) instanceof FlatVectorsFormat); - assertEquals(provider.getStaticIntParam("VERSION_CURRENT"), 0); + // Lucene 10.4: Lucene99HnswVectorsFormat.VERSION_CURRENT == 1 (GroupVarInt graph encoding). + assertEquals(1, provider.getStaticIntParam("VERSION_CURRENT")); assertNotEquals(provider.getSimilarityFunctions().size(), 0); } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 96be6c01..281b3777 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -194,6 +194,10 @@ public void testSingleVectorIndex() throws Exception { @After public void afterTest() throws Exception { + // If setup was skipped (e.g. cuVS not supported), @Before didn't run and indexDirPath is null. + if (indexDirPath == null) { + return; + } File indexDirPathFile = indexDirPath.toFile(); if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { FileUtils.deleteDirectory(indexDirPathFile); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java index 04ad20a0..bd8c5248 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -165,6 +165,11 @@ public void testCagraToHnswSerializationAndSearchWithFallbackWriter() throws Exc public static void afterClass() throws Exception { // Reset resources for other tests to work setCuVSResourcesInstance(cuVSResourcesOrNull()); + // If setup was skipped (e.g. cuVS not supported), @BeforeClass didn't run and indexDirPath is + // null. + if (indexDirPath == null) { + return; + } File indexDirPathFile = indexDirPath.toFile(); if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { FileUtils.deleteDirectory(indexDirPathFile); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 1c75cb23..bcfa1a19 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/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,17 @@ protected Codec getCodec() { return TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); } + @Override + protected boolean supportsFloatVectorFallback() { + return false; + } + + @Override + protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOException { + // CuVS reader uses native/GPU memory; off-heap accounting is not aligned with base test + // expectations (getOffHeapByteSize keys/values). Skip to avoid false failures. + } + public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); @@ -111,7 +124,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/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 2428fc50..1a26f687 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -39,6 +39,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(); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 4c3c2809..aa4ebde2 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/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); }