From a7bc89d26e7fbd6e89b3508f4f29efca72c6e407 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 31 Dec 2025 17:08:39 -0500 Subject: [PATCH 1/4] Changes for #73 - Make CAGRA graph build algo configurable --- build.sh | 2 +- ci/release/update-version.sh | 8 +- examples/README.md | 35 ++++ examples/pom.xml | 159 ++++++++++++++++++ .../assembly/jar-with-merged-services.xml | 22 +++ .../examples/AcceleratedHnswExample.java | 130 ++++++++++++++ .../examples/IndexAndSearchonGPUExample.java | 132 +++++++++++++++ .../nvidia/cuvs/lucene/examples/Utils.java | 39 +++++ .../src/main/resources/logging.properties | 9 + .../cuvs/lucene/CuVS2510GPUSearchCodec.java | 44 ++++- .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 22 ++- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 6 +- .../lucene/Lucene101AcceleratedHNSWCodec.java | 22 ++- .../Lucene99AcceleratedHNSWVectorsFormat.java | 17 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 6 +- .../services/org.apache.lucene.codecs.Codec | 5 + .../org.apache.lucene.codecs.KnnVectorsFormat | 16 +- ...TestCagraToHnswSerializationAndSearch.java | 4 +- ...ializationAndSearchWithFallbackWriter.java | 4 +- .../com/nvidia/cuvs/lucene/TestMerge.java | 3 + 20 files changed, 656 insertions(+), 29 deletions(-) create mode 100644 examples/README.md create mode 100644 examples/pom.xml create mode 100644 examples/src/main/assembly/jar-with-merged-services.xml create mode 100644 examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java create mode 100644 examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java create mode 100644 examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java create mode 100644 examples/src/main/resources/logging.properties rename src/{test => main}/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java (67%) rename src/{test => main}/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java (83%) create mode 100644 src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec diff --git a/build.sh b/build.sh index bf768069..a0054e43 100755 --- a/build.sh +++ b/build.sh @@ -9,7 +9,7 @@ ARGS="$*" NUMARGS=$# VERSION="26.02.0" # Note: The version is updated automatically when ci/release/update-version.sh is invoked -GROUP_ID="com.nvidia.cuvs" +GROUP_ID="com.nvidia.cuvs.lucene" function hasArg { (( NUMARGS != 0 )) && (echo " ${ARGS} " | grep -q " $1 ") diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index afe73f64..f98aefbc 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -105,7 +105,9 @@ echo "${RAPIDS_BRANCH_NAME}" > RAPIDS_BRANCH # Update Java version NEXT_FULL_JAVA_TAG="${NEXT_SHORT_TAG}.${PATCH_PEP440}" sed_runner "s/VERSION=\".*\"/VERSION=\"${NEXT_FULL_JAVA_TAG}\"/g" build.sh -sed_runner "/.*/s//${NEXT_FULL_JAVA_TAG}<\/version>/g" pom.xml +for FILE in pom.xml examples/pom.xml; do + sed_runner "/.*/s//${NEXT_FULL_JAVA_TAG}<\/version>/g" "${FILE}" +done sed_runner "s| CuVS [[:digit:]]\{2\}\.[[:digit:]]\{2\} | CuVS ${NEXT_SHORT_TAG} |g" README.md @@ -117,3 +119,7 @@ for FILE in .github/workflows/*.yaml; do sed_runner "/shared-workflows/ s|@.*|@${RAPIDS_BRANCH_NAME}|g" "${FILE}" sed_runner "s/:[0-9]*\\.[0-9]*-/:${NEXT_SHORT_TAG}-/g" "${FILE}" done + +# Update versions in the examples readme file +sed_runner "s/rapidsai\/ci-conda:[\.0-9]*-cuda/rapidsai\/ci-conda:${NEXT_SHORT_TAG}-cuda/g" "examples/README.md" +sed_runner "s/target\/examples-[\.0-9]*-jar/target\/examples-${NEXT_FULL_JAVA_TAG}-jar/g" "examples/README.md" diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 00000000..acebe9e2 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,35 @@ +# Examples + +This maven project contains basic examples that showcase how `cuvs-lucene` can be used. + +## Prerequisites +- [Docker](https://www.docker.com/) +- [Nvidia Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) +- A machine with an Nvidia GPU + +## Steps + +If you are currently in this directory (and to be in the `cuvs-lucene's` root directory) do: +```sh +cd .. +``` + +Then do: +```sh +docker run --rm --gpus all --pull=always --volume $PWD:$PWD --workdir $PWD -it rapidsai/ci-conda:26.02-cuda13.0.2-ubuntu24.04-py3.13 +``` + +Inside the docker container (and in the `cuvs-lucene's` root directory) do: +```sh +./ci/build_java.sh && conda activate java && cd examples +``` + +To run Accelerated HNSW example do: +```sh +mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.02.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.AcceleratedHnswExample +``` + +To run the Index and Search on GPU example do: +```sh +mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.02.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample +``` diff --git a/examples/pom.xml b/examples/pom.xml new file mode 100644 index 00000000..ecdb5b4f --- /dev/null +++ b/examples/pom.xml @@ -0,0 +1,159 @@ + + + + 4.0.0 + + com.nvidia.cuvs.lucene.examples + examples + 26.02.0 + examples + + + 22 + 22 + UTF-8 + UTF-8 + + + + + + + + + + org.apache.lucene + lucene-core + 10.2.0 + + + org.apache.lucene + lucene-codecs + 10.2.0 + test + + + org.apache.lucene + lucene-backward-codecs + 10.2.0 + + + commons-io + commons-io + 2.18.0 + + + com.nvidia.cuvs.lucene + cuvs-lucene + 26.02.0 + + + + + + + com.diffplug.spotless + spotless-maven-plugin + 2.44.5 + + + + apply + + validate + + + + + + src/**/*.java + + + 1.27.0 + + true + false + + + ../license-header.txt + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.6.0 + + + src/main/assembly/jar-with-merged-services.xml + + + + true + + + true + true + + + + + + make-assembly + package + + single + + + + + + + + + maven-clean-plugin + 3.4.0 + + + maven-resources-plugin + 3.3.1 + + + maven-compiler-plugin + 3.13.0 + + + maven-surefire-plugin + 3.3.0 + + + maven-jar-plugin + 3.4.2 + + + maven-install-plugin + 3.1.2 + + + maven-deploy-plugin + 3.1.2 + + + maven-site-plugin + 3.12.1 + + + maven-project-info-reports-plugin + 3.6.1 + + + + + diff --git a/examples/src/main/assembly/jar-with-merged-services.xml b/examples/src/main/assembly/jar-with-merged-services.xml new file mode 100644 index 00000000..dc670618 --- /dev/null +++ b/examples/src/main/assembly/jar-with-merged-services.xml @@ -0,0 +1,22 @@ + + jar-with-merged-services + + jar + + false + + + / + true + true + runtime + + + + + metaInf-services + + + diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java new file mode 100644 index 00000000..fd7a8808 --- /dev/null +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static com.nvidia.cuvs.lucene.examples.Utils.generateDataset; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +public class AcceleratedHnswExample { + + private static Logger log = Logger.getLogger(AcceleratedHnswExample.class.getName()); + + private static Random random; + private static Path indexDirPath; + + public static void main(String[] args) throws Exception { + + Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.log(Level.FINE, "Successfully opened index"); + + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + log.log( + Level.FINE, + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + } + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + } + } + + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java new file mode 100644 index 00000000..51f0aecd --- /dev/null +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java @@ -0,0 +1,132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static com.nvidia.cuvs.lucene.examples.Utils.generateDataset; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; +import com.nvidia.cuvs.lucene.GPUKnnFloatVectorQuery; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +public class IndexAndSearchonGPUExample { + + private static Logger log = Logger.getLogger(IndexAndSearchonGPUExample.class.getName()); + + private static Random random; + private static Path indexDirPath; + + public static void main(String[] args) throws Exception { + + Codec codec = new CuVS2510GPUSearchCodec(); + IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); + + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + + final int COMMIT_FREQ = 2000; + final String ID_FIELD = "id"; + final String VECTOR_FIELD = "vector_field"; + + int numDocs = 2000; + int dimension = 32; + int topK = 5; + int count = COMMIT_FREQ; + float[][] dataset = generateDataset(random, numDocs, dimension); + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + indexWriter.addDocument(document); + count -= 1; + if (count == 0) { + indexWriter.commit(); + count = COMMIT_FREQ; + } + } + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + log.log(Level.FINE, "Successfully opened index"); + + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + log.log( + Level.FINE, + VECTOR_FIELD + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + } + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + float[] queryVector = generateDataset(random, 1, dimension)[0]; + log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK, null, topK, 1); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String id = doc.get(ID_FIELD); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + } + } + + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java new file mode 100644 index 00000000..901f4b21 --- /dev/null +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import java.util.Random; + +public class Utils { + + public static float[][] generateDataset(Random random, int size, int dimensions) { + float[][] dataset = new float[size][dimensions]; + for (int i = 0; i < size; i++) { + for (int j = 0; j < dimensions; j++) { + dataset[i][j] = random.nextFloat() * 100; + } + } + return dataset; + } + + public static float[] generateRandomVector(int dimensions, Random random) { + float[] vector = new float[dimensions]; + for (int i = 0; i < dimensions; i++) { + vector[i] = random.nextFloat() * 100; + } + return vector; + } + + public static float[][] generateQueries(Random random, int dimensions, int numQueries) { + // Generate random query vectors + float[][] queries = new float[numQueries][dimensions]; + for (int i = 0; i < numQueries; i++) { + for (int j = 0; j < dimensions; j++) { + queries[i][j] = random.nextFloat() * 100; + } + } + return queries; + } +} diff --git a/examples/src/main/resources/logging.properties b/examples/src/main/resources/logging.properties new file mode 100644 index 00000000..53f2c427 --- /dev/null +++ b/examples/src/main/resources/logging.properties @@ -0,0 +1,9 @@ +# Set root level +.level=FINEST + +# Define handlers +handlers=java.util.logging.ConsoleHandler + +java.util.logging.ConsoleHandler.level=FINEST +java.util.logging.ConsoleHandler.formatter=java.util.logging.SimpleFormatter +java.util.logging.SimpleFormatter.format=[%1$tF %1$tT] %4$s %2$s %5$s%6$s%n diff --git a/src/test/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java similarity index 67% rename from src/test/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 57b1e3c9..82ba6135 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.LibraryException; import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.lang.reflect.InvocationTargetException; @@ -27,6 +28,8 @@ public class CuVS2510GPUSearchCodec extends FilterCodec { private static final int DEFAULT_CUVS_WRITER_THREADS = 1; private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; private static final int DEFAULT_GRAPH_DEGREE = 64; + private static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = + CagraGraphBuildAlgo.NN_DESCENT; private static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; private KnnVectorsFormat format; @@ -60,13 +63,46 @@ public CuVS2510GPUSearchCodec() */ public CuVS2510GPUSearchCodec(String name, Codec delegate) { super(name, delegate); + initializeFormat( + DEFAULT_INTERMEDIATE_GRAPH_DEGREE, + DEFAULT_GRAPH_DEGREE, + DEFAULT_CUVS_WRITER_THREADS, + DEFAULT_CAGRA_GRAPH_BUILD_ALGO, + DEFAULT_INDEX_TYPE); + } + + public CuVS2510GPUSearchCodec( + int cuvsWriterThreads, + int intermediateGraphDegree, + int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, + IndexType indexType) + throws ClassNotFoundException, + NoSuchMethodException, + SecurityException, + InstantiationException, + IllegalAccessException, + IllegalArgumentException, + InvocationTargetException { + this(NAME, LuceneProvider.getCodec("101")); + initializeFormat( + cuvsWriterThreads, intermediateGraphDegree, graphDegree, cagraGraphBuildAlgo, indexType); + } + + private void initializeFormat( + int cuvsWriterThreads, + int intermediateGraphDegree, + int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, + IndexType indexType) { try { format = new CuVS2510GPUVectorsFormat( - DEFAULT_CUVS_WRITER_THREADS, - DEFAULT_INTERMEDIATE_GRAPH_DEGREE, - DEFAULT_GRAPH_DEGREE, - DEFAULT_INDEX_TYPE); + cuvsWriterThreads, + intermediateGraphDegree, + graphDegree, + cagraGraphBuildAlgo, + indexType); setKnnFormat(format); } catch (LibraryException ex) { log.log( diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 99639b18..2406a633 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -6,6 +6,7 @@ import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; @@ -39,6 +40,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { static final int DEFAULT_WRITER_THREADS = 32; static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; static final int DEFAULT_GRAPH_DEGREE = 64; + static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = CagraGraphBuildAlgo.NN_DESCENT; static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; static CuVSResources resources = cuVSResourcesOrNull(); @@ -49,6 +51,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { final int cuvsWriterThreads; final int intGraphDegree; final int graphDegree; + final CagraGraphBuildAlgo cagraGraphBuildAlgo; final CuVS2510GPUVectorsWriter.IndexType indexType; // the index type to build, when writing static { @@ -71,6 +74,7 @@ public CuVS2510GPUVectorsFormat() { DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, + DEFAULT_CAGRA_GRAPH_BUILD_ALGO, DEFAULT_INDEX_TYPE); } @@ -80,16 +84,22 @@ public CuVS2510GPUVectorsFormat() { * @param cuvsWriterThreads the number of cuVS writer threads to use * @param intGraphDegree the intermediate graph degree for building the CAGRA index * @param graphDegree the graph degree for building the CAGRA index + * @param cagraGraphBuildAlgo the CAGRA graph build algorithm to use * @param indexType the {@link com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType} * * @throws LibraryException if the native library fails to load */ public CuVS2510GPUVectorsFormat( - int cuvsWriterThreads, int intGraphDegree, int graphDegree, IndexType indexType) { + int cuvsWriterThreads, + int intGraphDegree, + int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, + IndexType indexType) { super("CuVS2510GPUVectorsFormat"); this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; this.indexType = indexType; } @@ -101,7 +111,14 @@ public CuVS2510GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOE checkSupported(); var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); return new CuVS2510GPUVectorsWriter( - state, cuvsWriterThreads, intGraphDegree, graphDegree, indexType, resources, flatWriter); + state, + cuvsWriterThreads, + intGraphDegree, + graphDegree, + cagraGraphBuildAlgo, + indexType, + resources, + flatWriter); } /** @@ -130,6 +147,7 @@ public String toString() { sb.append("(cuvsWriterThreads=").append(cuvsWriterThreads); sb.append("intGraphDegree=").append(intGraphDegree); sb.append("graphDegree=").append(graphDegree); + sb.append("cagraGraphBuildAlgo=").append(cagraGraphBuildAlgo); sb.append("resources=").append(resources); sb.append(")"); return sb.toString(); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 033bbc14..66c5bd92 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -77,6 +77,7 @@ public class CuVS2510GPUVectorsWriter extends KnnVectorsWriter { private final int cuvsWriterThreads; private final int intGraphDegree; private final int graphDegree; + private final CagraGraphBuildAlgo cagraGraphBuildAlgo; private final CuVSResources resources; private final IndexType indexType; @@ -142,6 +143,7 @@ public boolean bruteForce() { * @param cuvsWriterThreads the number of cuVS writer threads * @param intGraphDegree the intermediate graph degree for building the CAGRA index * @param graphDegree the graph degree for building the CAGRA index + * @param cagraGraphBuildAlgo the CAGRA graph build algorithm to use * @param indexType the IndexType * @param resources instance of the CuVSResources * @param flatVectorsWriter instance of FlatVectorsWriter @@ -153,6 +155,7 @@ public CuVS2510GPUVectorsWriter( int cuvsWriterThreads, int intGraphDegree, int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, IndexType indexType, CuVSResources resources, FlatVectorsWriter flatVectorsWriter) @@ -162,6 +165,7 @@ public CuVS2510GPUVectorsWriter( this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; this.resources = resources; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; @@ -248,7 +252,7 @@ private CagraIndexParams cagraIndexParams(int size) { .withNumWriterThreads(cuvsWriterThreads) .withIntermediateGraphDegree(intGraphDegree) .withGraphDegree(graphDegree) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withCagraGraphBuildAlgo(cagraGraphBuildAlgo) .build(); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java similarity index 83% rename from src/test/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java rename to src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index 9b0db693..d73112fd 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.LibraryException; import java.lang.reflect.InvocationTargetException; import java.util.logging.Level; @@ -24,6 +25,8 @@ public class Lucene101AcceleratedHNSWCodec extends FilterCodec { private static final int DEFAULT_CUVS_WRITER_THREADS = 1; private static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; private static final int DEFAULT_GRAPH_DEGREE = 64; + private static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = + CagraGraphBuildAlgo.NN_DESCENT; private static final int DEFAULT_HNSW_LAYERS = 1; private static final String NAME = "Lucene101AcceleratedHNSWCodec"; private static final LuceneProvider lucene99Provider; @@ -62,6 +65,7 @@ public Lucene101AcceleratedHNSWCodec( int cuvsWriterThreads, int intGraphDegree, int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, int hnswLayers, int maxConn, int beamWidth) @@ -74,7 +78,13 @@ public Lucene101AcceleratedHNSWCodec( InvocationTargetException { this(NAME, LuceneProvider.getCodec("101")); initializeFormat( - cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, maxConn, beamWidth); + cuvsWriterThreads, + intGraphDegree, + graphDegree, + cagraGraphBuildAlgo, + hnswLayers, + maxConn, + beamWidth); } private void initializeFormatDefaultValues() { @@ -82,6 +92,7 @@ private void initializeFormatDefaultValues() { DEFAULT_CUVS_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, + DEFAULT_CAGRA_GRAPH_BUILD_ALGO, DEFAULT_HNSW_LAYERS, maxConn, beamWidth); @@ -91,13 +102,20 @@ private void initializeFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, int hnswLayers, int maxConn, int beamWidth) { try { format = new Lucene99AcceleratedHNSWVectorsFormat( - cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, maxConn, beamWidth); + cuvsWriterThreads, + intGraphDegree, + graphDegree, + cagraGraphBuildAlgo, + hnswLayers, + maxConn, + beamWidth); setKnnFormat(format); } catch (LibraryException ex) { log.log( diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 5128ce7f..2bc62ded 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -6,6 +6,7 @@ import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.LibraryException; import java.io.IOException; @@ -32,6 +33,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { static final int DEFAULT_WRITER_THREADS = 32; static final int DEFAULT_INTERMEDIATE_GRAPH_DEGREE = 128; static final int DEFAULT_GRAPH_DEGREE = 64; + static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = CagraGraphBuildAlgo.NN_DESCENT; static final int DEFAULT_HNSW_GRAPH_LAYERS = 1; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; @@ -51,6 +53,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private final int cuvsWriterThreads; private final int intGraphDegree; private final int graphDegree; + private final CagraGraphBuildAlgo cagraGraphBuildAlgo; private final int hnswLayers; private final int maxConn; private final int beamWidth; @@ -78,6 +81,7 @@ public Lucene99AcceleratedHNSWVectorsFormat() { DEFAULT_WRITER_THREADS, DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, + DEFAULT_CAGRA_GRAPH_BUILD_ALGO, DEFAULT_HNSW_GRAPH_LAYERS, MAX_CONN, BEAM_WIDTH); @@ -89,6 +93,7 @@ public Lucene99AcceleratedHNSWVectorsFormat() { * @param cuvsWriterThreads number of cuVS threads to use while building the CAGRA index * @param intGraphDegree the intermediate graph degree while building the CAGRA index * @param graphDegree the graph degree to use while building the CAGRA index + * @param cagraGraphBuildAlgo the CAGRA graph build algorithm to use * @param hnswLayers the number of HNSW layers to construct in the HNSW graph * @param maxConn the maximum connections for the HNSW graph * @param beamWidth the beam width to use while building the HNSW graph @@ -97,6 +102,7 @@ public Lucene99AcceleratedHNSWVectorsFormat( int cuvsWriterThreads, int intGraphDegree, int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, int hnswLayers, int maxConn, int beamWidth) { @@ -104,6 +110,7 @@ public Lucene99AcceleratedHNSWVectorsFormat( this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; this.hnswLayers = hnswLayers; this.maxConn = maxConn; this.beamWidth = beamWidth; @@ -118,7 +125,14 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException if (supported()) { log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter( - state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter); + state, + cuvsWriterThreads, + intGraphDegree, + graphDegree, + cagraGraphBuildAlgo, + hnswLayers, + resources, + flatWriter); } else { log.log( Level.WARNING, @@ -165,6 +179,7 @@ public String toString() { sb.append("(cuvsWriterThreads=").append(cuvsWriterThreads); sb.append("intGraphDegree=").append(intGraphDegree); sb.append("graphDegree=").append(graphDegree); + sb.append("cagraGraphBuildAlgo=").append(cagraGraphBuildAlgo); sb.append("hnswLayers=").append(hnswLayers); sb.append("resources=").append(resources); sb.append(")"); diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 8528ef45..c754668e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -81,6 +81,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final int cuvsWriterThreads; private final int intGraphDegree; private final int graphDegree; + private final CagraGraphBuildAlgo cagraGraphBuildAlgo; private final int hnswLayers; // Number of layers to create in CAGRA->HNSW conversion private final CuVSResources resources; private final FlatVectorsWriter flatVectorsWriter; // for writing the raw vectors @@ -109,6 +110,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { * @param cuvsWriterThreads number of cuVS threads to use while building the intermediate CAGRA index * @param intGraphDegree the intermediate graph degree to use while building the CAGRA index * @param graphDegree the graph degree to use while building the CAGRA index + * @param cagraGraphBuildAlgo the CAGRA graph build algorithm to use * @param hnswLayers the number of hnsw layers to construct while building the HNSW graph * @param resources instance of the {@link com.nvidia.cuvs.CuVSResources} * @param flatVectorsWriter instance of the {@link org.apache.lucene.codecs.hnsw.FlatVectorsWriter} @@ -119,6 +121,7 @@ public Lucene99AcceleratedHNSWVectorsWriter( int cuvsWriterThreads, int intGraphDegree, int graphDegree, + CagraGraphBuildAlgo cagraGraphBuildAlgo, int hnswLayers, CuVSResources resources, FlatVectorsWriter flatVectorsWriter) @@ -127,6 +130,7 @@ public Lucene99AcceleratedHNSWVectorsWriter( this.cuvsWriterThreads = cuvsWriterThreads; this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; this.hnswLayers = hnswLayers; this.resources = resources; this.flatVectorsWriter = flatVectorsWriter; @@ -210,7 +214,7 @@ private CagraIndexParams cagraIndexParams() { .withNumWriterThreads(cuvsWriterThreads) .withIntermediateGraphDegree(intGraphDegree) .withGraphDegree(graphDegree) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withCagraGraphBuildAlgo(cagraGraphBuildAlgo) .build(); } diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec new file mode 100644 index 00000000..f01d2e8d --- /dev/null +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec +com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec 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 747a1575..4e20fcb5 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 @@ -1,17 +1,5 @@ -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 33516340..ed63c350 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -7,6 +7,7 @@ import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import java.io.File; import java.io.IOException; import java.lang.reflect.InvocationTargetException; @@ -68,7 +69,8 @@ public void testCagraToHnswSerializationAndSearch() IllegalAccessException, IllegalArgumentException, InvocationTargetException { - Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + Codec codec = + new Lucene101AcceleratedHNSWCodec(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, 3, 16, 100); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); final int COMMIT_FREQ = 2000; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java index 746437d6..49422455 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -8,6 +8,7 @@ import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import java.io.File; import java.io.IOException; import java.lang.reflect.InvocationTargetException; @@ -72,7 +73,8 @@ public void testCagraToHnswSerializationAndSearchWithFallbackWriter() IllegalAccessException, IllegalArgumentException, InvocationTargetException { - Codec codec = new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100); + Codec codec = + new Lucene101AcceleratedHNSWCodec(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, 3, 16, 100); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); final int COMMIT_FREQ = 2000; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index 9baea0e6..bebc5ead 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -6,6 +6,7 @@ import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.ArrayList; @@ -749,6 +750,7 @@ public void testMergeBruteForceIndex() throws IOException { 32, // writer threads 128, // intermediate graph degree 64, // graph degree + CagraGraphBuildAlgo.NN_DESCENT, IndexType.BRUTE_FORCE); // Use brute force index IndexWriterConfig config = @@ -905,6 +907,7 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { 32, // writer threads 128, // intermediate graph degree 64, // graph degree + CagraGraphBuildAlgo.NN_DESCENT, IndexType.CAGRA_AND_BRUTE_FORCE); // Use combined CAGRA + brute force IndexWriterConfig config = From 15728c33e1b91cf48f45f268125c765559546cb1 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 31 Dec 2025 17:10:36 -0500 Subject: [PATCH 2/4] Changes for #79 - Refactor GPUPerLeafCuVSKnnCollector, plug in iTopK and search width values --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 20 +++-- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 4 +- .../lucene/GPUPerLeafCuVSKnnCollector.java | 80 +++---------------- .../cuvs/lucene/TestCuVSDeletedDocuments.java | 15 ++-- .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 7 +- .../TestCuVSRandomizedVectorSearch.java | 10 ++- 6 files changed, 47 insertions(+), 89 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 56a9880d..9f912e76 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -173,7 +173,7 @@ private void validateFieldEntry(FieldInfo info, FieldEntry fieldEntry) { /** * Reads the fieldInfo for each index field and loads FieldEntry in a map. * - * @param meta intance of the ChecksumIndexInput + * @param meta instance of the ChecksumIndexInput * @throws IOException */ private void readFields(ChecksumIndexInput meta) throws IOException { @@ -428,11 +428,19 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits Map result; if (knnCollector.k() <= 1024 && cuvsIndex.getCagraIndex() != null) { - CagraSearchParams searchParams = - new CagraSearchParams.Builder() - .withItopkSize(topK) // TODO: params - .withSearchWidth(1) - .build(); + + CagraSearchParams searchParams; + if (knnCollector instanceof GPUPerLeafCuVSKnnCollector) { + GPUPerLeafCuVSKnnCollector collector = (GPUPerLeafCuVSKnnCollector) knnCollector; + searchParams = + new CagraSearchParams.Builder() + .withItopkSize(Math.max(collector.getiTopK(), topK)) + .withSearchWidth(collector.getSearchWidth()) + .build(); + } else { + // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK + searchParams = new CagraSearchParams.Builder().withItopkSize(topK).build(); + } var query = new CagraQuery.Builder(resources) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 7dc984b2..3173c764 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -7,7 +7,6 @@ import java.io.IOException; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.search.KnnCollector; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; @@ -49,7 +48,8 @@ protected TopDocs approximateSearch( KnnCollectorManager knnCollectorManager) throws IOException { - KnnCollector results = new GPUPerLeafCuVSKnnCollector(k, iTopK, searchWidth); + GPUPerLeafCuVSKnnCollector results = + new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth); LeafReader reader = context.reader(); reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs); diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 9d18b516..57b76294 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -4,92 +4,36 @@ */ package com.nvidia.cuvs.lucene; -import java.util.ArrayList; -import java.util.List; -import org.apache.lucene.search.KnnCollector; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.search.TotalHits; -import org.apache.lucene.search.knn.KnnSearchStrategy; +import org.apache.lucene.search.TopKnnCollector; /** * KnnCollector for cuVS used for search on the GPU. * * @since 25.10 */ -class GPUPerLeafCuVSKnnCollector implements KnnCollector { +class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { - public List scoreDocs; - public int topK = 0; - public int iTopK = topK; // TODO getter, no setter - public int searchWidth = 1; // TODO getter, no setter - public int results = 0; + private int iTopK; + private int searchWidth; /** * Initializes {@link GPUPerLeafCuVSKnnCollector} * - * @param topK the topk value + * @param topK the topK value * @param iTopK the iTopK value * @param searchWidth the search width */ - public GPUPerLeafCuVSKnnCollector(int topK, int iTopK, int searchWidth) { - super(); - this.topK = topK; - this.iTopK = iTopK; + public GPUPerLeafCuVSKnnCollector(int topK, int visitLimit, int iTopK, int searchWidth) { + super(topK, visitLimit); + this.iTopK = iTopK > topK ? iTopK : topK; this.searchWidth = searchWidth; - scoreDocs = new ArrayList(); } - @Override - public boolean earlyTerminated() { - // TODO: may need implementation - return false; + public int getiTopK() { + return iTopK; } - @Override - public void incVisitedCount(int count) { - // TODO: may need implementation - } - - @Override - public long visitedCount() { - // TODO: may need implementation - return 0; - } - - @Override - public long visitLimit() { - // TODO: may need implementation - return 0; - } - - @Override - public int k() { - return topK; - } - - @Override - @SuppressWarnings("cast") - public boolean collect(int docId, float similarity) { - scoreDocs.add(new ScoreDoc(docId, similarity)); - return true; - } - - @Override - public float minCompetitiveSimilarity() { - // TODO: may need implementation - return 0; - } - - @Override - public TopDocs topDocs() { - return new TopDocs( - new TotalHits(scoreDocs.size(), TotalHits.Relation.EQUAL_TO), - scoreDocs.toArray(new ScoreDoc[scoreDocs.size()])); - } - - @Override - public KnnSearchStrategy getSearchStrategy() { - return KnnSearchStrategy.Patience.DEFAULT; + public int getSearchWidth() { + return searchWidth; } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index 04c807a0..37b472d9 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -26,7 +26,6 @@ import org.apache.lucene.index.Term; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TermQuery; @@ -93,7 +92,8 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Use a random vector for query float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify we got results @@ -169,7 +169,8 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti IndexSearcher searcher = newSearcher(reader); float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify results @@ -183,7 +184,7 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Test filtered search with deletions Query filter = new TermQuery(new Term("category", "A")); - Query filteredQuery = + GPUKnnFloatVectorQuery filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; @@ -233,7 +234,8 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { IndexSearcher searcher = newSearcher(reader); float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); TopDocs results = searcher.search(query, topK); assertEquals( @@ -294,7 +296,8 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio IndexSearcher searcher = newSearcher(reader); float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; Set resultIds = new HashSet<>(); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index a48b8f6c..a7883bf0 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -22,7 +22,6 @@ import org.apache.lucene.index.Term; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TermQuery; @@ -116,7 +115,8 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { float[] queryVector = dataset[0]; int topK = random.nextInt(5, TOP_K_LIMIT); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify we get exactly TOP_K results @@ -153,7 +153,8 @@ public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOExcepti // This should further restrict our results to even numbers 0, 2, 4, 6, 8 Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 - Query filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + GPUKnnFloatVectorQuery filteredQuery = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; // Should only get document 8 (the only one that matches the filter and has a vector) diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index abcd6b67..86102ce5 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -25,7 +25,6 @@ import org.apache.lucene.index.Term; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TermQuery; @@ -123,7 +122,8 @@ public void testVectorSearch() throws IOException { log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); log.log(Level.FINE, "TopK: " + topK); - Query query = new KnnFloatVectorQuery("vector", queries[0], topK); + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery("vector", queries[0], topK, null, topK, 1); int correct[] = new int[topK]; for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); @@ -183,7 +183,8 @@ public void testVectorSearchWithFilter() throws IOException { if (dataset.length < topK) topK = dataset.length; // Find a document that has a vector by doing a search first - Query unfiltered = new KnnFloatVectorQuery("vector", dataset[0], 1); + GPUKnnFloatVectorQuery unfiltered = + new GPUKnnFloatVectorQuery("vector", dataset[0], 1, null, 1, 1); ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; // Skip test if no vectors found at all @@ -197,7 +198,8 @@ public void testVectorSearchWithFilter() throws IOException { Query filter = new TermQuery(new Term("id", targetDocId)); // Test the new constructor with filter - Query filteredQuery = new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + GPUKnnFloatVectorQuery filteredQuery = + new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; From 4aadec54ae692558231cf54fc178467408054cf3 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Sat, 3 Jan 2026 23:13:51 -0500 Subject: [PATCH 3/4] Refactor tests --- ...ndSearch.java => TestAcceleratedHNSW.java} | 127 +- .../TestAcceleratedHNSWDeletedDocuments.java | 178 ++- .../lucene/TestAcceleratedHNSWFallback.java | 158 +++ .../cuvs/lucene/TestAcceleratedHNSWGaps.java | 151 +++ .../TestAcceleratedHNSWRandomizedSearch.java | 162 +++ ... => TestAcceleratedHNSWVectorsFormat.java} | 110 +- .../nvidia/cuvs/lucene/TestBackCompat.java | 2 +- ...ializationAndSearchWithFallbackWriter.java | 176 --- ...stCuVSAcceleratedHNSWDeletedDocuments.java | 334 ----- .../lucene/TestCuVSAcceleratedHNSWGaps.java | 190 --- .../cuvs/lucene/TestCuVSDeletedDocuments.java | 183 ++- .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 152 +-- .../TestCuVSRandomizedHNSWVectorSearch.java | 217 ---- .../TestCuVSRandomizedVectorSearch.java | 204 ++-- .../cuvs/lucene/TestCuVSVectorsFormat.java | 106 +- .../nvidia/cuvs/lucene/TestDataProvider.java | 89 ++ .../com/nvidia/cuvs/lucene/TestMerge.java | 1081 ++++++++--------- .../com/nvidia/cuvs/lucene/TestUtils.java | 83 +- 18 files changed, 1623 insertions(+), 2080 deletions(-) rename src/test/java/com/nvidia/cuvs/lucene/{TestCagraToHnswSerializationAndSearch.java => TestAcceleratedHNSW.java} (61%) create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java rename src/test/java/com/nvidia/cuvs/lucene/{TestLucene99AcceleratedHNSWVectorsFormat.java => TestAcceleratedHNSWVectorsFormat.java} (56%) delete mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java delete mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java delete mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java delete mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java similarity index 61% rename from src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java rename to src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java index ed63c350..819bcea6 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java @@ -4,17 +4,19 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import java.io.File; -import java.io.IOException; -import java.lang.reflect.InvocationTargetException; import java.nio.file.Path; import java.nio.file.Paths; import java.util.Arrays; -import java.util.HashSet; +import java.util.List; import java.util.Random; import java.util.UUID; import java.util.logging.Level; @@ -37,6 +39,7 @@ import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.junit.After; @@ -44,101 +47,89 @@ import org.junit.Test; @SuppressSysoutChecks(bugUrl = "") -public class TestCagraToHnswSerializationAndSearch extends LuceneTestCase { +public class TestAcceleratedHNSW extends LuceneTestCase { - private static Logger log = - Logger.getLogger(TestCagraToHnswSerializationAndSearch.class.getName()); + private static final Logger log = Logger.getLogger(TestAcceleratedHNSW.class.getName()); private static Random random; private static Path indexDirPath; + private static String randomID; + private static Codec codec; + private static TestDataProvider dataProvider; @Before public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - // Fixed seed so that we can validate against the same result. - random = new Random(222); + assumeTrue( + "cuVS not supported so skipping these tests", + Lucene99AcceleratedHNSWVectorsFormat.supported()); + random = new Random(); indexDirPath = Paths.get(UUID.randomUUID().toString()); + randomID = UUID.randomUUID().toString(); + dataProvider = new TestDataProvider(random); + codec = + new Lucene101AcceleratedHNSWCodec(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, 3, 16, 100); } @Test - public void testCagraToHnswSerializationAndSearch() - throws IOException, - ClassNotFoundException, - NoSuchMethodException, - SecurityException, - InstantiationException, - IllegalAccessException, - IllegalArgumentException, - InvocationTargetException { - Codec codec = - new Lucene101AcceleratedHNSWCodec(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, 3, 16, 100); - IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); - - final int COMMIT_FREQ = 2000; - final String ID_FIELD = "id"; - final String VECTOR_FIELD = "vector_field"; - - int numDocs = 2000; - int dimension = 32; - int topK = 5; - int count = COMMIT_FREQ; - float[][] dataset = generateDataset(random, numDocs, dimension); - + public void testAcceleratedHNSW() throws Exception { // Indexing try (Directory indexDirectory = FSDirectory.open(indexDirPath); - IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { - for (int i = 0; i < numDocs; i++) { + RandomIndexWriter indexWriter = createWriter(random, indexDirectory, codec)) { + for (int i = 0; i < dataProvider.getDatasetSize(); i++) { Document document = new Document(); document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + document.add( + new KnnFloatVectorField(VECTOR_FIELD1, dataProvider.getDataset1()[i], EUCLIDEAN)); + document.add( + new KnnFloatVectorField(VECTOR_FIELD2, dataProvider.getDataset2()[i], EUCLIDEAN)); indexWriter.addDocument(document); - count -= 1; - if (count == 0) { - indexWriter.commit(); - count = COMMIT_FREQ; - } } + indexWriter.commit(); } // Searching try (Directory indexDirectory = FSDirectory.open(indexDirPath); DirectoryReader reader = DirectoryReader.open(indexDirectory)) { - log.log(Level.FINE, "Successfully opened index"); + + int datasetSize = dataProvider.getDatasetSize(); + int dimensions = dataProvider.getDimensions(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; int vectorCount = 0; for (LeafReaderContext leafReaderContext : reader.leaves()) { LeafReader leafReader = leafReaderContext.reader(); - FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(knnValues); log.log( Level.FINE, - VECTOR_FIELD + VECTOR_FIELD1 + " field: " + knnValues.size() + " vectors, " + knnValues.dimension() + " dimensions"); vectorCount += knnValues.size(); - assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimensions); } - assertTrue("Dataset size mismatch", vectorCount == numDocs); + assertTrue("Dataset size mismatch", vectorCount == datasetSize); log.log(Level.FINE, "Testing vector search queries..."); IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateDataset(random, 1, dimension)[0]; - log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + log.log(Level.FINER, "Query vector: " + Arrays.toString(queryVector)); - KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); TopDocs results = searcher.search(query, topK); log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); - Integer[] expected = new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}; - HashSet expectedIds = new HashSet(Arrays.asList(expected)); + List> expected = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}); for (int i = 0; i < results.scoreDocs.length; i++) { ScoreDoc scoreDoc = results.scoreDocs[i]; Document doc = searcher.storedFields().document(scoreDoc.doc); - String id = doc.get(ID_FIELD); + int id = Integer.valueOf(doc.get(ID_FIELD)); log.log( Level.FINE, " Rank " @@ -149,8 +140,7 @@ public void testCagraToHnswSerializationAndSearch() + id + "), score=" + scoreDoc.score); - assertTrue( - "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); + assertTrue("Id: " + id + " expected but not found", expected.get(0).contains(id)); } assertTrue("TopK results not returned", results.scoreDocs.length == topK); } @@ -158,27 +148,17 @@ public void testCagraToHnswSerializationAndSearch() @Test public void testSingleVectorIndex() throws Exception { - // Test single vector index support with dummy HNSW graph - // TODO: This test can be removed once https://github.com/rapidsai/cuvs/pull/1256 is merged - // and CAGRA natively supports single vector indexes - Codec codec = new Lucene101AcceleratedHNSWCodec(); - - final String ID_FIELD = "id"; - final String VECTOR_FIELD = "vector_field"; + try (Directory indexDirectory = newDirectory()) { - int dimension = 32; - float[] vector = generateDataset(random, 1, dimension)[0]; + int dimensions = dataProvider.getDimensions(); + float[] queryVector = dataProvider.getQueries(1)[0]; - // Index a single document with a vector - this should now work with dummy HNSW graph - try (Directory indexDirectory = newDirectory()) { IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); try (IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { Document document = new Document(); - document.add(new StringField(ID_FIELD, "0", Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + document.add(new StringField(ID_FIELD, randomID, Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD1, queryVector, EUCLIDEAN)); indexWriter.addDocument(document); - - // This should now succeed by creating a dummy HNSW graph for the single vector indexWriter.commit(); } @@ -186,18 +166,19 @@ public void testSingleVectorIndex() throws Exception { try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { assertEquals(1, reader.numDocs()); LeafReader leafReader = getOnlyLeafReader(reader); - FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(knnValues); assertEquals(1, knnValues.size()); - assertEquals(dimension, knnValues.dimension()); + assertEquals(dimensions, knnValues.dimension()); // Test search functionality IndexSearcher searcher = new IndexSearcher(reader); - KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, vector, 1); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, 1); TopDocs results = searcher.search(query, 1); assertEquals(1, results.totalHits.value()); assertEquals(1, results.scoreDocs.length); - assertEquals(0, results.scoreDocs[0].doc); + Document doc = reader.storedFields().document(results.scoreDocs[0].doc); + assertEquals(randomID, doc.get(ID_FIELD)); } } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index c1814fae..9406ea09 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -4,8 +4,12 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.TestDataProvider.CATEGORY_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.createWriterConfig; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; import java.util.ArrayList; @@ -22,9 +26,7 @@ import org.apache.lucene.document.StringField; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; @@ -32,8 +34,6 @@ import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; @@ -44,58 +44,68 @@ @SuppressSysoutChecks(bugUrl = "") public class TestAcceleratedHNSWDeletedDocuments extends LuceneTestCase { - protected static Logger log = + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWDeletedDocuments.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + private static Codec codec; private static Random random; + private static float deletionProbability; + private static float vectorProbability; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); + assumeTrue( + "cuVS not supported so skipping these tests", + Lucene99AcceleratedHNSWVectorsFormat.supported()); + codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); random = random(); + dataProvider = new TestDataProvider(random); + deletionProbability = random.nextFloat() * 0.4f + 0.1f; + vectorProbability = random.nextFloat() * 0.5f + 0.3f; + log.log( + Level.FINE, + "deletion probability: " + + deletionProbability + + ", vector probability: " + + vectorProbability); } @Test public void testVectorSearchWithDeletedDocuments() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200, 1000); // 200-1200 documents - int dimensions = random.nextInt(64, 256); // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); Set deletedDocs = new HashSet<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); // Create index with all documents having vectors - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); } // Delete documents randomly based on probability for (int i = 0; i < datasetSize; i++) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log(Level.FINE, "Number of documents deleted: " + deletedDocs.size()); + // Search and verify deleted documents are not returned try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); // Use a random vector for query - float[] queryVector = generateRandomVector(dimensions, random); - - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getTopK(); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify we got results @@ -103,8 +113,7 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify no deleted documents in results for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); assertFalse( "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); @@ -113,7 +122,7 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify deleted documents are truly deleted for (int deletedId : deletedDocs) { TopDocs result = - searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + searcher.search(new TermQuery(new Term(ID_FIELD, String.valueOf(deletedId))), 1); assertEquals( "Deleted document " + deletedId + " should not be found", 0, @@ -127,29 +136,25 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200) + 50; // 50-250 documents - int dimensions = random.nextInt(256) + 64; // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors - float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); Set docsWithoutVectors = new HashSet<>(); Set deletedDocs = new HashSet<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create index with mixed documents - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); // Randomly assign categories String category = random.nextBoolean() ? "A" : "B"; - doc.add(new StringField("category", category, Field.Store.YES)); + doc.add(new StringField(CATEGORY_FIELD, category, Field.Store.YES)); // Randomly decide whether to add vectors if (random.nextFloat() < vectorProbability) { - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); } else { docsWithoutVectors.add(i); } @@ -159,40 +164,45 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Delete documents randomly for (int i = 0; i < datasetSize; i++) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log( + Level.FINE, + " Documents deleted: " + + deletedDocs.size() + + " Documents without vectors: " + + docsWithoutVectors.size()); + // Test vector search behavior try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify results for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); assertFalse("Deleted document should not appear", deletedDocs.contains(id)); assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); } // Test filtered search with deletions - Query filter = new TermQuery(new Term("category", "A")); - Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); + Query filter = new TermQuery(new Term(CATEGORY_FIELD, "A")); + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; for (ScoreDoc hit : filteredHits) { Document doc = reader.storedFields().document(hit.doc); - String category = doc.get("category"); + String category = doc.get(CATEGORY_FIELD); assertEquals("Should only match category A", "A", category); - int id = Integer.parseInt(doc.get("id")); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse( "Deleted document should not appear in filtered results", deletedDocs.contains(id)); } @@ -204,37 +214,35 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test - int dimensions = random.nextInt(128) + 32; // 32-160 dimensions - int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results - - float[][] dataset = generateDataset(random, datasetSize, dimensions); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create and delete all documents - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { + // Add all documents for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); } writer.commit(); // Delete all documents for (int i = 0; i < datasetSize; i++) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); } writer.commit(); - writer.forceMerge(1); // Force merge to apply deletions + writer.forceMerge(1); } // Verify search returns no results try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); TopDocs results = searcher.search(query, topK); assertEquals( @@ -249,22 +257,20 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200) + 50; // 50-250 documents - int dimensions = random.nextInt(256) + 64; // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - float[][] dataset = generateDataset(random, datasetSize, dimensions); List activeDocIds = new ArrayList<>(); // Initial indexing - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { - int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); for (int i = 0; i < initialDocs; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -273,7 +279,7 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio List candidatesForDeletion = new ArrayList<>(activeDocIds); for (int docId : candidatesForDeletion) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(docId))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(docId))); activeDocIds.remove(Integer.valueOf(docId)); } } @@ -281,9 +287,8 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // Add new documents with higher IDs for (int i = initialDocs; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -293,15 +298,13 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // Verify search behavior after deletions and additions try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; Set resultIds = new HashSet<>(); for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); resultIds.add(id); assertTrue("Result should be from active documents", activeDocIds.contains(id)); } @@ -316,19 +319,4 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio } } } - - private RandomIndexWriter createWriter(Directory directory) throws IOException { - return new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - } - - private IndexWriterConfig createWriterConfig() { - return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy()); - } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java new file mode 100644 index 00000000..71d4c4e7 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java @@ -0,0 +1,158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWFallback extends LuceneTestCase { + + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWFallback.class.getName()); + private static Random random; + private static Path indexDirPath; + private static Codec codec; + private static TestDataProvider dataProvider; + + @Before + public void beforeTest() throws Exception { + assumeTrue( + "cuVS not supported so skipping these tests", + Lucene99AcceleratedHNSWVectorsFormat.supported()); + // Set resources to null to simulate that cuVS is not supported. + Lucene99AcceleratedHNSWVectorsFormat.setResources(null); + + random = new Random(); + dataProvider = new TestDataProvider(random); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + codec = + new Lucene101AcceleratedHNSWCodec(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, 3, 16, 100); + } + + @Test + public void testAcceleratedHNSWFallback() throws Exception { + + int datasetSize = dataProvider.getDatasetSize(); + int dimensions = dataProvider.getDimensions(); + float[][] dataset = dataProvider.getDataset1(); + float[][] dataset2 = dataProvider.getDataset2(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + // Indexing + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + RandomIndexWriter indexWriter = createWriter(random, indexDirectory, codec)) { + for (int i = 0; i < datasetSize; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + document.add(new KnnFloatVectorField(VECTOR_FIELD2, dataset2[i], EUCLIDEAN)); + indexWriter.addDocument(document); + } + indexWriter.commit(); + } + + // Searching + try (Directory indexDirectory = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(indexDirectory)) { + int vectorCount = 0; + for (LeafReaderContext leafReaderContext : reader.leaves()) { + LeafReader leafReader = leafReaderContext.reader(); + FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); + assertNotNull(knnValues); + log.log( + Level.FINE, + VECTOR_FIELD1 + + " field: " + + knnValues.size() + + " vectors, " + + knnValues.dimension() + + " dimensions"); + vectorCount += knnValues.size(); + assertTrue("Vector dimension mismatch", knnValues.dimension() == dimensions); + } + assertTrue("Dataset size mismatch", vectorCount == datasetSize); + + log.log(Level.FINE, "Testing vector search queries..."); + IndexSearcher searcher = new IndexSearcher(reader); + + log.log(Level.FINER, "Query vector: " + Arrays.toString(queryVector)); + + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs results = searcher.search(query, topK); + + log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); + List> expected = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}); + + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + int id = Integer.valueOf(doc.get(ID_FIELD)); + log.log( + Level.FINE, + " Rank " + + (i + 1) + + ": doc " + + scoreDoc.doc + + " (id=" + + id + + "), score=" + + scoreDoc.score); + assertTrue("Id: " + id + " expected but not found", expected.get(0).contains(id)); + } + assertTrue("TopK results not returned", results.scoreDocs.length == topK); + } + } + + @After + public void afterTest() throws Exception { + // Reset resources for other tests to work + Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull()); + + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java new file mode 100644 index 00000000..87cc5d8a --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java @@ -0,0 +1,151 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.IOException; +import java.util.List; +import java.util.Random; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.StoredFields; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWGaps extends LuceneTestCase { + + private static final Logger log = Logger.getLogger(TestAcceleratedHNSWGaps.class.getName()); + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue( + "cuVS not supported so skipping these tests", + Lucene99AcceleratedHNSWVectorsFormat.supported()); + codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + directory = newDirectory(); + random = random(); + dataProvider = new TestDataProvider(random); + RandomIndexWriter writer = createWriter(random, directory, codec); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + + // Create documents where only even-numbered documents have vectors + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); + + // Only add vectors to even-numbered documents + if (i % 2 == 0) { + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + } + + writer.addDocument(doc); + } + + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @Test + public void testVectorSearchWithAlternatingDocuments() throws IOException { + + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + + // Perform search + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + + // Verify we get exactly TOP_K results + assertEquals("Should return exactly " + topK + " results", topK, hits.length); + + // Verify all returned documents have vectors (even-numbered IDs) + StoredFields storedFields = reader.storedFields(); + for (ScoreDoc hit : hits) { + int id = Integer.parseInt(storedFields.document(hit.doc).get(ID_FIELD)); + assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); + log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); + } + + // Verify the results match expected top-k + List expectedIds = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}).get(0); + for (ScoreDoc hit : hits) { + int id = Integer.parseInt(storedFields.document(hit.doc).get(ID_FIELD)); + assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); + } + + log.log(Level.FINE, "Alternating documents test passed with " + hits.length + " results"); + } + + @Test + public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { + + int datasetSize = dataProvider.getDatasetSize(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + String randomEvenInRange = String.valueOf(random.nextInt(datasetSize / 2 + 1) * 2); + log.log(Level.FINE, "Randomly chosen even value is: " + randomEvenInRange); + Query filter = new TermQuery(new Term(ID_FIELD, randomEvenInRange)); + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Should only get document (the only one that matches the filter and has a vector) + assertEquals("Should return exactly 1 result", 1, filteredHits.length); + + String docId = reader.storedFields().document(filteredHits[0].doc).get(ID_FIELD); + assertEquals("Should only return document " + randomEvenInRange, randomEvenInRange, docId); + + log.log( + Level.FINE, + "Filtered alternating documents test passed with " + filteredHits.length + " results"); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java new file mode 100644 index 00000000..a8725f95 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java @@ -0,0 +1,162 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.util.English; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWRandomizedSearch extends LuceneTestCase { + + private static final Logger log = + Logger.getLogger(TestAcceleratedHNSWRandomizedSearch.class.getName()); + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; + + @BeforeClass + public static void beforeClass() throws Exception { + assumeTrue( + "cuVS not supported so skipping these tests", + Lucene99AcceleratedHNSWVectorsFormat.supported()); + directory = newDirectory(); + random = random(); + dataProvider = new TestDataProvider(random); + codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + RandomIndexWriter writer = createWriter(random, directory, codec); + + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + float[][] dataset2 = dataProvider.getDataset2(); + + // Add documents + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); + boolean skipVector = random.nextInt(10) < 4; + if (!skipVector || datasetSize < 100) { + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, dataset2[i], EUCLIDEAN)); + } + writer.addDocument(doc); + } + writer.commit(); + reader = writer.getReader(); + searcher = newSearcher(reader); + writer.close(); + } + + @Test + public void testVectorSearch() throws IOException { + + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + int numQueries = dataProvider.getNumQueries(); + float[][] queries = dataProvider.getQueries(numQueries); + + // Generate queries and expected results for each + List> expected = generateExpectedTopK(topK, dataset, queries); + + for (int i = 0; i < numQueries; i++) { + log.log(Level.FINE, "Running query: " + (i + 1) + " of " + numQueries); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queries[i], topK); + + // Perform search + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); + log.log(Level.FINE, "EXPECTED: " + expected.get(i)); + + // Iterate through the results and assert + for (ScoreDoc hit : hits) { + Document doc = reader.storedFields().document(hit.doc); + int docId = Integer.parseInt(doc.get(ID_FIELD)); + log.log(Level.FINE, "\t" + doc.get(ID_FIELD) + ": " + hit.score); + assertTrue("Result returned was not in topk*2: " + doc, expected.get(i).contains(docId)); + } + } + } + + @Test + public void testVectorSearchWithFilter() throws IOException { + // Find a document that has a vector by doing a search first + + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + Query unfiltered = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, 1); + ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; + + assertTrue( + "Need at least one document with vector for filtering test", unfilteredHits.length > 0); + + Document doc = reader.storedFields().document(unfilteredHits[0].doc); + String targetDocId = doc.get(ID_FIELD); + + // Create a filter that matches only the document we know has a vector + Query filter = new TermQuery(new Term(ID_FIELD, targetDocId)); + + // Test the new constructor with filter + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); + + ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; + + // Ensure we got some results + assertTrue("Should have at least one result", filteredHits.length > 0); + + // Verify that all results match the filter + for (ScoreDoc hit : filteredHits) { + String docId = reader.storedFields().document(hit.doc).get(ID_FIELD); + assertEquals("All results should match the filter", targetDocId, docId); + } + + log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); + } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java similarity index 56% rename from src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java rename to src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java index f834d564..18fe43ed 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java @@ -4,9 +4,13 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.util.List; +import java.util.Random; import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; @@ -26,11 +30,18 @@ import org.junit.Ignore; @SuppressSysoutChecks(bugUrl = "") -public class TestLucene99AcceleratedHNSWVectorsFormat extends BaseKnnVectorsFormatTestCase { +public class TestAcceleratedHNSWVectorsFormat extends BaseKnnVectorsFormatTestCase { + + private static TestDataProvider dataProvider; + private static Random random; @BeforeClass public static void beforeClass() { - assumeTrue("cuVS is not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); + assumeTrue( + "cuVS not supported so skipping these tests", + Lucene99AcceleratedHNSWVectorsFormat.supported()); + random = random(); + dataProvider = new TestDataProvider(random); } @Override @@ -39,27 +50,27 @@ protected Codec getCodec() { } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { - float[][] f = new float[][] {randomVector(384), randomVector(384)}; + final int numDocs = 2; + float[][] vectors = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); - w.addDocument(doc1); - w.commit(); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); - w.addDocument(doc2); - w.flush(); - w.commit(); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors[i], EUCLIDEAN)); + w.addDocument(doc); + w.commit(); + } // sanity - verify one doc per leaf try (DirectoryReader reader = DirectoryReader.open(w)) { List subReaders = reader.leaves(); - assertEquals(2, subReaders.size()); - assertEquals(1, subReaders.get(0).reader().getFloatVectorValues("f").size()); - assertEquals(1, subReaders.get(1).reader().getFloatVectorValues("f").size()); + assertEquals(numDocs, subReaders.size()); + for (int i = 0; i < numDocs; i++) { + assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(VECTOR_FIELD1).size()); + } } // now merge to a single segment @@ -68,49 +79,56 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertEquals(numDocs, values.size()); + for (int i = 0; i < numDocs; i++) { + assertArrayEquals(vectors[i], values.vectorValue(i), 0.0f); + } } } } // Basic test for multiple vectors fields per document public void testTwoVectorFieldsPerDoc() throws Exception { - float[][] f1 = new float[][] {randomVector(384), randomVector(384)}; - float[][] f2 = new float[][] {randomVector(384), randomVector(384)}; + + final int numDocs = 2; + float[][] vectors1 = dataProvider.getVectors(numDocs); + float[][] vectors2 = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); - doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); - w.addDocument(doc1); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); - doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); - w.addDocument(doc2); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors1[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, vectors2[i], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f1"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); - - values = r.getFloatVectorValues("f2"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); - 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); + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors1[i], values.vectorValue(i), 0.0f); + } + + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD2); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors2[i], values.vectorValue(i), 0.0f); + } + + // Check boundary condition - search with a 0 topK + float[] target = dataProvider.getVectors(1)[0]; + var topDocs = r.searchNearestVectors(VECTOR_FIELD1, target, 0, null, 10); assertEquals(0, topDocs.scoreDocs.length); assertEquals(0, topDocs.totalHits.value()); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e660..02e36962 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -13,7 +13,7 @@ import org.junit.Test; /** - * Tests the backward compatibility mechanism. + * Tests the backward compatibility mechanism for {@link LuceneProvider} * * @since 25.12 */ diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java deleted file mode 100644 index 49422455..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ /dev/null @@ -1,176 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import java.io.File; -import java.io.IOException; -import java.lang.reflect.InvocationTargetException; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Arrays; -import java.util.HashSet; -import java.util.Random; -import java.util.UUID; -import java.util.logging.Level; -import java.util.logging.Logger; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.junit.Test; - -@SuppressSysoutChecks(bugUrl = "") -public class TestCagraToHnswSerializationAndSearchWithFallbackWriter extends LuceneTestCase { - - private static Logger log = - Logger.getLogger(TestCagraToHnswSerializationAndSearchWithFallbackWriter.class.getName()); - - private static Random random; - private static Path indexDirPath; - - @BeforeClass - public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - // Set resources to null to simulate that cuVS is not supported. - Lucene99AcceleratedHNSWVectorsFormat.setResources(null); - // Fixed seed so that we can validate against the same result. - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @Test - public void testCagraToHnswSerializationAndSearchWithFallbackWriter() - throws IOException, - ClassNotFoundException, - NoSuchMethodException, - SecurityException, - InstantiationException, - IllegalAccessException, - IllegalArgumentException, - InvocationTargetException { - Codec codec = - new Lucene101AcceleratedHNSWCodec(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, 3, 16, 100); - IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); - - final int COMMIT_FREQ = 2000; - final String ID_FIELD = "id"; - final String VECTOR_FIELD = "vector_field"; - - int numDocs = 2000; - int dimension = 32; - int topK = 5; - int count = COMMIT_FREQ; - float[][] dataset = generateDataset(random, numDocs, dimension); - - // Indexing - try (Directory indexDirectory = FSDirectory.open(indexDirPath); - IndexWriter indexWriter = new IndexWriter(indexDirectory, config)) { - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - indexWriter.addDocument(document); - count -= 1; - if (count == 0) { - indexWriter.commit(); - count = COMMIT_FREQ; - } - } - } - - // Searching - try (Directory indexDirectory = FSDirectory.open(indexDirPath)) { - try (DirectoryReader reader = DirectoryReader.open(indexDirectory)) { - log.log(Level.FINE, "Successfully opened index"); - - int vectorCount = 0; - for (LeafReaderContext leafReaderContext : reader.leaves()) { - LeafReader leafReader = leafReaderContext.reader(); - FloatVectorValues knnValues = leafReader.getFloatVectorValues(VECTOR_FIELD); - assertNotNull(knnValues); - log.log( - Level.FINE, - VECTOR_FIELD - + " field: " - + knnValues.size() - + " vectors, " - + knnValues.dimension() - + " dimensions"); - vectorCount += knnValues.size(); - assertTrue("Vector dimension mismatch", knnValues.dimension() == dimension); - } - assertTrue("Dataset size mismatch", vectorCount == numDocs); - - log.log(Level.FINE, "Testing vector search queries..."); - IndexSearcher searcher = new IndexSearcher(reader); - - float[] queryVector = generateDataset(random, 1, dimension)[0]; - log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); - - KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK); - TopDocs results = searcher.search(query, topK); - - log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); - Integer[] expected = new Integer[] {1869, 1411, 1497, 351, 554}; - HashSet expectedIds = new HashSet(Arrays.asList(expected)); - - for (int i = 0; i < results.scoreDocs.length; i++) { - ScoreDoc scoreDoc = results.scoreDocs[i]; - Document doc = searcher.storedFields().document(scoreDoc.doc); - String id = doc.get(ID_FIELD); - log.log( - Level.FINE, - " Rank " - + (i + 1) - + ": doc " - + scoreDoc.doc - + " (id=" - + id - + "), score=" - + scoreDoc.score); - assertTrue( - "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); - } - assertTrue("TopK results not returned", results.scoreDocs.length == topK); - - } catch (Exception e) { - e.printStackTrace(); - } - } - } - - @AfterClass - public static void afterClass() throws Exception { - // Reset resources for other tests to work - Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull()); - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } -} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java deleted file mode 100644 index ba53c8b0..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java +++ /dev/null @@ -1,334 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Random; -import java.util.Set; -import java.util.logging.Level; -import java.util.logging.Logger; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TermQuery; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; -import org.apache.lucene.tests.index.RandomIndexWriter; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.tests.util.TestUtil; -import org.junit.BeforeClass; -import org.junit.Test; - -@SuppressSysoutChecks(bugUrl = "") -public class TestCuVSAcceleratedHNSWDeletedDocuments extends LuceneTestCase { - - protected static Logger log = - Logger.getLogger(TestCuVSAcceleratedHNSWDeletedDocuments.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - private static Random random; - - @BeforeClass - public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - random = random(); - } - - @Test - public void testVectorSearchWithDeletedDocuments() throws IOException { - - try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200, 1000); // 200-1200 documents - int dimensions = random.nextInt(64, 256); // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); - Set deletedDocs = new HashSet<>(); - - // Create index with all documents having vectors - try (RandomIndexWriter writer = createWriter(directory)) { - for (int i = 0; i < datasetSize; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - writer.addDocument(doc); - } - - // Delete documents randomly based on probability - for (int i = 0; i < datasetSize; i++) { - if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); - deletedDocs.add(i); - } - } - writer.commit(); - } - - // Search and verify deleted documents are not returned - try (DirectoryReader reader = DirectoryReader.open(directory)) { - IndexSearcher searcher = newSearcher(reader); - // Use a random vector for query - float[] queryVector = generateRandomVector(dimensions, random); - - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - - // Verify we got results - assertTrue("Should have search results", hits.length > 0); - - // Verify no deleted documents in results - for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); - assertFalse( - "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); - log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); - } - - // Verify deleted documents are truly deleted - for (int deletedId : deletedDocs) { - TopDocs result = - searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); - assertEquals( - "Deleted document " + deletedId + " should not be found", - 0, - result.totalHits.value()); - } - } - } - } - - @Test - public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { - - try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200) + 50; // 50-250 documents - int dimensions = random.nextInt(256) + 64; // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors - float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); - Set docsWithoutVectors = new HashSet<>(); - Set deletedDocs = new HashSet<>(); - - // Create index with mixed documents - try (RandomIndexWriter writer = createWriter(directory)) { - for (int i = 0; i < datasetSize; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - // Randomly assign categories - String category = random.nextBoolean() ? "A" : "B"; - doc.add(new StringField("category", category, Field.Store.YES)); - - // Randomly decide whether to add vectors - if (random.nextFloat() < vectorProbability) { - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - } else { - docsWithoutVectors.add(i); - } - writer.addDocument(doc); - } - - // Delete documents randomly - for (int i = 0; i < datasetSize; i++) { - if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); - deletedDocs.add(i); - } - } - writer.commit(); - } - - // Test vector search behavior - try (DirectoryReader reader = DirectoryReader.open(directory)) { - IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - - // Verify results - for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); - assertFalse("Deleted document should not appear", deletedDocs.contains(id)); - assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); - log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); - } - - // Test filtered search with deletions - Query filter = new TermQuery(new Term("category", "A")); - Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); - ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; - - for (ScoreDoc hit : filteredHits) { - Document doc = reader.storedFields().document(hit.doc); - String category = doc.get("category"); - assertEquals("Should only match category A", "A", category); - int id = Integer.parseInt(doc.get("id")); - assertFalse( - "Deleted document should not appear in filtered results", deletedDocs.contains(id)); - } - } - } - } - - @Test - public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { - - try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test - int dimensions = random.nextInt(128) + 32; // 32-160 dimensions - int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results - - float[][] dataset = generateDataset(random, datasetSize, dimensions); - - // Create and delete all documents - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { - for (int i = 0; i < datasetSize; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - writer.addDocument(doc); - } - writer.commit(); - - // Delete all documents - for (int i = 0; i < datasetSize; i++) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); - } - writer.commit(); - writer.forceMerge(1); // Force merge to apply deletions - } - - // Verify search returns no results - try (DirectoryReader reader = DirectoryReader.open(directory)) { - IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); - TopDocs results = searcher.search(query, topK); - - assertEquals( - "Should return no results when all documents are deleted", - 0, - results.totalHits.value()); - } - } - } - - @Test - public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { - - try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200) + 50; // 50-250 documents - int dimensions = random.nextInt(256) + 64; // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); - List activeDocIds = new ArrayList<>(); - - // Initial indexing - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { - int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset - for (int i = 0; i < initialDocs; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - writer.addDocument(doc); - activeDocIds.add(i); - } - - // Delete some documents randomly - List candidatesForDeletion = new ArrayList<>(activeDocIds); - for (int docId : candidatesForDeletion) { - if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(docId))); - activeDocIds.remove(Integer.valueOf(docId)); - } - } - - // Add new documents with higher IDs - for (int i = initialDocs; i < datasetSize; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - writer.addDocument(doc); - activeDocIds.add(i); - } - writer.commit(); - } - - // Verify search behavior after deletions and additions - try (DirectoryReader reader = DirectoryReader.open(directory)) { - IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - - Set resultIds = new HashSet<>(); - for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); - resultIds.add(id); - assertTrue("Result should be from active documents", activeDocIds.contains(id)); - } - - log.log( - Level.FINE, - "Search returned " - + hits.length - + " results from " - + activeDocIds.size() - + " active documents"); - } - } - } - - private RandomIndexWriter createWriter(Directory directory) throws IOException { - return new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - } - - private IndexWriterConfig createWriterConfig() { - return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy()); - } -} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java deleted file mode 100644 index 6cac6804..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java +++ /dev/null @@ -1,190 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; - -import java.io.IOException; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.TreeMap; -import java.util.logging.Level; -import java.util.logging.Logger; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TermQuery; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; -import org.apache.lucene.tests.index.RandomIndexWriter; -import org.apache.lucene.tests.util.English; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.tests.util.TestUtil; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.junit.Test; - -@SuppressSysoutChecks(bugUrl = "") -public class TestCuVSAcceleratedHNSWGaps extends LuceneTestCase { - - protected static Logger log = Logger.getLogger(TestCuVSAcceleratedHNSWGaps.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - static IndexSearcher searcher; - static IndexReader reader; - static Directory directory; - static Random random; - - static int DATASET_SIZE_LIMIT = 1000; - static int DIMENSIONS_LIMIT = 2048; - static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; - - static int datasetSize; - static int dimension; - static float[][] dataset; - - @BeforeClass - public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - directory = newDirectory(); - random = random(); - - RandomIndexWriter writer = - new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - - log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - datasetSize = random.nextInt(100, DATASET_SIZE_LIMIT); - dimension = random.nextInt(8, DIMENSIONS_LIMIT); - dataset = generateDataset(random, datasetSize, dimension); - - // Create documents where only even-numbered documents have vectors - for (int i = 0; i < datasetSize; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); - - // Only add vectors to even-numbered documents - if (i % 2 == 0) { - doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - } - - writer.addDocument(doc); - } - - reader = writer.getReader(); - searcher = newSearcher(reader); - writer.close(); - } - - @AfterClass - public static void afterClass() throws Exception { - if (reader != null) reader.close(); - if (directory != null) directory.close(); - searcher = null; - reader = null; - directory = null; - log.log(Level.FINE, "Test finished"); - } - - @Test - public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - - // Use the first vector (from document 0) as query - float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); - - Query query = new KnnFloatVectorQuery("vector", queryVector, topK); - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - - // Verify we get exactly TOP_K results - assertEquals("Should return exactly " + topK + " results", topK, hits.length); - - // Verify all returned documents have vectors (even-numbered IDs) - for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); - assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); - log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); - } - - // Verify the results match expected top-k based on Euclidean distance - List expectedIds = calculateExpectedTopK(queryVector, topK, dataset); - for (int i = 0; i < hits.length; i++) { - String docId = reader.storedFields().document(hits[i].doc).get("id"); - int id = Integer.parseInt(docId); - assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); - } - - log.log(Level.FINE, "Alternating document test passed with " + hits.length + " results"); - } - - @Test - public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - - // Use the first vector (from document 0) as query - float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); - - // Create a filter that only matches documents with ID less than 10 - // This should further restrict our results to even numbers 0, 2, 4, 6, 8 - Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 - - Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); - ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; - - // Should only get document 8 (the only one that matches the filter and has a vector) - assertEquals("Should return exactly 1 result", 1, filteredHits.length); - - String docId = reader.storedFields().document(filteredHits[0].doc).get("id"); - assertEquals("Should only return document 8", "8", docId); - - log.log( - Level.FINE, - "Filtered alternating document test passed with " + filteredHits.length + " results"); - } - - public static List calculateExpectedTopK(float[] query, int topK, float[][] dataset) { - Map distances = new TreeMap<>(); - - // Calculate distances only for documents that have vectors (even-numbered) - for (int i = 0; i < dataset.length; i += 2) { - double distance = 0; - for (int j = 0; j < dataset[0].length; j++) { - distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); - } - distances.put(i, distance); - } - - // Sort by distance and return top-k - return distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .limit(topK) - .toList(); - } -} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index 37b472d9..baf43a4d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -4,8 +4,12 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; +import static com.nvidia.cuvs.lucene.TestDataProvider.CATEGORY_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.createWriterConfig; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; import java.util.ArrayList; @@ -22,18 +26,15 @@ import org.apache.lucene.document.StringField; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.tests.store.BaseDirectoryWrapper; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.tests.util.TestUtil; @@ -43,57 +44,69 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSDeletedDocuments extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); - - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + private static final Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName()); + private static Codec codec; private static Random random; + private static TestDataProvider dataProvider; + private static float deletionProbability; + private static float vectorProbability; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported so skipping these tests", CuVS2510GPUVectorsFormat.supported()); random = random(); + dataProvider = new TestDataProvider(random); + codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + deletionProbability = random.nextFloat() * 0.4f + 0.1f; + vectorProbability = random.nextFloat() * 0.5f + 0.3f; + + log.log( + Level.FINE, + ", deletion probability: " + + deletionProbability + + ", vector probability: " + + vectorProbability); } @Test public void testVectorSearchWithDeletedDocuments() throws IOException { - try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200, 1000); // 200-1200 documents - int dimensions = random.nextInt(64, 256); // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float deletionProbability = random.nextFloat() * 0.4f + 0.1f; // 10-50% deletion rate + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - float[][] dataset = generateDataset(random, datasetSize, dimensions); + try (BaseDirectoryWrapper directory = newDirectory()) { Set deletedDocs = new HashSet<>(); // Create index with all documents having vectors - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { + // Add documents for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); } + writer.commit(); // Delete documents randomly based on probability for (int i = 0; i < datasetSize; i++) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log(Level.FINE, "Number of documents deleted: " + deletedDocs.size()); + // Search and verify deleted documents are not returned try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - // Use a random vector for query - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify we got results @@ -101,8 +114,8 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify no deleted documents in results for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse( "Deleted document " + id + " should not appear in results", deletedDocs.contains(id)); log.log(Level.FINE, "Found non-deleted document: " + id + ", Score: " + hit.score); @@ -110,8 +123,8 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { // Verify deleted documents are truly deleted for (int deletedId : deletedDocs) { - TopDocs result = - searcher.search(new TermQuery(new Term("id", String.valueOf(deletedId))), 1); + Query termQuery = new TermQuery(new Term(ID_FIELD, String.valueOf(deletedId))); + TopDocs result = searcher.search(termQuery, 1); assertEquals( "Deleted document " + deletedId + " should not be found", 0, @@ -125,29 +138,25 @@ public void testVectorSearchWithDeletedDocuments() throws IOException { public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200) + 50; // 50-250 documents - int dimensions = random.nextInt(256) + 64; // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float vectorProbability = random.nextFloat() * 0.5f + 0.3f; // 30-80% have vectors - float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); Set docsWithoutVectors = new HashSet<>(); Set deletedDocs = new HashSet<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create index with mixed documents - try (RandomIndexWriter writer = createWriter(directory)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); // Randomly assign categories String category = random.nextBoolean() ? "A" : "B"; - doc.add(new StringField("category", category, Field.Store.YES)); + doc.add(new StringField(CATEGORY_FIELD, category, Field.Store.YES)); // Randomly decide whether to add vectors if (random.nextFloat() < vectorProbability) { - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); } else { docsWithoutVectors.add(i); } @@ -157,42 +166,42 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti // Delete documents randomly for (int i = 0; i < datasetSize; i++) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); deletedDocs.add(i); } } writer.commit(); } + log.log(Level.FINE, "Number of deleted documents: " + deletedDocs.size()); + // Test vector search behavior try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; // Verify results for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse("Deleted document should not appear", deletedDocs.contains(id)); assertFalse("Document without vector should not appear", docsWithoutVectors.contains(id)); log.log(Level.FINE, "Found document with vector: " + id + ", Score: " + hit.score); } // Test filtered search with deletions - Query filter = new TermQuery(new Term("category", "A")); - GPUKnnFloatVectorQuery filteredQuery = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filter = new TermQuery(new Term(CATEGORY_FIELD, "A")); + Query filteredQuery = + new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; for (ScoreDoc hit : filteredHits) { Document doc = reader.storedFields().document(hit.doc); - String category = doc.get("category"); + String category = doc.get(CATEGORY_FIELD); assertEquals("Should only match category A", "A", category); - int id = Integer.parseInt(doc.get("id")); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertFalse( "Deleted document should not appear in filtered results", deletedDocs.contains(id)); } @@ -204,26 +213,24 @@ public void testVectorSearchWithMixedDeletedAndMissingVectors() throws IOExcepti public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(20) + 5; // 5-25 documents for this test - int dimensions = random.nextInt(128) + 32; // 32-160 dimensions - int topK = Math.min(random.nextInt(10) + 5, datasetSize); // 5-15 results - - float[][] dataset = generateDataset(random, datasetSize, dimensions); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Create and delete all documents - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); } writer.commit(); // Delete all documents for (int i = 0; i < datasetSize; i++) { - writer.deleteDocuments(new Term("id", String.valueOf(i))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(i))); } writer.commit(); writer.forceMerge(1); // Force merge to apply deletions @@ -232,10 +239,8 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { // Verify search returns no results try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); TopDocs results = searcher.search(query, topK); assertEquals( @@ -250,22 +255,19 @@ public void testVectorSearchAfterAllDocumentsDeleted() throws IOException { public void testVectorSearchWithPartialDeletionAndReindexing() throws IOException { try (Directory directory = newDirectory()) { - int datasetSize = random.nextInt(200) + 50; // 50-250 documents - int dimensions = random.nextInt(256) + 64; // 64-320 dimensions - int topK = Math.min(random.nextInt(20) + 5, datasetSize / 2); // 5-25 results - float deletionProbability = random.nextFloat() * 0.3f + 0.1f; // 10-40% deletion rate - - float[][] dataset = generateDataset(random, datasetSize, dimensions); List activeDocIds = new ArrayList<>(); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; // Initial indexing - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { - int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); // 50-75% of dataset + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig(random, codec))) { + int initialDocs = datasetSize / 2 + random.nextInt(datasetSize / 4); for (int i = 0; i < initialDocs; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -274,7 +276,7 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio List candidatesForDeletion = new ArrayList<>(activeDocIds); for (int docId : candidatesForDeletion) { if (random.nextFloat() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(docId))); + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(docId))); activeDocIds.remove(Integer.valueOf(docId)); } } @@ -282,9 +284,8 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // Add new documents with higher IDs for (int i = initialDocs; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add( - new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); writer.addDocument(doc); activeDocIds.add(i); } @@ -294,16 +295,13 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio // Verify search behavior after deletions and additions try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = newSearcher(reader); - float[] queryVector = generateRandomVector(dimensions, random); - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; Set resultIds = new HashSet<>(); for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + int id = Integer.parseInt(reader.storedFields().document(hit.doc).get(ID_FIELD)); resultIds.add(id); assertTrue("Result should be from active documents", activeDocIds.contains(id)); } @@ -318,19 +316,4 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio } } } - - private RandomIndexWriter createWriter(Directory directory) throws IOException { - return new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - } - - private IndexWriterConfig createWriterConfig() { - return newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy()); - } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index a7883bf0..95861a1c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -4,13 +4,16 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.TreeMap; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -20,14 +23,11 @@ import org.apache.lucene.document.StringField; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TermQuery; import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.English; import org.apache.lucene.tests.util.LuceneTestCase; @@ -40,101 +40,71 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSGaps extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); - - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); - static IndexSearcher searcher; - static IndexReader reader; - static Directory directory; - static Random random; - - static int DATASET_SIZE_LIMIT = 1000; - static int DIMENSIONS_LIMIT = 2048; - static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; - - static int datasetSize; - static int dimension; - static float[][] dataset; + private static final Logger log = Logger.getLogger(TestCuVSGaps.class.getName()); + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported so skipping these tests", CuVS2510GPUVectorsFormat.supported()); directory = newDirectory(); random = random(); + dataProvider = new TestDataProvider(random); - RandomIndexWriter writer = - new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - - log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - datasetSize = random.nextInt(100, DATASET_SIZE_LIMIT); - dimension = random.nextInt(8, DIMENSIONS_LIMIT); - dataset = generateDataset(random, datasetSize, dimension); + codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + RandomIndexWriter writer = createWriter(random, directory, codec); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); // Create documents where only even-numbered documents have vectors for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); // Only add vectors to even-numbered documents if (i % 2 == 0) { - doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); } - writer.addDocument(doc); } - + writer.commit(); reader = writer.getReader(); searcher = newSearcher(reader); writer.close(); } - @AfterClass - public static void afterClass() throws Exception { - if (reader != null) reader.close(); - if (directory != null) directory.close(); - searcher = null; - reader = null; - directory = null; - log.log(Level.FINE, "Test finished"); - } - @Test public void testVectorSearchWithAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); - // Use the first vector (from document 0) as query - float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); + Query query = new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, null, topK, 1); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - // Verify we get exactly TOP_K results + // Verify we get exactly topK results assertEquals("Should return exactly " + topK + " results", topK, hits.length); // Verify all returned documents have vectors (even-numbered IDs) for (ScoreDoc hit : hits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - int id = Integer.parseInt(docId); + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); } - // Verify the results match expected top-k based on Euclidean distance - List expectedIds = calculateExpectedTopK(queryVector, topK, dataset); - for (int i = 0; i < hits.length; i++) { - String docId = reader.storedFields().document(hits[i].doc).get("id"); - int id = Integer.parseInt(docId); + // Verify the results match expected top-k + List expectedIds = + generateExpectedTopK(topK, dataset, new float[][] {queryVector}).get(0); + for (ScoreDoc hit : hits) { + Document doc = reader.storedFields().document(hit.doc); + int id = Integer.parseInt(doc.get(ID_FIELD)); assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); } @@ -143,48 +113,36 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { @Test public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); - // Use the first vector (from document 0) as query - float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + int datasetSize = dataProvider.getDatasetSize(); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; - // Create a filter that only matches documents with ID less than 10 - // This should further restrict our results to even numbers 0, 2, 4, 6, 8 - Query filter = new TermQuery(new Term("id", "8")); // Only match document 8 + String randomEvenInRange = String.valueOf(random.nextInt(datasetSize / 2 + 1) * 2); + log.log(Level.FINE, "Randomly chosen even value is: " + randomEvenInRange); + Query filter = new TermQuery(new Term(ID_FIELD, randomEvenInRange)); - GPUKnnFloatVectorQuery filteredQuery = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = + new GPUKnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter, topK, 1); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; - // Should only get document 8 (the only one that matches the filter and has a vector) + // Should only get document (the only one that matches the filter and has a vector) assertEquals("Should return exactly 1 result", 1, filteredHits.length); - - String docId = reader.storedFields().document(filteredHits[0].doc).get("id"); - assertEquals("Should only return document 8", "8", docId); + Document doc = reader.storedFields().document(filteredHits[0].doc); + assertEquals( + "Should only return document " + randomEvenInRange, randomEvenInRange, doc.get(ID_FIELD)); log.log( Level.FINE, "Filtered alternating document test passed with " + filteredHits.length + " results"); } - public static List calculateExpectedTopK(float[] query, int topK, float[][] dataset) { - Map distances = new TreeMap<>(); - - // Calculate distances only for documents that have vectors (even-numbered) - for (int i = 0; i < dataset.length; i += 2) { - double distance = 0; - for (int j = 0; j < dataset[0].length; j++) { - distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); - } - distances.put(i, distance); - } - - // Sort by distance and return top-k - return distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .limit(topK) - .toList(); + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java deleted file mode 100644 index 385ca8f5..00000000 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java +++ /dev/null @@ -1,217 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.TreeMap; -import java.util.logging.Level; -import java.util.logging.Logger; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TermQuery; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; -import org.apache.lucene.tests.index.RandomIndexWriter; -import org.apache.lucene.tests.util.English; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.tests.util.TestUtil; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.junit.Test; - -@SuppressSysoutChecks(bugUrl = "") -public class TestCuVSRandomizedHNSWVectorSearch extends LuceneTestCase { - - protected static Logger log = - Logger.getLogger(TestCuVSRandomizedHNSWVectorSearch.class.getName()); - - static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - static IndexSearcher searcher; - static IndexReader reader; - static Directory directory; - - static int DATASET_SIZE_LIMIT = 1000; - static int DIMENSIONS_LIMIT = 2048; - static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 - static float[][] dataset; - - @BeforeClass - public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - directory = newDirectory(); - - RandomIndexWriter writer = - new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - - log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - Random random = random(); - int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; - int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; - dataset = generateDataset(random, datasetSize, dimensions); - for (int i = 0; i < datasetSize; i++) { - Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); - boolean skipVector = - random.nextInt(10) - < 4; // some documents won't have vectors to test deleted/missing vectors - if (!skipVector - || datasetSize < 100) { // about 10th of the documents shouldn't have a single vector - doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - doc.add(new KnnFloatVectorField("vector2", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - } - - writer.addDocument(doc); - } - - reader = writer.getReader(); - searcher = newSearcher(reader); - writer.close(); - } - - @AfterClass - public static void afterClass() throws Exception { - if (reader != null) reader.close(); - if (directory != null) directory.close(); - searcher = null; - reader = null; - directory = null; - log.log(Level.FINE, "Test finished"); - } - - @Test - public void testVectorSearch() throws IOException { - Random random = random(); - int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; - int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - - if (dataset.length < topK) topK = dataset.length; - - float[][] queries = generateQueries(random, dataset[0].length, numQueries); - List> expected = generateExpectedResults(topK, dataset, queries); - - log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); - log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); - log.log(Level.FINE, "TopK: " + topK); - - Query query = new KnnFloatVectorQuery("vector", queries[0], topK); - int correct[] = new int[topK]; - for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); - - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); - log.log(Level.FINE, "EXPECTD: " + expected.get(0)); - - for (ScoreDoc hit : hits) { - log.log( - Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); - } - - for (ScoreDoc hit : hits) { - int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); - assertTrue("Result returned was not in topk*2: " + doc, expected.get(0).contains(doc)); - } - } - - private static List> generateExpectedResults( - int topK, float[][] dataset, float[][] queries) { - List> neighborsResult = new ArrayList<>(); - int dimensions = dataset[0].length; - - for (float[] query : queries) { - Map distances = new TreeMap<>(); - for (int j = 0; j < dataset.length; j++) { - double distance = 0; - for (int k = 0; k < dimensions; k++) { - distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); - } - distances.put(j, (distance)); - } - - Map sorted = new TreeMap(distances); - log.log(Level.FINE, "EXPECTED: " + sorted); - - // Sort by distance and select the topK nearest neighbors - List neighbors = - distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .toList(); - neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); - } - - log.log(Level.FINE, "Expected results generated successfully."); - return neighborsResult; - } - - @Test - public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported()); - - Random random = random(); - int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - - if (dataset.length < topK) topK = dataset.length; - - // Find a document that has a vector by doing a search first - Query unfiltered = new KnnFloatVectorQuery("vector", dataset[0], 1); - ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; - - // Skip test if no vectors found at all - assumeTrue( - "Need at least one document with vector for filtering test", unfilteredHits.length > 0); - - String targetDocId = reader.storedFields().document(unfilteredHits[0].doc).get("id"); - float[] queryVector = dataset[0]; - - // Create a filter that matches only the document we know has a vector - Query filter = new TermQuery(new Term("id", targetDocId)); - - // Test the new constructor with filter - Query filteredQuery = new KnnFloatVectorQuery("vector", queryVector, topK, filter); - - ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; - - // Ensure we got some results - assertTrue("Should have at least one result", filteredHits.length > 0); - - // Verify that all results match the filter - for (ScoreDoc hit : filteredHits) { - String docId = reader.storedFields().document(hit.doc).get("id"); - assertEquals("All results should match the filter", targetDocId, docId); - } - - log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); - } -} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 86102ce5..dda6982f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -4,16 +4,18 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.TEXT_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.generateExpectedTopK; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.io.IOException; -import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.TreeMap; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -23,14 +25,12 @@ import org.apache.lucene.document.StringField; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.Term; -import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TermQuery; import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.analysis.MockAnalyzer; -import org.apache.lucene.tests.analysis.MockTokenizer; import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.English; import org.apache.lucene.tests.util.LuceneTestCase; @@ -43,50 +43,36 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSRandomizedVectorSearch extends LuceneTestCase { - protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); + private static final Logger log = + Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName()); - static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); - static IndexSearcher searcher; - static IndexReader reader; - static Directory directory; - - static int DATASET_SIZE_LIMIT = 1000; - static int DIMENSIONS_LIMIT = 2048; - static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 - static float[][] dataset; + private static Codec codec; + private static IndexSearcher searcher; + private static IndexReader reader; + private static Directory directory; + private static Random random; + private static TestDataProvider dataProvider; @BeforeClass public static void beforeClass() throws Exception { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported so skipping these tests", CuVS2510GPUVectorsFormat.supported()); + codec = TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); directory = newDirectory(); + random = random(); + dataProvider = new TestDataProvider(random); + + RandomIndexWriter writer = createWriter(random, directory, codec); + int datasetSize = dataProvider.getDatasetSize(); + float[][] dataset = dataProvider.getDataset1(); - RandomIndexWriter writer = - new RandomIndexWriter( - random(), - directory, - newIndexWriterConfig(new MockAnalyzer(random(), MockTokenizer.SIMPLE, true)) - .setMaxBufferedDocs(TestUtil.nextInt(random(), 100, 1000)) - .setCodec(codec) - .setMergePolicy(newTieredMergePolicy())); - - log.log(Level.FINE, "Merge Policy: " + writer.w.getConfig().getMergePolicy()); - - Random random = random(); - int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; - int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; - dataset = generateDataset(random, datasetSize, dimensions); for (int i = 0; i < datasetSize; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add(newTextField("field", English.intToEnglish(i), Field.Store.YES)); - boolean skipVector = - random.nextInt(10) - < 4; // some documents won't have vectors to test deleted/missing vectors - if (!skipVector - || datasetSize < 100) { // about 10th of the documents shouldn't have a single vector - doc.add(new KnnFloatVectorField("vector", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); - doc.add(new KnnFloatVectorField("vector2", dataset[i], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(newTextField(TEXT_FIELD, English.intToEnglish(i), Field.Store.YES)); + boolean skipVector = random.nextInt(10) < 4; + if (!skipVector || datasetSize < 100) { + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, dataset[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, dataset[i], EUCLIDEAN)); } writer.addDocument(doc); @@ -97,109 +83,55 @@ public static void beforeClass() throws Exception { writer.close(); } - @AfterClass - public static void afterClass() throws Exception { - if (reader != null) reader.close(); - if (directory != null) directory.close(); - searcher = null; - reader = null; - directory = null; - log.log(Level.FINE, "Test finished"); - } - @Test public void testVectorSearch() throws IOException { - Random random = random(); - int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; - int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - - if (dataset.length < topK) topK = dataset.length; - - float[][] queries = generateQueries(random, dataset[0].length, numQueries); - List> expected = generateExpectedResults(topK, dataset, queries); - - log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); - log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); - log.log(Level.FINE, "TopK: " + topK); - - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery("vector", queries[0], topK, null, topK, 1); - int correct[] = new int[topK]; - for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); - - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); - log.log(Level.FINE, "EXPECTD: " + expected.get(0)); - - for (ScoreDoc hit : hits) { - log.log( - Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); - } - - for (ScoreDoc hit : hits) { - int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); - assertTrue("Result returned was not in topk*2: " + doc, expected.get(0).contains(doc)); - } - } - - private static List> generateExpectedResults( - int topK, float[][] dataset, float[][] queries) { - List> neighborsResult = new ArrayList<>(); - int dimensions = dataset[0].length; - - for (float[] query : queries) { - Map distances = new TreeMap<>(); - for (int j = 0; j < dataset.length; j++) { - double distance = 0; - for (int k = 0; k < dimensions; k++) { - distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); - } - distances.put(j, (distance)); + float[][] dataset = dataProvider.getDataset1(); + int topK = dataProvider.getTopK(); + int numQueries = dataProvider.getNumQueries(); + float[][] queries = dataProvider.getQueries(numQueries); + + // Generate queries and expected results for each + List> expected = generateExpectedTopK(topK, dataset, queries); + + for (int i = 0; i < numQueries; i++) { + log.log(Level.FINE, "Running query: " + (i + 1) + " of " + numQueries); + Query query = new KnnFloatVectorQuery(VECTOR_FIELD1, queries[i], topK); + + // Perform search + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); + log.log(Level.FINE, "EXPECTED: " + expected.get(i)); + + // Iterate through the results and assert + for (ScoreDoc hit : hits) { + Document doc = reader.storedFields().document(hit.doc); + int docId = Integer.parseInt(doc.get(ID_FIELD)); + log.log(Level.FINE, "\t" + doc.get(ID_FIELD) + ": " + hit.score); + assertTrue("Result returned was not in topk*2: " + doc, expected.get(i).contains(docId)); } - - Map sorted = new TreeMap(distances); - log.log(Level.FINE, "EXPECTED: " + sorted); - - // Sort by distance and select the topK nearest neighbors - List neighbors = - distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .toList(); - neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); } - - log.log(Level.FINE, "Expected results generated successfully."); - return neighborsResult; } @Test public void testVectorSearchWithFilter() throws IOException { - assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported()); - - Random random = random(); - int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - - if (dataset.length < topK) topK = dataset.length; - // Find a document that has a vector by doing a search first - GPUKnnFloatVectorQuery unfiltered = - new GPUKnnFloatVectorQuery("vector", dataset[0], 1, null, 1, 1); + int topK = dataProvider.getTopK(); + float[] queryVector = dataProvider.getQueries(1)[0]; + + Query unfiltered = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, 1); ScoreDoc[] unfilteredHits = searcher.search(unfiltered, 1).scoreDocs; - // Skip test if no vectors found at all - assumeTrue( + assertTrue( "Need at least one document with vector for filtering test", unfilteredHits.length > 0); - String targetDocId = reader.storedFields().document(unfilteredHits[0].doc).get("id"); - float[] queryVector = dataset[0]; + Document doc = reader.storedFields().document(unfilteredHits[0].doc); + String targetDocId = doc.get(ID_FIELD); // Create a filter that matches only the document we know has a vector - Query filter = new TermQuery(new Term("id", targetDocId)); + Query filter = new TermQuery(new Term(ID_FIELD, targetDocId)); // Test the new constructor with filter - GPUKnnFloatVectorQuery filteredQuery = - new GPUKnnFloatVectorQuery("vector", queryVector, topK, filter, topK, 1); + Query filteredQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK, filter); ScoreDoc[] filteredHits = searcher.search(filteredQuery, topK).scoreDocs; @@ -208,10 +140,20 @@ public void testVectorSearchWithFilter() throws IOException { // Verify that all results match the filter for (ScoreDoc hit : filteredHits) { - String docId = reader.storedFields().document(hit.doc).get("id"); + String docId = reader.storedFields().document(hit.doc).get(ID_FIELD); assertEquals("All results should match the filter", targetDocId, docId); } log.log(Level.FINE, "Prefiltering test passed with " + filteredHits.length + " results"); } + + @AfterClass + public static void afterClass() throws Exception { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + searcher = null; + reader = null; + directory = null; + log.log(Level.FINE, "Test finished"); + } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 1b3de41d..67ba654a 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -4,9 +4,13 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD2; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import java.util.List; +import java.util.Random; import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; @@ -28,9 +32,14 @@ @SuppressSysoutChecks(bugUrl = "") public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase { + private static TestDataProvider dataProvider; + private static Random random; + @BeforeClass public static void beforeClass() { - assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported so skipping these tests", CuVS2510GPUVectorsFormat.supported()); + random = random(); + dataProvider = new TestDataProvider(random); } @Override @@ -39,27 +48,27 @@ protected Codec getCodec() { } public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { - float[][] f = new float[][] {randomVector(384), randomVector(384)}; + final int numDocs = 2; + float[][] vectors = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); - w.addDocument(doc1); - w.commit(); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); - w.addDocument(doc2); - w.flush(); - w.commit(); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors[i], EUCLIDEAN)); + w.addDocument(doc); + w.commit(); + } // sanity - verify one doc per leaf try (DirectoryReader reader = DirectoryReader.open(w)) { List subReaders = reader.leaves(); - assertEquals(2, subReaders.size()); - assertEquals(1, subReaders.get(0).reader().getFloatVectorValues("f").size()); - assertEquals(1, subReaders.get(1).reader().getFloatVectorValues("f").size()); + assertEquals(numDocs, subReaders.size()); + for (int i = 0; i < numDocs; i++) { + assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(VECTOR_FIELD1).size()); + } } // now merge to a single segment @@ -68,49 +77,56 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertEquals(numDocs, values.size()); + for (int i = 0; i < numDocs; i++) { + assertArrayEquals(vectors[i], values.vectorValue(i), 0.0f); + } } } } // Basic test for multiple vectors fields per document public void testTwoVectorFieldsPerDoc() throws Exception { - float[][] f1 = new float[][] {randomVector(384), randomVector(384)}; - float[][] f2 = new float[][] {randomVector(384), randomVector(384)}; + + final int numDocs = 2; + float[][] vectors1 = dataProvider.getVectors(numDocs); + float[][] vectors2 = dataProvider.getVectors(numDocs); + try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { - Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); - doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); - doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); - w.addDocument(doc1); - Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); - doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); - doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); - w.addDocument(doc2); + + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vectors1[i], EUCLIDEAN)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD2, vectors2[i], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f1"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); - - values = r.getFloatVectorValues("f2"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); - 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); + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD1); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors1[i], values.vectorValue(i), 0.0f); + } + + for (int i = 0; i < numDocs; i++) { + FloatVectorValues values = r.getFloatVectorValues(VECTOR_FIELD2); + assertNotNull(values); + assertEquals(2, values.size()); + assertArrayEquals(vectors2[i], values.vectorValue(i), 0.0f); + } + + // Check boundary condition - search with a 0 topK + float[] target = dataProvider.getVectors(1)[0]; + var topDocs = r.searchNearestVectors(VECTOR_FIELD1, target, 0, null, 10); assertEquals(0, topDocs.scoreDocs.length); assertEquals(0, topDocs.totalHits.value()); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java b/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java new file mode 100644 index 00000000..c4480e84 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVectors; + +import java.util.Random; +import java.util.logging.Level; +import java.util.logging.Logger; + +public class TestDataProvider { + + private static final Logger log = Logger.getLogger(TestDataProvider.class.getName()); + private static final int DATASET_SIZE_LIMIT = 1000; + private static final int DATASET_SIZE_MIN = 200; + private static final int DIMENSIONS_LIMIT = 256; + private static final int DIMENSIONS_MIN = 8; + private static final int TOP_K_LIMIT = 64; + private static final int TOP_K_MIN = 2; + private static final int QUERIES_LIMIT = 50; + private static final int QUERIES_MIN = 2; + + public static final String ID_FIELD = "id"; + public static final String TEXT_FIELD = "some_text_field"; + public static final String CATEGORY_FIELD = "category_field"; + public static final String VECTOR_FIELD1 = "vector_field1"; + public static final String VECTOR_FIELD2 = "vector_field2"; + + private int datasetSize; + private int dimensions; + private int topK; + private float[][] dataset1; + private float[][] dataset2; + private int numQueries; + private Random random; + + public TestDataProvider(Random random) { + datasetSize = random.nextInt(DATASET_SIZE_MIN, DATASET_SIZE_LIMIT); + dimensions = random.nextInt(DIMENSIONS_MIN, DIMENSIONS_LIMIT); + topK = Math.min(random.nextInt(TOP_K_MIN, TOP_K_LIMIT), datasetSize); + dataset1 = generateRandomVectors(random, datasetSize, dimensions); + dataset2 = generateRandomVectors(random, datasetSize, dimensions); + numQueries = random.nextInt(QUERIES_MIN, QUERIES_LIMIT); + this.random = random; + log.log(Level.FINE, "Dataset size: " + datasetSize + "x" + dimensions + ", topK: " + topK); + } + + public int getDatasetSize() { + return datasetSize; + } + + public int getDimensions() { + return dimensions; + } + + public int getTopK() { + return topK; + } + + public float[][] getDataset1() { + return dataset1; + } + + public float[][] getDataset2() { + return dataset2; + } + + public float[][] getQueries(int numQueries) { + return generateRandomVectors(random, numQueries, dimensions); + } + + public float[][] getVectors(int numVectors) { + return generateRandomVectors(random, numVectors, dimensions); + } + + public int getRandom(int min, int max) { + return random.nextInt(min, max); + } + + public double getRandom(double min, double max) { + return random.nextDouble(min, max); + } + + public int getNumQueries() { + return numQueries; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index bebc5ead..13dac867 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -4,16 +4,25 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestDataProvider.ID_FIELD; +import static com.nvidia.cuvs.lucene.TestDataProvider.VECTOR_FIELD1; +import static com.nvidia.cuvs.lucene.TestUtils.createWriter; +import static com.nvidia.cuvs.lucene.TestUtils.createWriterConfig; +import static com.nvidia.cuvs.lucene.TestUtils.generateRandomText; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.io.IOException; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; import java.util.Random; +import java.util.Set; import java.util.logging.Level; import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; import org.apache.lucene.document.KnnFloatVectorField; @@ -24,10 +33,10 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.SortedDocValues; import org.apache.lucene.index.Term; import org.apache.lucene.index.TieredMergePolicy; -import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.ScoreDoc; @@ -36,6 +45,7 @@ import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.util.BytesRef; @@ -54,31 +64,23 @@ public class TestMerge extends LuceneTestCase { private static final Logger log = Logger.getLogger(TestMerge.class.getName()); - - private static final int MIN_VECTOR_DIMENSION = 64; - private static final int MAX_VECTOR_DIMENSION = 256; - private static final int TOP_K_LIMIT = 64; + private static Random random; + private static TestDataProvider dataProvider; + private static Directory directory; + private static Codec codec; @BeforeClass public static void beforeClass() { - assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported()); + assumeTrue("cuVS not supported so skipping these tests", CuVS2510GPUVectorsFormat.supported()); + random = random(); + codec = alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); } - private Directory directory; - private int vectorDimension; - @Before public void setUp() throws Exception { super.setUp(); directory = newDirectory(); - - // Randomize vector dimension for each test - vectorDimension = - MIN_VECTOR_DIMENSION + random().nextInt(MAX_VECTOR_DIMENSION - MIN_VECTOR_DIMENSION + 1); - // Ensure dimension is multiple of 4 for better performance - vectorDimension = (vectorDimension / 4) * 4; - - log.log(Level.FINE, "Using randomized vector dimension: " + vectorDimension); + dataProvider = new TestDataProvider(random); } @After @@ -94,16 +96,12 @@ public void tearDown() throws Exception { **/ @Test public void testMergeManyDocumentsMultipleSegments() throws IOException { - log.log(Level.FINE, "Starting testMergeManyDocumentsMultipleSegments"); - // Randomize configuration parameters - int maxBufferedDocs = 5 + random().nextInt(16); // 5-20 docs per buffer - int totalBatches = 8 + random().nextInt(8); // 8-15 batches - int docsPerBatch = 15 + random().nextInt(11); // 15-25 docs per batch + int maxBufferedDocs = dataProvider.getRandom(5, 16); + int totalBatches = dataProvider.getRandom(8, 16); + int docsPerBatch = dataProvider.getRandom(15, 25); int totalDocuments = totalBatches * docsPerBatch; - - // Randomize vector presence probability (60-85%) - double vectorProbability = 0.6 + (random().nextDouble() * 0.25); + double vectorProbability = dataProvider.getRandom(0.6, 0.8); log.log( Level.FINE, @@ -118,37 +116,27 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { + ", vectorProbability=" + vectorProbability); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) - .setMaxBufferedDocs(maxBufferedDocs) // Randomized buffer size - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - - List expectedVectors = new ArrayList<>(); List expectedDocIds = new ArrayList<>(); int documentsWithVectors = 0; - try (IndexWriter writer = new IndexWriter(directory, config)) { + try (RandomIndexWriter writer = createWriter(random, directory, codec)) { // Add documents in multiple batches to create many segments for (int batch = 0; batch < totalBatches; batch++) { for (int i = 0; i < docsPerBatch; i++) { int docId = batch * docsPerBatch + i; Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); - doc.add(new NumericDocValuesField("batch", batch)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); // Randomly decide if document has vector if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); - expectedVectors.add(vector); + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); expectedDocIds.add(docId); documentsWithVectors++; } - writer.addDocument(doc); } - writer.commit(); // Create a new segment + writer.commit(); } int documentsWithoutVectors = totalDocuments - documentsWithVectors; @@ -157,84 +145,66 @@ public void testMergeManyDocumentsMultipleSegments() throws IOException { log.log(Level.FINE, "Documents with vectors: " + documentsWithVectors); log.log(Level.FINE, "Documents without vectors: " + documentsWithoutVectors); - // Force merge to trigger merge logic writer.forceMerge(1); log.log(Level.FINE, "Forced merge to single segment completed"); } // Verify the merged index try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); - // Verify vector search works correctly after merge - if (documentsWithVectors > 0) { - IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); - - // Randomize search parameters - int searchK = - Math.min(5 + random().nextInt(10), Math.min(documentsWithVectors, TOP_K_LIMIT)); + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getTopK(); - KnnFloatVectorQuery query = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs results = searcher.search(query, searchK); + KnnFloatVectorQuery query = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs results = searcher.search(query, topK); - assertTrue("Should find some results after merge", results.scoreDocs.length > 0); - assertTrue( - "Should find reasonable number of results", - results.scoreDocs.length <= documentsWithVectors); + assertTrue("Should find some results after merge", results.scoreDocs.length > 0); + assertTrue("Should find some results", results.scoreDocs.length <= documentsWithVectors); - log.log( - Level.FINE, - "Vector search returned " - + results.scoreDocs.length - + " results out of " - + documentsWithVectors - + " documents with vectors"); - - // Verify all returned documents have valid IDs - for (ScoreDoc scoreDoc : results.scoreDocs) { - int docId = Integer.parseInt(searcher.storedFields().document(scoreDoc.doc).get("id")); - assertTrue("Document ID should be valid", docId >= 0 && docId < totalDocuments); - } - } else { - log.log(Level.FINE, "No documents with vectors - skipping vector search verification"); + log.log( + Level.FINE, + "Vector search returned " + + results.scoreDocs.length + + " results out of " + + documentsWithVectors + + " documents with vectors"); + + // Verify all returned documents have valid IDs + for (ScoreDoc scoreDoc : results.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + int docId = Integer.parseInt(doc.get(ID_FIELD)); + assertTrue("Document ID should be valid", expectedDocIds.contains(docId)); } - - log.log(Level.FINE, "Merge verification completed successfully"); } } /** - * Test merging with index sorting enabled using text-based sorting and SortingMergePolicy + * Test merging with index sorting enabled using SortingMergePolicy **/ @Test - public void testMergeWithIndexSorting() throws IOException { - log.log(Level.FINE, "Starting testMergeWithIndexSorting with text-based sorting"); - + public void testMergeWithIndexSortingStringField() throws IOException { // Randomize sort field type - SortField.Type sortType = random().nextBoolean() ? SortField.Type.STRING : SortField.Type.LONG; - String sortFieldName = sortType == SortField.Type.STRING ? "text_sort_key" : "numeric_sort_key"; - + final String SORT_FIELD_NAME = "text_sort_key"; + final String ORIGINAL_ORDER = "original_order"; // Configure index sorting by a randomized field - Sort indexSort = new Sort(new SortField(sortFieldName, sortType)); + Sort indexSort = new Sort(new SortField(SORT_FIELD_NAME, SortField.Type.STRING)); // Randomize merge policy parameters TieredMergePolicy mergePolicy = new TieredMergePolicy(); - mergePolicy.setMaxMergedSegmentMB(128 + random().nextInt(257)); // 128-384 MB - mergePolicy.setSegmentsPerTier(3 + random().nextInt(4)); // 3-6 segments per tier + mergePolicy.setMaxMergedSegmentMB(dataProvider.getRandom(128, 385)); + mergePolicy.setSegmentsPerTier(dataProvider.getRandom(3, 7)); // Randomize writer configuration parameters - int maxBufferedDocs = 10 + random().nextInt(16); // 10-25 docs per buffer - int totalDocuments = 80 + random().nextInt(81); // 80-160 documents - int segmentSize = 15 + random().nextInt(11); // 15-25 docs per segment - double vectorProbability = 0.65 + (random().nextDouble() * 0.25); // 65-90% have vectors + int maxBufferedDocs = dataProvider.getRandom(10, 26); + int totalDocuments = dataProvider.getRandom(80, 161); + int segmentSize = dataProvider.getRandom(15, 26); + double vectorProbability = dataProvider.getRandom(0.65, 0.91); - log.log( - Level.FINE, - "Randomized sorting parameters: sortType=" + sortType + ", sortFieldName=" + sortFieldName); log.log( Level.FINE, "Randomized config: maxBufferedDocs=" @@ -248,46 +218,29 @@ public void testMergeWithIndexSorting() throws IOException { IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) - .setIndexSort(indexSort) // This automatically enables sorting during merges + .setCodec(codec) + .setIndexSort(indexSort) .setMergePolicy(mergePolicy) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - // List documents = new ArrayList<>(); - try (IndexWriter writer = new IndexWriter(directory, config)) { // Create documents with randomized sort keys + int numDocsWithVectors = 0; for (int i = 0; i < totalDocuments; i++) { - float[] vector = null; - - // Randomly decide if document has vector - if (random().nextDouble() < vectorProbability) { - vector = generateRandomVector(vectorDimension, random()); - } Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); - doc.add(new StringField("original_order", String.valueOf(i), Field.Store.YES)); - - // Add sort field based on randomized type - if (sortType == SortField.Type.STRING) { - // Randomize text sort key length (4-12 characters) - int keyLength = 4 + random().nextInt(9); - String textSortKey = generateRandomText(random(), keyLength); - doc.add(new SortedDocValuesField(sortFieldName, new BytesRef(textSortKey))); - doc.add(new StringField(sortFieldName + "_stored", textSortKey, Field.Store.YES)); - } else { - // Use numeric sort key with wider range - long numericSortKey = random().nextLong() % 100000; // Can be negative for more variety - doc.add(new NumericDocValuesField(sortFieldName, numericSortKey)); - doc.add( - new StringField( - sortFieldName + "_stored", String.valueOf(numericSortKey), Field.Store.YES)); - } + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new StringField(ORIGINAL_ORDER, String.valueOf(i), Field.Store.YES)); - if (vector != null) { - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + String textSortKey = generateRandomText(random, dataProvider.getRandom(4, 21)); + doc.add(new SortedDocValuesField(SORT_FIELD_NAME, new BytesRef(textSortKey))); + doc.add(new StringField(SORT_FIELD_NAME + "_stored", textSortKey, Field.Store.YES)); + + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + numDocsWithVectors++; } writer.addDocument(doc); @@ -305,7 +258,13 @@ public void testMergeWithIndexSorting() throws IOException { } } - log.log(Level.FINE, "Created " + totalDocuments + " documents with text-based index sorting"); + log.log( + Level.FINE, + "Number of documents with vectors is: " + + numDocsWithVectors + + " out of a total of " + + totalDocuments + + " documents"); // Force merge with sorting - this will use the sorting merge policy writer.forceMerge(1); @@ -314,95 +273,34 @@ public void testMergeWithIndexSorting() throws IOException { // Verify the merged and sorted index try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); - // Verify documents are sorted correctly by the randomized sort field - log.log( - Level.FINE, - "Verifying document sorting order using sortType: " - + sortType - + ", field: " - + sortFieldName); - - if (sortType == SortField.Type.STRING) { - // Verify string-based sorting - String previousSortKey = ""; - SortedDocValues sortedValues = leafReader.getSortedDocValues(sortFieldName); - - for (int docId = 0; docId < leafReader.maxDoc(); docId++) { - String currentSortKey = ""; - if (sortedValues != null && sortedValues.advanceExact(docId)) { - currentSortKey = sortedValues.lookupOrd(sortedValues.ordValue()).utf8ToString(); - } + String previousSortKey = ""; + SortedDocValues sortedValues = leafReader.getSortedDocValues(SORT_FIELD_NAME); - assertTrue( - "Documents should be sorted by " - + sortFieldName - + ": '" - + previousSortKey - + "' should be <= '" - + currentSortKey - + "'", - previousSortKey.compareTo(currentSortKey) <= 0); - previousSortKey = currentSortKey; - - // Log first 10 documents to verify sorting - if (docId < 10) { - IndexSearcher searcher = new IndexSearcher(reader); - String originalOrder = searcher.storedFields().document(docId).get("original_order"); - log.log( - Level.FINE, - "DocId: " - + docId - + ", OriginalOrder: " - + originalOrder - + ", SortKey: '" - + currentSortKey - + "'"); - } + for (int docId = 0; docId < leafReader.maxDoc(); docId++) { + String currentSortKey = ""; + if (sortedValues != null && sortedValues.advanceExact(docId)) { + currentSortKey = sortedValues.lookupOrd(sortedValues.ordValue()).utf8ToString(); } - } else { - // Verify numeric-based sorting - long previousSortKey = Long.MIN_VALUE; - var numericValues = leafReader.getNumericDocValues(sortFieldName); - - for (int docId = 0; docId < leafReader.maxDoc(); docId++) { - long currentSortKey = Long.MIN_VALUE; - if (numericValues != null && numericValues.advanceExact(docId)) { - currentSortKey = numericValues.longValue(); - } - assertTrue( - "Documents should be sorted by " - + sortFieldName - + ": " - + previousSortKey - + " should be <= " - + currentSortKey, - previousSortKey <= currentSortKey); - previousSortKey = currentSortKey; - - // Log first 10 documents to verify sorting - if (docId < 10) { - IndexSearcher searcher = new IndexSearcher(reader); - String originalOrder = searcher.storedFields().document(docId).get("original_order"); - log.log( - Level.FINE, - "DocId: " - + docId - + ", OriginalOrder: " - + originalOrder - + ", SortKey: " - + currentSortKey); - } - } + assertTrue( + "Documents should be sorted by " + + SORT_FIELD_NAME + + ": '" + + previousSortKey + + "' should be <= '" + + currentSortKey + + "'", + previousSortKey.compareTo(currentSortKey) <= 0); + previousSortKey = currentSortKey; } // Count total vectors by checking if vector field exists and has values - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int documentsWithVectors = vectorValues != null ? vectorValues.size() : 0; log.log( @@ -412,10 +310,10 @@ public void testMergeWithIndexSorting() throws IOException { // Test vector search on sorted index if (documentsWithVectors > 0) { IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); + float[] queryVector = dataProvider.getQueries(1)[0]; KnnFloatVectorQuery query = - new KnnFloatVectorQuery("vector", queryVector, Math.min(10, documentsWithVectors)); + new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, Math.min(10, documentsWithVectors)); TopDocs results = searcher.search(query, 10); assertTrue("Should find results in sorted index", results.scoreDocs.length > 0); @@ -425,12 +323,11 @@ public void testMergeWithIndexSorting() throws IOException { // Verify that returned documents maintain sort order if we check their sort keys log.log(Level.FINE, "Verifying vector search results maintain sorting consistency..."); - for (int i = 0; i < Math.min(3, results.scoreDocs.length); i++) { + for (int i = 0; i < Math.min(5, results.scoreDocs.length); i++) { ScoreDoc scoreDoc = results.scoreDocs[i]; - String originalOrder = - searcher.storedFields().document(scoreDoc.doc).get("original_order"); - String sortKey = - searcher.storedFields().document(scoreDoc.doc).get(sortFieldName + "_stored"); + Document doc = searcher.storedFields().document(scoreDoc.doc); + String originalOrder = doc.get(ORIGINAL_ORDER); + String sortKey = doc.get(SORT_FIELD_NAME + "_stored"); log.log( Level.FINE, "Result " @@ -445,61 +342,212 @@ public void testMergeWithIndexSorting() throws IOException { + scoreDoc.score); } } - - log.log(Level.FINE, "Text-based index sorting verification completed successfully"); } } /** - * Test merging segments with various patterns of missing vectors + * Test merging with index sorting enabled using SortingMergePolicy **/ @Test - public void testMergeWithMissingVectors() throws IOException { - log.log(Level.FINE, "Starting testMergeWithMissingVectors"); + public void testMergeWithIndexSortingLongField() throws IOException { + final String SORT_FIELD_NAME = "numeric_sort_key"; + final String ORIGINAL_ORDER = "original_order"; + Sort indexSort = new Sort(new SortField(SORT_FIELD_NAME, SortField.Type.LONG)); - // Randomize configuration - int maxBufferedDocs = 10 + random().nextInt(11); // 10-20 docs per buffer - int numSegments = 3 + random().nextInt(3); // 3-5 segments + // Randomize merge policy parameters + TieredMergePolicy mergePolicy = new TieredMergePolicy(); + mergePolicy.setMaxMergedSegmentMB(dataProvider.getRandom(128, 385)); + mergePolicy.setSegmentsPerTier(dataProvider.getRandom(3, 7)); + + // Randomize writer configuration parameters + int maxBufferedDocs = dataProvider.getRandom(10, 26); + int totalDocuments = dataProvider.getRandom(80, 161); + int segmentSize = dataProvider.getRandom(15, 26); + double vectorProbability = dataProvider.getRandom(0.65, 0.91); + + log.log( + Level.FINE, + "Randomized config: maxBufferedDocs=" + + maxBufferedDocs + + ", totalDocuments=" + + totalDocuments + + ", segmentSize=" + + segmentSize + + ", vectorProbability=" + + vectorProbability); IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) + .setCodec(codec) + .setIndexSort(indexSort) + .setMergePolicy(mergePolicy) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - log.log( - Level.FINE, - "Randomized parameters: maxBufferedDocs=" - + maxBufferedDocs - + ", numSegments=" - + numSegments); + try (IndexWriter writer = new IndexWriter(directory, config)) { + int numDocsWithVectors = 0; + for (int i = 0; i < totalDocuments; i++) { + + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); + doc.add(new StringField(ORIGINAL_ORDER, String.valueOf(i), Field.Store.YES)); + + long numericSortKey = random.nextLong() % 100000; + doc.add(new NumericDocValuesField(SORT_FIELD_NAME, numericSortKey)); + doc.add( + new StringField( + SORT_FIELD_NAME + "_stored", String.valueOf(numericSortKey), Field.Store.YES)); + + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + numDocsWithVectors++; + } + + writer.addDocument(doc); + + // Commit based on randomized segment size + if ((i + 1) % segmentSize == 0) { + writer.commit(); + log.log( + Level.FINE, + "Committed segment " + + ((i + 1) / segmentSize) + + " with " + + (i + 1) + + " total documents"); + } + } + + log.log( + Level.FINE, + "Number of documents with vectors is: " + + numDocsWithVectors + + " out of a total of " + + totalDocuments + + " documents"); + + // Force merge with sorting - this will use the sorting merge policy + writer.forceMerge(1); + log.log(Level.FINE, "Forced merge with text-based sorting completed"); + } + + // Verify the merged and sorted index + try (DirectoryReader reader = DirectoryReader.open(directory)) { + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); + assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); + + // Verify numeric-based sorting + long previousSortKey = Long.MIN_VALUE; + var numericValues = leafReader.getNumericDocValues(SORT_FIELD_NAME); + + for (int docId = 0; docId < leafReader.maxDoc(); docId++) { + long currentSortKey = Long.MIN_VALUE; + if (numericValues != null && numericValues.advanceExact(docId)) { + currentSortKey = numericValues.longValue(); + } + + assertTrue( + "Documents should be sorted by " + + SORT_FIELD_NAME + + ": " + + previousSortKey + + " should be <= " + + currentSortKey, + previousSortKey <= currentSortKey); + previousSortKey = currentSortKey; + + // Log first 10 documents to verify sorting + if (docId < 10) { + IndexSearcher searcher = new IndexSearcher(reader); + String originalOrder = searcher.storedFields().document(docId).get(ORIGINAL_ORDER); + log.log( + Level.FINE, + "DocId: " + + docId + + ", OriginalOrder: " + + originalOrder + + ", SortKey: " + + currentSortKey); + } + } + + // Count total vectors by checking if vector field exists and has values + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); + int documentsWithVectors = vectorValues != null ? vectorValues.size() : 0; + + log.log( + Level.FINE, + "Found " + documentsWithVectors + " documents with vectors after sorted merge"); + + // Test vector search on sorted index + if (documentsWithVectors > 0) { + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + + KnnFloatVectorQuery query = + new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, Math.min(10, documentsWithVectors)); + TopDocs results = searcher.search(query, 10); + + assertTrue("Should find results in sorted index", results.scoreDocs.length > 0); + log.log( + Level.FINE, + "Vector search on sorted index returned " + results.scoreDocs.length + " results"); + + // Verify that returned documents maintain sort order if we check their sort keys + log.log(Level.FINE, "Verifying vector search results maintain sorting consistency..."); + for (int i = 0; i < Math.min(5, results.scoreDocs.length); i++) { + ScoreDoc scoreDoc = results.scoreDocs[i]; + Document doc = searcher.storedFields().document(scoreDoc.doc); + String originalOrder = doc.get(ORIGINAL_ORDER); + String sortKey = doc.get(SORT_FIELD_NAME + "_stored"); + log.log( + Level.FINE, + "Result " + + i + + ": DocId=" + + scoreDoc.doc + + ", OriginalOrder=" + + originalOrder + + ", SortKey='" + + sortKey + + "', Score=" + + scoreDoc.score); + } + } + } + } + + /** + * Test merging segments with various patterns of missing vectors + **/ + @Test + public void testMergeWithMissingVectors() throws IOException { + int numSegments = dataProvider.getRandom(3, 13); + IndexWriterConfig config = createWriterConfig(random, codec); + log.log(Level.FINE, "Randomized parameters: numSegments=" + numSegments); int totalExpectedVectors = 0; int totalDocuments = 0; try (IndexWriter writer = new IndexWriter(directory, config)) { for (int seg = 0; seg < numSegments; seg++) { - // Randomize segment characteristics - int docsInSegment = 15 + random().nextInt(16); // 15-30 docs per segment - double vectorProbability = random().nextDouble(); // 0-100% vector probability - String segmentType = "seg_" + seg + "_prob_" + String.format("%.2f", vectorProbability); - + int docsInSegment = dataProvider.getRandom(15, 100); + double vectorProbability = dataProvider.getRandom(0.1, 0.6); int segmentVectorCount = 0; for (int i = 0; i < docsInSegment; i++) { Document doc = new Document(); - doc.add(new StringField("id", "seg" + seg + "_" + i, Field.Store.YES)); - doc.add(new StringField("segment", segmentType, Field.Store.YES)); - doc.add(new NumericDocValuesField("segment_num", seg)); - doc.add(new NumericDocValuesField("doc_in_segment", i)); + doc.add(new StringField(ID_FIELD, String.valueOf(i), Field.Store.YES)); // Randomly add vector based on segment's probability - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); segmentVectorCount++; } - writer.addDocument(doc); } @@ -516,7 +564,7 @@ public void testMergeWithMissingVectors() throws IOException { + " documents, " + segmentVectorCount + " with vectors (probability: " - + String.format("%.2f", vectorProbability) + + vectorProbability + ")"); } @@ -527,13 +575,13 @@ public void testMergeWithMissingVectors() throws IOException { // Verify the merged index handles missing vectors correctly try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); // Count actual vectors in merged index - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; log.log( @@ -550,31 +598,22 @@ public void testMergeWithMissingVectors() throws IOException { // Test vector search if we have vectors if (actualVectorCount > 0) { IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); - - // Randomize search parameters - int searchK = Math.min(5 + random().nextInt(10), Math.min(actualVectorCount, TOP_K_LIMIT)); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = Math.min(dataProvider.getTopK(), actualVectorCount); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); - - assertTrue("Should find some vector results", vectorResults.scoreDocs.length > 0); - assertTrue( - "Should not find more vectors than exist", - vectorResults.scoreDocs.length <= actualVectorCount); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); log.log( Level.FINE, - "Found " - + vectorResults.scoreDocs.length - + " vector results out of " - + actualVectorCount - + " available"); + "Found " + numResults + " vector results out of " + actualVectorCount + " available"); + assertEquals("Search should return exactly topK results", topK, numResults); } else { log.log(Level.FINE, "No vectors in merged index - skipping vector search"); } - - log.log(Level.FINE, "Missing vectors test completed successfully"); } } @@ -583,20 +622,14 @@ public void testMergeWithMissingVectors() throws IOException { **/ @Test public void testMergeWithDeletions() throws IOException { - log.log(Level.FINE, "Starting testMergeWithDeletions"); - - // Randomize configuration parameters - int maxBufferedDocs = 15 + random().nextInt(11); // 15-25 docs per buffer - int numSegments = 3 + random().nextInt(4); // 3-6 segments - int docsPerSegment = 20 + random().nextInt(21); // 20-40 docs per segment - double vectorProbability = 0.7 + (random().nextDouble() * 0.25); // 70-95% have vectors - double deletionProbability = 0.2 + (random().nextDouble() * 0.3); // 20-50% deletion rate + int numSegments = dataProvider.getRandom(3, 7); + int docsPerSegment = dataProvider.getRandom(20, 41); + double vectorProbability = dataProvider.getRandom(0.7, 0.95); + double deletionProbability = dataProvider.getRandom(0.2, 0.5); log.log( Level.FINE, - "Randomized parameters: maxBufferedDocs=" - + maxBufferedDocs - + ", numSegments=" + "Randomized parameters: numSegments=" + numSegments + ", docsPerSegment=" + docsPerSegment @@ -605,15 +638,11 @@ public void testMergeWithDeletions() throws IOException { + ", deletionProbability=" + deletionProbability); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) - .setMaxBufferedDocs(maxBufferedDocs) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - + IndexWriterConfig config = createWriterConfig(random, codec); List expectedRemainingDocs = new ArrayList<>(); List deletedDocs = new ArrayList<>(); int totalDocuments = numSegments * docsPerSegment; + int numDocsWithVectors = 0; try (IndexWriter writer = new IndexWriter(directory, config)) { // Create multiple segments with documents @@ -621,17 +650,14 @@ public void testMergeWithDeletions() throws IOException { for (int i = 0; i < docsPerSegment; i++) { int docId = seg * docsPerSegment + i; Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); - doc.add(new StringField("segment", "seg_" + seg, Field.Store.YES)); - doc.add(new NumericDocValuesField("doc_num", docId)); - doc.add(new NumericDocValuesField("segment_num", seg)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); // Randomly add vectors - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getQueries(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + numDocsWithVectors++; } - writer.addDocument(doc); } writer.commit(); @@ -646,12 +672,13 @@ public void testMergeWithDeletions() throws IOException { + " documents each (" + totalDocuments + " total)"); + log.log(Level.FINE, "Number of docs with vectors: " + numDocsWithVectors); // Delete documents randomly and track which ones are deleted int deletedCount = 0; for (int docId = 0; docId < totalDocuments; docId++) { - if (random().nextDouble() < deletionProbability) { - writer.deleteDocuments(new Term("id", String.valueOf(docId))); + if (random.nextDouble() < deletionProbability) { + writer.deleteDocuments(new Term(ID_FIELD, String.valueOf(docId))); deletedDocs.add(docId); deletedCount++; } else { @@ -664,7 +691,7 @@ public void testMergeWithDeletions() throws IOException { "Deleted " + deletedCount + " documents (" - + String.format("%.1f", (100.0 * deletedCount / totalDocuments)) + + (100.0 * deletedCount / totalDocuments) + "%), remaining: " + expectedRemainingDocs.size()); @@ -677,9 +704,9 @@ public void testMergeWithDeletions() throws IOException { // Verify the merged index correctly handles deletions try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); int expectedRemaining = expectedRemainingDocs.size(); assertEquals( "Should have correct number of documents after deletions", @@ -692,31 +719,30 @@ public void testMergeWithDeletions() throws IOException { // Test that we can find expected remaining documents for (int i = 0; i < Math.min(10, expectedRemainingDocs.size()); i++) { int docId = expectedRemainingDocs.get(i); - TopDocs result = searcher.search(new TermQuery(new Term("id", String.valueOf(docId))), 1); + TopDocs result = + searcher.search(new TermQuery(new Term(ID_FIELD, String.valueOf(docId))), 1); assertEquals("Should find remaining document " + docId, 1, (int) result.totalHits.value()); } // Test that actually deleted documents are not found - int deletedDocsToCheck = Math.min(10, deletedDocs.size()); // Check up to 10 deleted docs + int deletedDocsToCheck = Math.min(10, deletedDocs.size()); for (int i = 0; i < deletedDocsToCheck; i++) { int docId = deletedDocs.get(i); - TopDocs result = searcher.search(new TermQuery(new Term("id", String.valueOf(docId))), 1); + TopDocs result = + searcher.search(new TermQuery(new Term(ID_FIELD, String.valueOf(docId))), 1); assertEquals( "Should not find deleted document " + docId, 0, (int) result.totalHits.value()); } // Test vector search works after deletions - float[] queryVector = generateRandomVector(vectorDimension, random()); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, 10); - TopDocs vectorResults = searcher.search(vectorQuery, 10); - - assertTrue( - "Should find some vector results after deletions", vectorResults.scoreDocs.length > 0); - - log.log( - Level.FINE, - "Found " + vectorResults.scoreDocs.length + " vector results after deletions"); - log.log(Level.FINE, "Deletion merge verification completed successfully"); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = Math.min(1, numDocsWithVectors); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results after deletions", numResults > 0); + assertEquals("Search should return exactly topK documents", numResults, topK); + log.log(Level.FINE, "Found " + numResults + " vector results after deletions"); } } @@ -725,13 +751,10 @@ public void testMergeWithDeletions() throws IOException { * */ @Test public void testMergeBruteForceIndex() throws IOException { - log.log(Level.FINE, "Starting testMergeBruteForceIndex"); - - // Randomize configuration parameters - int maxBufferedDocs = 8 + random().nextInt(8); // 8-15 docs per buffer - int numSegments = 3 + random().nextInt(3); // 3-5 segments - int docsPerSegment = 12 + random().nextInt(9); // 12-20 docs per segment - double vectorProbability = 0.8 + (random().nextDouble() * 0.2); // 80-100% have vectors + int numSegments = dataProvider.getRandom(3, 10); + int docsPerSegment = dataProvider.getRandom(20, 100); + double vectorProbability = dataProvider.getRandom(0.2, 0.7); + int maxBufferedDocs = dataProvider.getRandom(8, 17); log.log( Level.FINE, @@ -747,11 +770,7 @@ public void testMergeBruteForceIndex() throws IOException { // Configure with brute force index type CuVS2510GPUVectorsFormat bruteForceFormat = new CuVS2510GPUVectorsFormat( - 32, // writer threads - 128, // intermediate graph degree - 64, // graph degree - CagraGraphBuildAlgo.NN_DESCENT, - IndexType.BRUTE_FORCE); // Use brute force index + 32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, IndexType.BRUTE_FORCE); IndexWriterConfig config = new IndexWriterConfig() @@ -761,27 +780,23 @@ public void testMergeBruteForceIndex() throws IOException { int totalDocuments = numSegments * docsPerSegment; int totalExpectedVectors = 0; + Set docIDsHavingVectors = new HashSet(); try (IndexWriter writer = new IndexWriter(directory, config)) { - // Create multiple segments with brute force index for (int seg = 0; seg < numSegments; seg++) { int segmentVectorCount = 0; - for (int i = 0; i < docsPerSegment; i++) { int docId = seg * docsPerSegment + i; Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); - doc.add(new StringField("segment", "seg_" + seg, Field.Store.YES)); - doc.add(new NumericDocValuesField("segment_num", seg)); - doc.add(new NumericDocValuesField("doc_in_segment", i)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); // Randomly add vectors based on probability - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getVectors(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + docIDsHavingVectors.add(docId); segmentVectorCount++; } - writer.addDocument(doc); } @@ -816,13 +831,13 @@ public void testMergeBruteForceIndex() throws IOException { // Verify the merged brute force index try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); // Count actual vectors in merged index - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; log.log( @@ -837,58 +852,47 @@ public void testMergeBruteForceIndex() throws IOException { assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); // Test brute force vector search (exact search) - if (actualVectorCount > 0) { - IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); - - // Search for reasonable number of results - int searchK = Math.min(8 + random().nextInt(8), Math.min(actualVectorCount, TOP_K_LIMIT)); + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getRandom(1, actualVectorCount); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results in brute force index", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); + log.log( + Level.FINE, + "Brute force search found " + + numResults + + " results out of " + + actualVectorCount + + " available vectors"); + + assertEquals("Search should return exactly topK documents", topK, numResults); + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + String docId = doc.get(ID_FIELD); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); assertTrue( - "Should find some vector results in brute force index", - vectorResults.scoreDocs.length > 0); - assertTrue( - "Should not find more vectors than exist", - vectorResults.scoreDocs.length <= actualVectorCount); - - log.log( - Level.FINE, - "Brute force search found " - + vectorResults.scoreDocs.length - + " results out of " - + actualVectorCount - + " available vectors"); - - // Verify all returned documents are valid - for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { - String docId = searcher.storedFields().document(scoreDoc.doc).get("id"); - assertNotNull("Document should have valid ID", docId); - assertTrue("Score should be positive", scoreDoc.score > 0); - } - } else { - log.log(Level.FINE, "No vectors in brute force merged index - skipping vector search"); + "Document does not have a vector", + docIDsHavingVectors.contains(Integer.parseInt(docId))); } - - log.log(Level.FINE, "Brute force merge verification completed successfully"); } } /** - * Test merging segments for {@link IndexType#CAGRA_AND_BRUTE_FORCE} + * Test merging segments for {@link IndexType#CAGRA} * */ @Test - public void testMergeCagraAndBruteForceIndex() throws IOException { - log.log(Level.FINE, "Starting testMergeCagraAndBruteForceIndex"); - - // Use moderate dataset size - int maxBufferedDocs = 15 + random().nextInt(10); // 15-24 docs per buffer - int numSegments = - 4; // Fixed 4 segments: alternating CAGRA vs small segments (brute force fallback) - int docsPerSegment = 20 + random().nextInt(11); // 20-30 docs per segment - double vectorProbability = 0.9 + (random().nextDouble() * 0.1); // 90-100% have vectors + public void testMergeCagraIndex() throws IOException { + int numSegments = dataProvider.getRandom(3, 10); + int docsPerSegment = dataProvider.getRandom(20, 100); + double vectorProbability = dataProvider.getRandom(0.2, 0.7); + int maxBufferedDocs = dataProvider.getRandom(8, 17); log.log( Level.FINE, @@ -901,45 +905,35 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { + ", vectorProbability=" + vectorProbability); - // Configure with CAGRA + brute force combined index type - CuVS2510GPUVectorsFormat combinedFormat = - new CuVS2510GPUVectorsFormat( - 32, // writer threads - 128, // intermediate graph degree - 64, // graph degree - CagraGraphBuildAlgo.NN_DESCENT, - IndexType.CAGRA_AND_BRUTE_FORCE); // Use combined CAGRA + brute force + // Configure with CAGRA index type + CuVS2510GPUVectorsFormat cagraAndBruteForceFormat = + new CuVS2510GPUVectorsFormat(32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, IndexType.CAGRA); IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(combinedFormat)) + .setCodec(alwaysKnnVectorsFormat(cagraAndBruteForceFormat)) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); int totalDocuments = numSegments * docsPerSegment; int totalExpectedVectors = 0; + Set docIDsHavingVectors = new HashSet(); try (IndexWriter writer = new IndexWriter(directory, config)) { - // Create segments that will result in mixed index types during merge for (int seg = 0; seg < numSegments; seg++) { int segmentVectorCount = 0; - for (int i = 0; i < docsPerSegment; i++) { int docId = seg * docsPerSegment + i; Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); - doc.add(new StringField("segment", "mixed_seg_" + seg, Field.Store.YES)); - doc.add(new StringField("index_type", "cagra_and_brute_force", Field.Store.YES)); - doc.add(new NumericDocValuesField("segment_num", seg)); - doc.add(new NumericDocValuesField("doc_in_segment", i)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); - // Add vectors based on probability - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + // Randomly add vectors based on probability + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getVectors(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + docIDsHavingVectors.add(docId); segmentVectorCount++; } - writer.addDocument(doc); } @@ -948,7 +942,7 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { log.log( Level.FINE, - "Created CAGRA+brute force segment " + "Created CAGRA segment " + seg + ": " + docsPerSegment @@ -961,31 +955,31 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { Level.FINE, "Created " + numSegments - + " CAGRA+brute force segments with " + + " CAGRA segments with " + totalDocuments + " total documents and " + totalExpectedVectors + " vectors"); - // Force merge all CAGRA+brute force segments + // Force merge all CAGRA segments writer.forceMerge(1); - log.log(Level.FINE, "Forced merge of CAGRA+brute force segments completed"); + log.log(Level.FINE, "Forced merge of brute force segments completed"); } - // Verify the merged CAGRA+brute force index + // Verify the merged CAGRA index try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); // Count actual vectors in merged index - var vectorValues = leafReader.getFloatVectorValues("vector"); + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; log.log( Level.FINE, - "CAGRA+brute force merge results: Total documents: " + "CAGRA merge results: Total documents: " + totalDocuments + ", Expected vectors: " + totalExpectedVectors @@ -994,209 +988,174 @@ public void testMergeCagraAndBruteForceIndex() throws IOException { assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); - // Test CAGRA+brute force index vector search - if (actualVectorCount > 0) { - IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); - - // Search for reasonable number of results - int searchK = Math.min(12 + random().nextInt(8), Math.min(actualVectorCount, TOP_K_LIMIT)); + // Test CAGRA vector search (exact search) + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getRandom(1, actualVectorCount); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results in CAGRA + brute force index", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); + log.log( + Level.FINE, + "CAGRA search found " + + numResults + + " results out of " + + actualVectorCount + + " available vectors"); + + assertEquals("Search should return exactly topK documents", topK, numResults); + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + String docId = doc.get(ID_FIELD); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); assertTrue( - "Should find some vector results in CAGRA+brute force index", - vectorResults.scoreDocs.length > 0); - assertTrue( - "Should not find more vectors than exist", - vectorResults.scoreDocs.length <= actualVectorCount); - - log.log( - Level.FINE, - "CAGRA+brute force index search found " - + vectorResults.scoreDocs.length - + " results out of " - + actualVectorCount - + " available vectors"); - - // Verify all returned documents are valid and have expected metadata - for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { - Document resultDoc = searcher.storedFields().document(scoreDoc.doc); - String docId = resultDoc.get("id"); - String indexType = resultDoc.get("index_type"); - - assertNotNull("Document should have valid ID", docId); - assertEquals( - "Document should be marked as CAGRA+brute force index type", - "cagra_and_brute_force", - indexType); - assertTrue("Score should be positive", scoreDoc.score > 0); - } - - // Test that the CAGRA+brute force index handles both approximate and exact search - // consistently - for (int trial = 0; trial < 3; trial++) { - float[] trialQueryVector = generateRandomVector(vectorDimension, random()); - KnnFloatVectorQuery trialQuery = - new KnnFloatVectorQuery("vector", trialQueryVector, Math.min(5, actualVectorCount)); - TopDocs trialResults = searcher.search(trialQuery, Math.min(5, actualVectorCount)); - - assertTrue("Trial " + trial + " should find results", trialResults.scoreDocs.length > 0); - log.log( - Level.FINE, - "Trial " + trial + " found " + trialResults.scoreDocs.length + " results"); - } - } else { - log.log( - Level.FINE, "No vectors in CAGRA+brute force merged index - skipping vector search"); + "Document does not have a vector", + docIDsHavingVectors.contains(Integer.parseInt(docId))); } - - log.log(Level.FINE, "CAGRA+brute force merge verification completed successfully"); } } /** - * Test large scale merge to stress test the system - **/ + * Test merging segments for {@link IndexType#CAGRA_AND_BRUTE_FORCE} + * */ @Test - public void testLargeScaleMerge() throws IOException { - assumeTrue( - "testLargeScaleMerge requires -DlargeScale=true", - Boolean.parseBoolean(System.getProperty("largeScale", "false"))); - - log.log(Level.FINE, "Starting testLargeScaleMerge"); - - // Randomize large scale parameters - int maxBufferedDocs = 40 + random().nextInt(21); // 40-60 docs per buffer - int segmentCount = 15 + random().nextInt(11); // 15-25 segments - int docsPerSegment = 30 + random().nextInt(21); // 30-50 docs per segment - int totalDocuments = segmentCount * docsPerSegment; + public void testMergeCagraAndBruteForceIndex() throws IOException { + int numSegments = dataProvider.getRandom(3, 10); + int docsPerSegment = dataProvider.getRandom(20, 100); + double vectorProbability = dataProvider.getRandom(0.2, 0.7); + int maxBufferedDocs = dataProvider.getRandom(8, 17); log.log( Level.FINE, - "Randomized large scale parameters: maxBufferedDocs=" + "Randomized parameters: maxBufferedDocs=" + maxBufferedDocs - + ", segmentCount=" - + segmentCount + + ", numSegments=" + + numSegments + ", docsPerSegment=" + docsPerSegment - + ", totalDocuments=" - + totalDocuments); + + ", vectorProbability=" + + vectorProbability); + + // Configure with CAGRA + brute force index type + CuVS2510GPUVectorsFormat cagraAndBruteForceFormat = + new CuVS2510GPUVectorsFormat( + 32, 128, 64, CagraGraphBuildAlgo.NN_DESCENT, IndexType.CAGRA_AND_BRUTE_FORCE); IndexWriterConfig config = new IndexWriterConfig() - .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) + .setCodec(alwaysKnnVectorsFormat(cagraAndBruteForceFormat)) .setMaxBufferedDocs(maxBufferedDocs) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - try (IndexWriter writer = new IndexWriter(directory, config)) { - for (int seg = 0; seg < segmentCount; seg++) { - log.log(Level.FINE, "Creating segment " + (seg + 1) + "/" + segmentCount); - - // Randomize vector probability per segment - double vectorProbability = - 0.5 + (random().nextDouble() * 0.4); // 50-90% vectors per segment + int totalDocuments = numSegments * docsPerSegment; + int totalExpectedVectors = 0; + Set docIDsHavingVectors = new HashSet(); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int seg = 0; seg < numSegments; seg++) { + int segmentVectorCount = 0; for (int i = 0; i < docsPerSegment; i++) { int docId = seg * docsPerSegment + i; Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(docId), Field.Store.YES)); - doc.add(new NumericDocValuesField("segment", seg)); - doc.add(new NumericDocValuesField("position", i)); + doc.add(new StringField(ID_FIELD, String.valueOf(docId), Field.Store.YES)); - // Add vector based on segment's randomized probability - if (random().nextDouble() < vectorProbability) { - float[] vector = generateRandomVector(vectorDimension, random()); - doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.COSINE)); + // Randomly add vectors based on probability + if (random.nextDouble() < vectorProbability) { + float[] vector = dataProvider.getVectors(1)[0]; + doc.add(new KnnFloatVectorField(VECTOR_FIELD1, vector, EUCLIDEAN)); + docIDsHavingVectors.add(docId); + segmentVectorCount++; } - writer.addDocument(doc); } + writer.commit(); + totalExpectedVectors += segmentVectorCount; + + log.log( + Level.FINE, + "Created CAGRA + brute force segment " + + seg + + ": " + + docsPerSegment + + " documents, " + + segmentVectorCount + + " with vectors"); } log.log( Level.FINE, - "Created " + segmentCount + " segments with " + totalDocuments + " total documents"); + "Created " + + numSegments + + " CAGRA + brute force segments with " + + totalDocuments + + " total documents and " + + totalExpectedVectors + + " vectors"); - // Force merge all segments - long startTime = System.currentTimeMillis(); + // Force merge all CAGRA + brute force segments writer.forceMerge(1); - long mergeTime = System.currentTimeMillis() - startTime; - - log.log(Level.FINE, "Large scale merge completed in " + mergeTime + "ms"); + log.log(Level.FINE, "Forced merge of brute force segments completed"); } - // Verify the large merged index + // Verify the merged CAGRA + brute force index try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals("Should have exactly one segment after merge", 1, reader.leaves().size()); - - LeafReader leafReader = reader.leaves().get(0).reader(); + List leaves = reader.leaves(); + assertEquals("Should have exactly one segment after merge", 1, leaves.size()); + LeafReader leafReader = leaves.get(0).reader(); assertEquals("Total documents should match", totalDocuments, leafReader.maxDoc()); - // Test vector search performance - var vectorValues = leafReader.getFloatVectorValues("vector"); + // Count actual vectors in merged index + var vectorValues = leafReader.getFloatVectorValues(VECTOR_FIELD1); int actualVectorCount = vectorValues != null ? vectorValues.size() : 0; - if (actualVectorCount > 0) { - IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateRandomVector(vectorDimension, random()); - - // Randomize search parameters for large scale test - int searchK = - Math.min(20 + random().nextInt(31), Math.min(actualVectorCount, TOP_K_LIMIT)); // 20-50 + log.log( + Level.FINE, + "CAGRA + brute force merge results: Total documents: " + + totalDocuments + + ", Expected vectors: " + + totalExpectedVectors + + ", Actual vectors: " + + actualVectorCount); - long searchStart = System.currentTimeMillis(); - KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery("vector", queryVector, searchK); - TopDocs vectorResults = searcher.search(vectorQuery, searchK); - long searchTime = System.currentTimeMillis() - searchStart; + assertEquals("Vector count should match expected", totalExpectedVectors, actualVectorCount); - assertTrue("Should find vector results in large index", vectorResults.scoreDocs.length > 0); - log.log( - Level.FINE, - "Vector search in large index returned " - + vectorResults.scoreDocs.length - + " results out of " - + actualVectorCount - + " vectors in " - + searchTime - + "ms"); - } else { - log.log(Level.FINE, "No vectors in large merged index - skipping vector search"); - } + // Test CAGRA + brute force vector search (exact search) + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = dataProvider.getQueries(1)[0]; + int topK = dataProvider.getRandom(1, actualVectorCount); - log.log(Level.FINE, "Large scale merge verification completed successfully"); - } - } + KnnFloatVectorQuery vectorQuery = new KnnFloatVectorQuery(VECTOR_FIELD1, queryVector, topK); + TopDocs vectorResults = searcher.search(vectorQuery, topK); + int numResults = vectorResults.scoreDocs.length; + assertTrue("Should find some vector results in CAGRA + brute force index", numResults > 0); + assertTrue("Should not find more vectors than exist", numResults <= actualVectorCount); - /** Helper method to generate random vectors */ - private float[] generateRandomVector(int dimension, Random random) { - float[] vector = new float[dimension]; - for (int i = 0; i < dimension; i++) { - vector[i] = (float) random().nextGaussian(); - } - // Normalize the vector - float norm = 0.0f; - for (float v : vector) { - norm += v * v; - } - norm = (float) Math.sqrt(norm); - if (norm > 0) { - for (int i = 0; i < dimension; i++) { - vector[i] /= norm; + log.log( + Level.FINE, + "CAGRA + Brute force search found " + + numResults + + " results out of " + + actualVectorCount + + " available vectors"); + + assertEquals("Search should return exactly topK documents", topK, numResults); + // Verify all returned documents are valid + for (ScoreDoc scoreDoc : vectorResults.scoreDocs) { + Document doc = searcher.storedFields().document(scoreDoc.doc); + String docId = doc.get(ID_FIELD); + assertNotNull("Document should have valid ID", docId); + assertTrue("Score should be positive", scoreDoc.score > 0); + assertTrue( + "Document does not have a vector", + docIDsHavingVectors.contains(Integer.parseInt(docId))); } } - return vector; - } - - /** Helper method to generate random text strings for sorting */ - private String generateRandomText(Random random, int length) { - StringBuilder sb = new StringBuilder(length); - String chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; - for (int i = 0; i < length; i++) { - sb.append(chars.charAt(random().nextInt(chars.length()))); - } - return sb.toString(); } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index 05e37f84..39101dde 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -4,11 +4,29 @@ */ package com.nvidia.cuvs.lucene; +import static org.apache.lucene.tests.util.LuceneTestCase.newIndexWriterConfig; +import static org.apache.lucene.tests.util.LuceneTestCase.newTieredMergePolicy; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.TreeMap; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockTokenizer; +import org.apache.lucene.tests.index.RandomIndexWriter; public class TestUtils { - public static float[][] generateDataset(Random random, int size, int dimensions) { + private static final Logger log = Logger.getLogger(TestUtils.class.getName()); + + public static float[][] generateRandomVectors(Random random, int size, int dimensions) { float[][] dataset = new float[size][dimensions]; for (int i = 0; i < size; i++) { for (int j = 0; j < dimensions; j++) { @@ -18,22 +36,59 @@ public static float[][] generateDataset(Random random, int size, int dimensions) return dataset; } - public static float[] generateRandomVector(int dimensions, Random random) { - float[] vector = new float[dimensions]; - for (int i = 0; i < dimensions; i++) { - vector[i] = random.nextFloat() * 100; + public static List> generateExpectedTopK( + int topK, float[][] dataset, float[][] queries) { + List> neighborsResult = new ArrayList<>(); + int dimensions = dataset[0].length; + + for (float[] query : queries) { + Map distances = new TreeMap<>(); + for (int j = 0; j < dataset.length; j++) { + double distance = 0; + for (int k = 0; k < dimensions; k++) { + distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); + } + distances.put(j, (distance)); + } + + Map sorted = new TreeMap(distances); + log.log(Level.FINER, "EXPECTED: " + sorted); + + // Sort by distance and select the topK nearest neighbors + List neighbors = + distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .map(Map.Entry::getKey) + .toList(); + neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, dataset.length))); } - return vector; + + log.log(Level.FINE, "Expected results generated successfully."); + return neighborsResult; } - public static float[][] generateQueries(Random random, int dimensions, int numQueries) { - // Generate random query vectors - float[][] queries = new float[numQueries][dimensions]; - for (int i = 0; i < numQueries; i++) { - for (int j = 0; j < dimensions; j++) { - queries[i][j] = random.nextFloat() * 100; - } + public static RandomIndexWriter createWriter(Random random, Directory directory, Codec codec) + throws IOException { + return new RandomIndexWriter( + random, + directory, + newIndexWriterConfig(new MockAnalyzer(random, MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy())); + } + + public static IndexWriterConfig createWriterConfig(Random random, Codec codec) { + return newIndexWriterConfig(new MockAnalyzer(random, MockTokenizer.SIMPLE, true)) + .setCodec(codec) + .setMergePolicy(newTieredMergePolicy()); + } + + public static String generateRandomText(Random random, int length) { + StringBuilder sb = new StringBuilder(length); + String chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; + for (int i = 0; i < length; i++) { + sb.append(chars.charAt(random.nextInt(chars.length()))); } - return queries; + return sb.toString(); } } From c136164fe2b0a2a6aa2b71929d2c3172a59c3125 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Mon, 5 Jan 2026 18:07:29 -0500 Subject: [PATCH 4/4] Update license headers --- ci/release/update-version.sh | 2 +- .../com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java | 2 +- .../nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java | 2 +- .../src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java | 2 +- .../java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java | 2 +- .../java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java | 2 +- .../java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java | 2 +- .../java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java | 2 +- .../java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java | 2 +- .../java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java | 2 +- .../com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java | 2 +- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java | 2 +- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java | 2 +- .../nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java | 2 +- .../com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java | 2 +- .../java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java | 2 +- .../nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java | 2 +- .../nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java | 2 +- .../java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java | 2 +- .../com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestMerge.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestUtils.java | 2 +- 27 files changed, 27 insertions(+), 27 deletions(-) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index f98aefbc..40401044 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -1,6 +1,6 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 ######################## diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java index fd7a8808..8ac19071 100644 --- a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene.examples; diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java index 51f0aecd..8c889620 100644 --- a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene.examples; diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java index 901f4b21..a7a14e22 100644 --- a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene.examples; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 82ba6135..7bada638 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 2406a633..58603c34 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 9f912e76..a016feeb 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 66c5bd92..bcc60ce5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 3173c764..6a9daae7 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 57b76294..421b2baa 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index d73112fd..558a94a4 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 2bc62ded..4ed16c41 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index c754668e..e7669650 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java index 819bcea6..9b437567 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSW.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index 9406ea09..704f697a 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java index 71d4c4e7..af36b55f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWFallback.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java index 87cc5d8a..a1a48fc4 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWGaps.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java index a8725f95..bc96a0d7 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWRandomizedSearch.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java index 18fe43ed..31c9ddd9 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 02e36962..b8751553 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index baf43a4d..1074ed76 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index 95861a1c..eed0eaa4 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index dda6982f..867c22d7 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 67ba654a..25f7a331 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java b/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java index c4480e84..954bd52d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestDataProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index 13dac867..760fe4b8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index 39101dde..1217e492 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene;