From 1c0f64890be6ee4bbefca2bf4f27ee3eab989e11 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 13 May 2026 16:23:28 -0400 Subject: [PATCH 01/13] Add PyLucene integration support --- README.md | 34 ++++++++++ examples/pylucene_smoke.py | 68 +++++++++++++++++++ pom.xml | 7 +- src/main/assembly/jar-with-dependencies.xml | 22 ++++++ .../jar-with-pylucene-dependencies.xml | 25 +++++++ .../cuvs/lucene/CuVS2510GPUSearchCodec.java | 4 +- .../lucene/Lucene101AcceleratedHNSWCodec.java | 4 +- ...neAcceleratedHNSWBinaryQuantizedCodec.java | 4 +- ...ratedHNSWBinaryQuantizedVectorsFormat.java | 34 ++++++---- ...neAcceleratedHNSWScalarQuantizedCodec.java | 4 +- .../nvidia/cuvs/lucene/LuceneProvider.java | 24 ++++++- .../nvidia/cuvs/lucene/TestBackCompat.java | 28 ++++++++ 12 files changed, 231 insertions(+), 27 deletions(-) create mode 100644 examples/pylucene_smoke.py create mode 100644 src/main/assembly/jar-with-dependencies.xml create mode 100644 src/main/assembly/jar-with-pylucene-dependencies.xml diff --git a/README.md b/README.md index d88d3401..9b327247 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,40 @@ mvn clean compile package ``` The artifacts would be built and available in the target / folder. +### Using with PyLucene + +PyLucene embeds a JVM and starts it with the classpath passed to `lucene.initVM(...)`. +Because PyLucene's generated Python module only exposes the Java classes it was built +to wrap, use Lucene's service provider lookup to load `cuvs-lucene` codecs from +Python instead of importing `com.nvidia.cuvs.lucene` classes directly. + +Build the PyLucene sidecar jar: + +```sh +mvn clean package -DskipTests +``` + +Then add the generated PyLucene jar to `lucene.CLASSPATH` before starting the VM: + +```python +import os +import lucene + +cuvs_lucene_jar = "target/cuvs-lucene-26.06.0-jar-with-pylucene-dependencies.jar" +lucene.initVM(classpath=os.pathsep.join([lucene.CLASSPATH, cuvs_lucene_jar])) + +from org.apache.lucene.codecs import Codec + +codec = Codec.forName("Lucene101AcceleratedHNSWCodec") +``` + +Use the returned `codec` with `IndexWriterConfig.setCodec(codec)`. The +`jar-with-pylucene-dependencies` artifact includes `cuvs-lucene` and its non-Lucene +runtime dependencies while leaving Lucene itself to PyLucene's own classpath. This +avoids loading a second copy of Lucene classes into the embedded JVM. The regular +`jar-with-dependencies` artifact also merges `META-INF/services` entries and is +available for non-PyLucene Java applications that want a standalone jar. + ### Running Tests ```sh export LD_LIBRARY_PATH={ PATH TO YOUR LOCAL libcuvs_c.so }:$LD_LIBRARY_PATH && mvn clean test diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py new file mode 100644 index 00000000..68bf6145 --- /dev/null +++ b/examples/pylucene_smoke.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path + +import lucene +from lucene import JArray + + +REPO_ROOT = Path(__file__).resolve().parents[1] +CUVS_LUCENE_JAR = ( + REPO_ROOT / "target" / "cuvs-lucene-26.06.0-jar-with-pylucene-dependencies.jar" +) +INDEX_PATH = Path("/tmp/cuvs-lucene-pylucene-index") + + +def fvec(values): + return JArray("float")([float(value) for value in values]) + + +def main(): + lucene.initVM(classpath=os.pathsep.join([lucene.CLASSPATH, str(CUVS_LUCENE_JAR)])) + + from java.nio.file import Paths + from org.apache.lucene.codecs import Codec + from org.apache.lucene.document import Document, Field, KnnFloatVectorField, StringField + from org.apache.lucene.index import ( + DirectoryReader, + IndexWriter, + IndexWriterConfig, + VectorSimilarityFunction, + ) + from org.apache.lucene.search import IndexSearcher, KnnFloatVectorQuery + from org.apache.lucene.store import MMapDirectory + + codec = Codec.forName("Lucene101AcceleratedHNSWCodec") + + directory = MMapDirectory(Paths.get(str(INDEX_PATH))) + config = IndexWriterConfig() + config.setCodec(codec) + config.setUseCompoundFile(False) + + writer = IndexWriter(directory, config) + for doc_id, vector in [("a", [0.1, 0.2, 0.3]), ("b", [0.2, 0.1, 0.4])]: + doc = Document() + doc.add(StringField("id", doc_id, Field.Store.YES)) + doc.add( + KnnFloatVectorField( + "vector", fvec(vector), VectorSimilarityFunction.EUCLIDEAN + ) + ) + writer.addDocument(doc) + writer.close() + + reader = DirectoryReader.open(directory) + searcher = IndexSearcher(reader) + query = KnnFloatVectorQuery("vector", fvec([0.1, 0.2, 0.3]), 2) + + for hit in searcher.search(query, 2).scoreDocs: + print(searcher.doc(hit.doc).get("id"), hit.score) + + reader.close() + directory.close() + + +if __name__ == "__main__": + main() diff --git a/pom.xml b/pom.xml index 19479f10..8b18da20 100644 --- a/pom.xml +++ b/pom.xml @@ -112,9 +112,10 @@ maven-assembly-plugin 3.6.0 - - jar-with-dependencies - + + src/main/assembly/jar-with-dependencies.xml + src/main/assembly/jar-with-pylucene-dependencies.xml + diff --git a/src/main/assembly/jar-with-dependencies.xml b/src/main/assembly/jar-with-dependencies.xml new file mode 100644 index 00000000..3722a871 --- /dev/null +++ b/src/main/assembly/jar-with-dependencies.xml @@ -0,0 +1,22 @@ + + jar-with-dependencies + + jar + + false + + + / + true + true + runtime + + + + + metaInf-services + + + diff --git a/src/main/assembly/jar-with-pylucene-dependencies.xml b/src/main/assembly/jar-with-pylucene-dependencies.xml new file mode 100644 index 00000000..d243257b --- /dev/null +++ b/src/main/assembly/jar-with-pylucene-dependencies.xml @@ -0,0 +1,25 @@ + + jar-with-pylucene-dependencies + + jar + + false + + + / + true + true + runtime + + org.apache.lucene:* + + + + + + metaInf-services + + + diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 2de761a7..2567fc4f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -29,7 +29,7 @@ public class CuVS2510GPUSearchCodec extends FilterCodec { * @throws Exception */ public CuVS2510GPUSearchCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(new GPUSearchParams.Builder().build()); } @@ -53,7 +53,7 @@ public CuVS2510GPUSearchCodec(String name, Codec delegate) { * @throws Exception Exception raised when initializing the codec */ public CuVS2510GPUSearchCodec(GPUSearchParams params) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(params); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index b4c5a33d..4944b761 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -30,7 +30,7 @@ public class Lucene101AcceleratedHNSWCodec extends FilterCodec { * @throws Exception */ public Lucene101AcceleratedHNSWCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); } /** @@ -52,7 +52,7 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { */ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index f2c1aa37..6af2d3db 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java @@ -25,7 +25,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedCodec extends FilterCodec { private KnnVectorsFormat format; public LuceneAcceleratedHNSWBinaryQuantizedCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); } public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) { @@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) { public LuceneAcceleratedHNSWBinaryQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 0f8d9602..d70efc5e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -27,21 +27,26 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector private static final Logger log = Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class.getName()); - private static final LuceneProvider LUCENE102_PROVIDER; - private static final LuceneProvider LUCENE99_PROVIDER; - private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - static { + private static LuceneProvider luceneProvider(String version) { try { - LUCENE99_PROVIDER = LuceneProvider.getInstance("99"); - LUCENE102_PROVIDER = LuceneProvider.getInstance("102"); - FLAT_VECTORS_FORMAT = - LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + return LuceneProvider.getInstance(version); } catch (Exception e) { - throw new ExceptionInInitializerError(e.getMessage()); + throw new UnsupportedOperationException( + "Lucene" + version + " vector formats are not available in this runtime", e); + } + } + + private static FlatVectorsFormat flatVectorsFormat() { + try { + return luceneProvider("102") + .getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + } catch (Exception e) { + throw new UnsupportedOperationException( + "Binary quantized vectors require Lucene102 vector formats", e); } } @@ -70,7 +75,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + var flatWriter = flatVectorsFormat().fieldsWriter(state); if (isSupported()) { log.log( Level.FINE, @@ -85,8 +90,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException "GPU based indexing not supported, falling back to using the" + " Lucene102HnswBinaryQuantizedVectorsFormat"); KnnVectorsFormat fallbackFormat = - LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance( - acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); + luceneProvider("102") + .getLuceneHnswBinaryQuantizedVectorsFormatInstance( + acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); return fallbackFormat.fieldsWriter(state); } catch (Exception e) { throw new RuntimeException(e.getMessage()); @@ -100,8 +106,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { try { - return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance( - state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + return luceneProvider("99") + .getLuceneHnswVectorsReaderInstance(state, flatVectorsFormat().fieldsReader(state)); } catch (Exception e) { throw new RuntimeException(e.getMessage()); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java index 0c7736a0..4e49027c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java @@ -25,7 +25,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedCodec extends FilterCodec { private KnnVectorsFormat format; public LuceneAcceleratedHNSWScalarQuantizedCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); } public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) { @@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) { public LuceneAcceleratedHNSWScalarQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 7635e323..df758cc5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -8,7 +8,9 @@ import java.lang.invoke.VarHandle; import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -79,7 +81,7 @@ public class LuceneProvider { private static String luceneCodec = BASE + codecs + "LuceneCodec"; private static String luceneCodecFallback = BASE + fallbackCodecs + "LuceneCodec"; - private static LuceneProvider instance; + private static final Map INSTANCES = new HashMap<>(); private static MethodHandles.Lookup lookup = MethodHandles.lookup(); @@ -92,9 +94,12 @@ public class LuceneProvider { private Class scalarQuantizedVectorsFormat; private Class hnswScalarQuantizedVectorsFormat; - public static LuceneProvider getInstance(String version) throws ClassNotFoundException { + public static synchronized LuceneProvider getInstance(String version) + throws ClassNotFoundException { + LuceneProvider instance = INSTANCES.get(version); if (instance == null) { instance = new LuceneProvider(version); + INSTANCES.put(version, instance); } return instance; } @@ -173,6 +178,21 @@ public static Codec getCodec(String version) return (Codec) codecClassConstructor.newInstance(); } + public static Codec getDefaultDelegateCodec() { + for (String version : List.of("101", "99")) { + try { + return getCodec(version); + } catch (ReflectiveOperationException + | SecurityException + | IllegalArgumentException + | LinkageError e) { + log.log(Level.FINE, "Unable to load Lucene" + version + "Codec", e); + } + } + log.log(Level.FINE, "Falling back to the runtime default codec"); + return Codec.getDefault(); + } + public FlatVectorsFormat getLuceneFlatVectorsFormatInstance(FlatVectorsScorer scorer) throws Exception { try { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e660..1407e0ab 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -6,6 +6,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertTrue; import org.apache.lucene.codecs.Codec; @@ -38,4 +40,30 @@ public void testExistingComponents() throws Exception { assertEquals(provider.getStaticIntParam("VERSION_CURRENT"), 0); assertNotEquals(provider.getSimilarityFunctions().size(), 0); } + + @Test + public void testProviderCachesByVersion() throws Exception { + LuceneProvider provider99 = LuceneProvider.getInstance("99"); + LuceneProvider provider102 = LuceneProvider.getInstance("102"); + assertNotSame(provider99, provider102); + } + + @Test + public void testDefaultDelegateCodec() { + assertNotNull(LuceneProvider.getDefaultDelegateCodec()); + } + + @Test + public void testServiceLoadedCodecsCanBeInstantiated() { + String[] codecNames = { + "Lucene101AcceleratedHNSWCodec", + "CuVS2510GPUSearchCodec", + "Lucene101AcceleratedHNSWBinaryQuantizedCodec", + "Lucene101AcceleratedHNSWScalarQuantizedCodec" + }; + for (String codecName : codecNames) { + assertTrue(Codec.availableCodecs().contains(codecName)); + assertEquals(codecName, Codec.forName(codecName).getName()); + } + } } From e27a53b8bebd9b50c04762f0233d0cfc5c33f272 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 19 May 2026 16:31:55 -0400 Subject: [PATCH 02/13] Fix benchmark version marker --- bench/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bench/pom.xml b/bench/pom.xml index 8f7262af..d9a13443 100644 --- a/bench/pom.xml +++ b/bench/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.benchmarks cuvs-lucene-benchmarks - 26.02.0 + 26.06.0 jar cuvs-lucene-benchmarks From 3cab3846def7cc9c28eaac2607acb7ee8d8de599 Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Wed, 8 Jul 2026 15:28:04 -0700 Subject: [PATCH 03/13] Fix PyLucene sidecar packaging and update cuVS version --- README.md | 43 +++- ci/test_pylucene_smoke.sh | 207 +++++++++++++++++ examples/pylucene_smoke.py | 217 +++++++++++++++--- .../jar-with-pylucene-dependencies.xml | 20 +- .../org.apache.lucene.codecs.KnnVectorsFormat | 2 - 5 files changed, 427 insertions(+), 62 deletions(-) create mode 100755 ci/test_pylucene_smoke.sh diff --git a/README.md b/README.md index 9b327247..84d5f6ed 100644 --- a/README.md +++ b/README.md @@ -31,14 +31,28 @@ Build the PyLucene sidecar jar: mvn clean package -DskipTests ``` -Then add the generated PyLucene jar to `lucene.CLASSPATH` before starting the VM: +Then start PyLucene with the base `cuvs-java` jar, the generated PyLucene +sidecar jar, and PyLucene's own Lucene classpath: ```python import os +from pathlib import Path + import lucene -cuvs_lucene_jar = "target/cuvs-lucene-26.06.0-jar-with-pylucene-dependencies.jar" -lucene.initVM(classpath=os.pathsep.join([lucene.CLASSPATH, cuvs_lucene_jar])) +cuvs_java_jar = Path(os.environ["CUVS_LUCENE_CUVS_JAVA_JAR"]) +cuvs_lucene_jar = next( + Path("target").glob("cuvs-lucene-*-jar-with-pylucene-dependencies.jar") +) +lucene.initVM( + classpath=os.pathsep.join( + [str(cuvs_java_jar), str(cuvs_lucene_jar), lucene.CLASSPATH] + ), + vmargs=[ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + ], +) from org.apache.lucene.codecs import Codec @@ -46,11 +60,24 @@ codec = Codec.forName("Lucene101AcceleratedHNSWCodec") ``` Use the returned `codec` with `IndexWriterConfig.setCodec(codec)`. The -`jar-with-pylucene-dependencies` artifact includes `cuvs-lucene` and its non-Lucene -runtime dependencies while leaving Lucene itself to PyLucene's own classpath. This -avoids loading a second copy of Lucene classes into the embedded JVM. The regular -`jar-with-dependencies` artifact also merges `META-INF/services` entries and is -available for non-PyLucene Java applications that want a standalone jar. +`jar-with-pylucene-dependencies` artifact includes only `cuvs-lucene` classes and +service descriptors. PyLucene must provide Lucene classes, and the base +multi-release `cuvs-java` jar must be present separately on the JVM classpath. Do +not use a native classifier `cuvs-java` jar here unless you also want to rely on +its embedded native libraries; the base jar uses native libraries from +`LD_LIBRARY_PATH`/`java.library.path`. + +To run an end-to-end smoke test against a local PyLucene environment: + +```sh +./ci/test_pylucene_smoke.sh +``` + +To run the same smoke through the GPU search codec with 2,000 documents: + +```sh +./ci/test_pylucene_smoke.sh --gpu-e2e +``` ### Running Tests ```sh diff --git a/ci/test_pylucene_smoke.sh b/ci/test_pylucene_smoke.sh new file mode 100755 index 00000000..db3f2a76 --- /dev/null +++ b/ci/test_pylucene_smoke.sh @@ -0,0 +1,207 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "${REPO_ROOT}" + +MVN_BIN="${MVN:-mvn}" +PYTHON_BIN="${PYTHON:-python3}" +SKIP_BUILD=0 +GPU_E2E=0 + +for arg in "$@"; do + case "${arg}" in + --gpu-e2e) + GPU_E2E=1 + ;; + --no-build) + SKIP_BUILD=1 + ;; + -h|--help) + echo "Usage: $0 [--no-build] [--gpu-e2e]" + echo + echo "Builds and checks the PyLucene sidecar jar, then runs examples/pylucene_smoke.py." + echo "Set CUVS_LUCENE_PYLUCENE_JAR to test an existing sidecar jar." + echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar if it is not in ~/.m2." + echo "Set PYTHON or MVN to override the Python or Maven executable." + echo + echo "--gpu-e2e indexes 2,000 rows through CuVS2510GPUSearchCodec and requires cuVS native support." + exit 0 + ;; + *) + echo "Unknown argument: ${arg}" >&2 + exit 2 + ;; + esac +done + +require_command() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "Required command not found: $1" >&2 + exit 127 + fi +} + +require_command "${PYTHON_BIN}" +require_command jar + +if [[ "${SKIP_BUILD}" -eq 0 && -z "${CUVS_LUCENE_PYLUCENE_JAR:-}" ]]; then + require_command "${MVN_BIN}" + "${MVN_BIN}" clean package -DskipTests +fi + +project_version="$( + sed -n 's/.*CUVS_LUCENE#VERSION_UPDATE_MARKER_START-->\([^<]*\)<\/version>.*/\1/p' pom.xml \ + | head -n 1 +)" +if [[ -z "${project_version}" ]]; then + echo "Unable to determine project version from pom.xml" >&2 + exit 1 +fi + +if [[ -n "${CUVS_LUCENE_PYLUCENE_JAR:-}" ]]; then + sidecar_jar="${CUVS_LUCENE_PYLUCENE_JAR}" +else + sidecar_jar="target/cuvs-lucene-${project_version}-jar-with-pylucene-dependencies.jar" +fi + +if [[ ! -f "${sidecar_jar}" ]]; then + echo "PyLucene sidecar jar not found: ${sidecar_jar}" >&2 + echo "Run without --no-build, or set CUVS_LUCENE_PYLUCENE_JAR to an existing jar." >&2 + exit 1 +fi + +case "${sidecar_jar}" in + /*) + sidecar_jar_abs="${sidecar_jar}" + ;; + *) + sidecar_jar_abs="${REPO_ROOT}/${sidecar_jar}" + ;; +esac + +find_cuvs_java_jar() { + if [[ -n "${CUVS_LUCENE_CUVS_JAVA_JAR:-}" ]]; then + printf '%s\n' "${CUVS_LUCENE_CUVS_JAVA_JAR}" + return + fi + + local m2_base="${HOME}/.m2/repository/com/nvidia/cuvs/cuvs-java/${project_version}" + local jar="${m2_base}/cuvs-java-${project_version}.jar" + if [[ -f "${jar}" ]]; then + printf '%s\n' "${jar}" + return + fi + + local m2_repo="${HOME}/.m2/repository/com/nvidia/cuvs/cuvs-java" + if [[ ! -d "${m2_repo}" ]]; then + return + fi + + find "${m2_repo}" \ + -type f \ + -name 'cuvs-java-*.jar' \ + ! -name '*sources*' \ + ! -name '*javadoc*' \ + ! -name '*x86_64*' \ + | sort -V \ + | tail -n 1 +} + +cuvs_java_jar="$(find_cuvs_java_jar)" +if [[ -z "${cuvs_java_jar}" || ! -f "${cuvs_java_jar}" ]]; then + echo "Base cuvs-java jar not found." >&2 + echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar, not a native classifier jar." >&2 + exit 1 +fi + +entries_file="$(mktemp)" +services_dir="$(mktemp -d)" +trap 'rm -f "${entries_file}"; rm -rf "${services_dir}"' EXIT + +jar tf "${sidecar_jar_abs}" >"${entries_file}" + +for service in \ + "META-INF/services/org.apache.lucene.codecs.Codec" \ + "META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat"; do + grep -qx "${service}" "${entries_file}" || { + echo "Missing service descriptor in ${sidecar_jar}: ${service}" >&2 + exit 1 + } +done + +grep -qx "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.class" "${entries_file}" || { + echo "Missing cuvs-lucene codec classes in ${sidecar_jar}" >&2 + exit 1 +} + +if grep -q "^org/apache/lucene/" "${entries_file}"; then + echo "${sidecar_jar} contains org.apache.lucene classes; PyLucene must provide Lucene." >&2 + exit 1 +fi + +if grep -q "^com/nvidia/cuvs/" "${entries_file}" \ + && grep "^com/nvidia/cuvs/" "${entries_file}" | grep -vq "^com/nvidia/cuvs/lucene/"; then + echo "${sidecar_jar} contains flattened cuvs-java classes; use the base cuvs-java jar separately." >&2 + exit 1 +fi + +if grep -q "^META-INF/versions/.*/com/nvidia/cuvs/" "${entries_file}"; then + echo "${sidecar_jar} contains flattened multi-release cuvs-java classes." >&2 + exit 1 +fi + +extra_lucene_services="$( + grep "^META-INF/services/org.apache.lucene." "${entries_file}" \ + | grep -v "^META-INF/services/org.apache.lucene.codecs.Codec$" \ + | grep -v "^META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat$" || true +)" +if [[ -n "${extra_lucene_services}" ]]; then + echo "${sidecar_jar} contains unexpected Lucene service descriptors:" >&2 + echo "${extra_lucene_services}" >&2 + exit 1 +fi + +( + cd "${services_dir}" + jar xf \ + "${sidecar_jar_abs}" \ + META-INF/services/org.apache.lucene.codecs.Codec \ + META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +) + +for descriptor in \ + "${services_dir}/META-INF/services/org.apache.lucene.codecs.Codec" \ + "${services_dir}/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat"; do + if grep -q "^org\\.apache\\.lucene\\." "${descriptor}"; then + echo "${descriptor#${services_dir}/} advertises Lucene-owned providers." >&2 + exit 1 + fi +done + +"${PYTHON_BIN}" -c "import lucene" >/dev/null 2>&1 || { + echo "Python cannot import PyLucene's lucene module." >&2 + echo "Activate or install a PyLucene environment compatible with this project's Lucene version." >&2 + exit 1 +} + +smoke_env=( + "CUVS_LUCENE_PYLUCENE_JAR=${sidecar_jar_abs}" + "CUVS_LUCENE_CUVS_JAVA_JAR=${cuvs_java_jar}" +) + +if [[ "${GPU_E2E}" -eq 1 ]]; then + smoke_env+=( + "CUVS_LUCENE_PYLUCENE_CODEC=CuVS2510GPUSearchCodec" + "CUVS_LUCENE_PYLUCENE_ROWS=2000" + "CUVS_LUCENE_PYLUCENE_DIMS=32" + "CUVS_LUCENE_PYLUCENE_TOPK=20" + "CUVS_LUCENE_EXPECT_CUVS_FILES=1" + ) +fi + +env "${smoke_env[@]}" "${PYTHON_BIN}" examples/pylucene_smoke.py diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py index 68bf6145..decdd11a 100644 --- a/examples/pylucene_smoke.py +++ b/examples/pylucene_smoke.py @@ -2,26 +2,118 @@ # SPDX-License-Identifier: Apache-2.0 import os +import sys +import tempfile from pathlib import Path import lucene -from lucene import JArray REPO_ROOT = Path(__file__).resolve().parents[1] -CUVS_LUCENE_JAR = ( - REPO_ROOT / "target" / "cuvs-lucene-26.06.0-jar-with-pylucene-dependencies.jar" -) -INDEX_PATH = Path("/tmp/cuvs-lucene-pylucene-index") +DEFAULT_CODEC = "Lucene101AcceleratedHNSWCodec" +ID_FIELD = "id" +VECTOR_FIELD = "vector" -def fvec(values): - return JArray("float")([float(value) for value in values]) +def find_cuvs_lucene_jar(): + configured = os.environ.get("CUVS_LUCENE_PYLUCENE_JAR") or os.environ.get( + "CUVS_LUCENE_JAR" + ) + if configured: + jar = Path(configured) + if not jar.exists(): + raise FileNotFoundError(f"Configured cuvs-lucene jar does not exist: {jar}") + return jar + + jars = sorted( + (REPO_ROOT / "target").glob( + "cuvs-lucene-*-jar-with-pylucene-dependencies.jar" + ) + ) + if not jars: + raise FileNotFoundError( + "No PyLucene sidecar jar found under target/. " + "Run `mvn clean package -DskipTests` first." + ) + return jars[-1] + + +def find_cuvs_java_jar(): + configured = os.environ.get("CUVS_LUCENE_CUVS_JAVA_JAR") or os.environ.get( + "CUVS_JAVA_JAR" + ) + if configured: + jar = Path(configured) + if not jar.exists(): + raise FileNotFoundError(f"Configured cuvs-java jar does not exist: {jar}") + return jar + + m2_repo = Path.home() / ".m2" / "repository" / "com" / "nvidia" / "cuvs" / "cuvs-java" + if not m2_repo.exists(): + raise FileNotFoundError( + "Unable to find cuvs-java in ~/.m2. Set CUVS_LUCENE_CUVS_JAVA_JAR " + "to the base cuvs-java jar, not a native classifier jar." + ) + + def is_base_cuvs_java_jar(jar): + return ( + jar.name.startswith("cuvs-java-") + and jar.name.endswith(".jar") + and "-x86_64-" not in jar.name + and "-sources" not in jar.name + and "-javadoc" not in jar.name + ) + + jars = sorted( + jar for jar in m2_repo.glob("*/*.jar") if is_base_cuvs_java_jar(jar) + ) + if not jars: + raise FileNotFoundError( + "Unable to find the base cuvs-java jar in ~/.m2. Set " + "CUVS_LUCENE_CUVS_JAVA_JAR explicitly." + ) + return jars[-1] + + +def vector_for(doc_id, dims): + x = ((doc_id + 1) * 2654435761) & 0xFFFFFFFF + values = [] + for dim in range(dims): + x = (1664525 * x + 1013904223 + dim * 17) & 0xFFFFFFFF + values.append((x / 4294967295.0) * 2.0 - 1.0) + return values + + +def fvec(jarray, values): + return jarray("float")([float(value) for value in values]) + + +def init_vm(cuvs_java_jar, cuvs_lucene_jar): + java_library_path = os.environ.get("JAVA_LIBRARY_PATH") or os.environ.get( + "LD_LIBRARY_PATH" + ) + vmargs = [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + ] + if java_library_path: + vmargs.append(f"-Djava.library.path={java_library_path}") + + lucene.initVM( + classpath=os.pathsep.join( + [str(cuvs_java_jar), str(cuvs_lucene_jar), lucene.CLASSPATH] + ), + vmargs=vmargs, + ) def main(): - lucene.initVM(classpath=os.pathsep.join([lucene.CLASSPATH, str(CUVS_LUCENE_JAR)])) + cuvs_lucene_jar = find_cuvs_lucene_jar() + cuvs_java_jar = find_cuvs_java_jar() + init_vm(cuvs_java_jar, cuvs_lucene_jar) + from lucene import JArray + from java.lang import Class from java.nio.file import Paths from org.apache.lucene.codecs import Codec from org.apache.lucene.document import Document, Field, KnnFloatVectorField, StringField @@ -32,37 +124,88 @@ def main(): VectorSimilarityFunction, ) from org.apache.lucene.search import IndexSearcher, KnnFloatVectorQuery - from org.apache.lucene.store import MMapDirectory - - codec = Codec.forName("Lucene101AcceleratedHNSWCodec") - - directory = MMapDirectory(Paths.get(str(INDEX_PATH))) - config = IndexWriterConfig() - config.setCodec(codec) - config.setUseCompoundFile(False) - - writer = IndexWriter(directory, config) - for doc_id, vector in [("a", [0.1, 0.2, 0.3]), ("b", [0.2, 0.1, 0.4])]: - doc = Document() - doc.add(StringField("id", doc_id, Field.Store.YES)) - doc.add( - KnnFloatVectorField( - "vector", fvec(vector), VectorSimilarityFunction.EUCLIDEAN - ) + from org.apache.lucene.store import FSDirectory + + # Forces an early, clear failure if cuvs-java was flattened or omitted. + Class.forName("com.nvidia.cuvs.spi.JDKProvider") + + codec_name = os.environ.get("CUVS_LUCENE_PYLUCENE_CODEC", DEFAULT_CODEC) + row_count = int(os.environ.get("CUVS_LUCENE_PYLUCENE_ROWS", "2")) + dims = int(os.environ.get("CUVS_LUCENE_PYLUCENE_DIMS", "3")) + top_k = int(os.environ.get("CUVS_LUCENE_PYLUCENE_TOPK", "2")) + expect_cuvs_files = os.environ.get("CUVS_LUCENE_EXPECT_CUVS_FILES") == "1" + + available_codecs = Codec.availableCodecs() + if not available_codecs.contains(codec_name): + raise AssertionError( + f"{codec_name} was not advertised by Lucene SPI. " + f"Available codecs: {available_codecs}" ) - writer.addDocument(doc) - writer.close() - reader = DirectoryReader.open(directory) - searcher = IndexSearcher(reader) - query = KnnFloatVectorQuery("vector", fvec([0.1, 0.2, 0.3]), 2) - - for hit in searcher.search(query, 2).scoreDocs: - print(searcher.doc(hit.doc).get("id"), hit.score) - - reader.close() - directory.close() + codec = Codec.forName(codec_name) + if codec.getName() != codec_name: + raise AssertionError(f"Expected codec {codec_name}, got {codec.getName()}") + + with tempfile.TemporaryDirectory(prefix="cuvs-lucene-pylucene-") as index_path: + directory = FSDirectory.open(Paths.get(index_path)) + config = IndexWriterConfig() + config.setCodec(codec) + config.setUseCompoundFile(False) + + writer = IndexWriter(directory, config) + try: + for doc_id in range(row_count): + doc = Document() + doc.add(StringField(ID_FIELD, f"doc-{doc_id}", Field.Store.YES)) + doc.add( + KnnFloatVectorField( + VECTOR_FIELD, + fvec(JArray, vector_for(doc_id, dims)), + VectorSimilarityFunction.EUCLIDEAN, + ) + ) + writer.addDocument(doc) + writer.commit() + finally: + writer.close() + + if expect_cuvs_files: + index_files = sorted(path.name for path in Path(index_path).iterdir()) + if not any(name.endswith(".vcag") for name in index_files): + raise AssertionError(f"No cuVS .vcag file found: {index_files}") + if not any(name.endswith(".vemc") for name in index_files): + raise AssertionError(f"No cuVS .vemc file found: {index_files}") + + reader = DirectoryReader.open(directory) + try: + searcher = IndexSearcher(reader) + stored_fields = searcher.storedFields() + query_ids = sorted({0, row_count // 2, row_count - 1}) + for query_id in query_ids: + query = KnnFloatVectorQuery( + VECTOR_FIELD, fvec(JArray, vector_for(query_id, dims)), top_k + ) + hits = searcher.search(query, top_k).scoreDocs + ids = [stored_fields.document(hit.doc).get(ID_FIELD) for hit in hits] + expected = f"doc-{query_id}" + if expected not in ids: + raise AssertionError( + f"Expected {expected} in top {top_k}, got {ids}" + ) + finally: + reader.close() + directory.close() + + print( + "PyLucene smoke test passed: " + f"loaded {codec_name} from {cuvs_lucene_jar.name}, " + f"used {cuvs_java_jar.name}, and searched {row_count} docs" + ) if __name__ == "__main__": - main() + try: + main() + except Exception as exc: + print(f"PyLucene smoke test failed: {exc}", file=sys.stderr) + raise diff --git a/src/main/assembly/jar-with-pylucene-dependencies.xml b/src/main/assembly/jar-with-pylucene-dependencies.xml index d243257b..de01eda5 100644 --- a/src/main/assembly/jar-with-pylucene-dependencies.xml +++ b/src/main/assembly/jar-with-pylucene-dependencies.xml @@ -6,20 +6,10 @@ jar false - - + + + ${project.build.outputDirectory} / - true - true - runtime - - org.apache.lucene:* - - - - - - metaInf-services - - + + diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 6625ac72..1f9ceeda 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,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat -org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat From 7ec17fe8a6c22107866e15ef7e26d44abdf65966 Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Thu, 16 Jul 2026 08:15:08 -0700 Subject: [PATCH 04/13] Expand PyLucene smoke coverage Convert the PyLucene smoke runner to a pytest-backed case matrix covering GPU CAGRA search, CAGRA-built HNSW, and forced CPU HNSW fallback across segment and force-merge topologies. Add named one-layer and three-layer CAGRA-to-HNSW codecs plus writer-path telemetry so the e2e suite can assert which CPU/GPU path was exercised. Tighten sidecar and Lucene delegate-codec validation to avoid misleading expected probe warnings while still failing on unsupported delegate codecs. --- README.md | 25 +- ci/test_pylucene_smoke.sh | 136 ++- examples/pylucene_smoke.py | 898 ++++++++++++++++-- examples/test_pylucene_smoke.py | 29 + .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 3 + ...ucene101AcceleratedHNSWBaseLayerCodec.java | 31 + .../lucene/Lucene101AcceleratedHNSWCodec.java | 14 +- ...cene101AcceleratedHNSWMultiLayerCodec.java | 31 + .../Lucene99AcceleratedHNSWVectorsFormat.java | 35 +- .../nvidia/cuvs/lucene/LuceneProvider.java | 29 +- .../ThreadLocalCuVSResourcesProvider.java | 12 +- .../services/org.apache.lucene.codecs.Codec | 2 + .../nvidia/cuvs/lucene/TestBackCompat.java | 20 +- 13 files changed, 1163 insertions(+), 102 deletions(-) create mode 100644 examples/test_pylucene_smoke.py create mode 100644 src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.java create mode 100644 src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.java diff --git a/README.md b/README.md index fe5fdb3b..82bc3597 100644 --- a/README.md +++ b/README.md @@ -69,18 +69,39 @@ not use a native classifier `cuvs-java` jar here unless you also want to rely on its embedded native libraries; the base jar uses native libraries from `LD_LIBRARY_PATH`/`java.library.path`. -To run an end-to-end smoke test against a local PyLucene environment: +To run the PyLucene pytest smoke suite against a local PyLucene environment: ```sh ./ci/test_pylucene_smoke.sh ``` -To run the same smoke through the GPU search codec with 2,000 documents: +To run an expanded GPU end-to-end pytest suite through CPU HNSW, +CAGRA-to-HNSW, and CAGRA search paths: ```sh ./ci/test_pylucene_smoke.sh --gpu-e2e ``` +The expanded suite runs the `gpu-basic`, `gpu-segments`, `cpu-hnsw`, and +`cagra-hnsw` case groups. The basic cases cover `hnsw`, `cagra`, `hnsw-single`, +and `cagra-single`. The segment cases cover 1-segment indexes, 10-segment +indexes, 10 segments force-merged to 1, and 100 segments force-merged to 10 for +both HNSW and CAGRA. The CPU HNSW cases force the accelerated HNSW codec through +its Lucene CPU fallback path in the same run, including 10 segments force-merged +to 1 and 100 segments force-merged to 10. The CAGRA-to-HNSW cases explicitly +cover one-layer and three-layer HNSW graphs built from CAGRA with NN_DESCENT, +`graphDegree=32`, and `intermediateGraphDegree=64`. The base matrix uses 2,000 +documents and 32 dimensions; high-segment cases use at least 257 rows per +segment to avoid expected cuVS graph-degree clamps on tiny per-segment datasets. +The suite checks Lucene SPI discovery, sidecar packaging, index file suffixes +(`.vex`/`.vem` for HNSW and `.vcag`/`.vemc` for CAGRA), indexed vector metadata, +unfiltered KNN, filtered KNN, missing-vector documents, deletions, and force +merge behavior. To run a subset or resize the test: + +```sh +./ci/test_pylucene_smoke.sh --gpu-e2e --cases=gpu-segments --rows=5000 --dims=64 --topk=20 +``` + ### Running Tests ```sh diff --git a/ci/test_pylucene_smoke.sh b/ci/test_pylucene_smoke.sh index baec802a..b6953b97 100755 --- a/ci/test_pylucene_smoke.sh +++ b/ci/test_pylucene_smoke.sh @@ -12,6 +12,10 @@ MVN_BIN="${MVN:-mvn}" PYTHON_BIN="${PYTHON:-python3}" SKIP_BUILD=0 GPU_E2E=0 +PYLUCENE_CASES="${CUVS_LUCENE_PYLUCENE_CASES:-}" +PYLUCENE_ROWS="${CUVS_LUCENE_PYLUCENE_ROWS:-}" +PYLUCENE_DIMS="${CUVS_LUCENE_PYLUCENE_DIMS:-}" +PYLUCENE_TOPK="${CUVS_LUCENE_PYLUCENE_TOPK:-}" for arg in "$@"; do case "${arg}" in @@ -21,15 +25,38 @@ for arg in "$@"; do --no-build) SKIP_BUILD=1 ;; + --cases=*) + PYLUCENE_CASES="${arg#--cases=}" + ;; + --rows=*) + PYLUCENE_ROWS="${arg#--rows=}" + ;; + --dims=*) + PYLUCENE_DIMS="${arg#--dims=}" + ;; + --topk=*) + PYLUCENE_TOPK="${arg#--topk=}" + ;; -h|--help) - echo "Usage: $0 [--no-build] [--gpu-e2e]" + echo "Usage: $0 [--no-build] [--gpu-e2e] [--cases=CASE[,CASE...]] [--rows=N] [--dims=N] [--topk=N]" echo - echo "Builds and checks the PyLucene sidecar jar, then runs examples/pylucene_smoke.py." + echo "Builds and checks the PyLucene sidecar jar, then runs the PyLucene pytest suite." echo "Set CUVS_LUCENE_PYLUCENE_JAR to test an existing sidecar jar." echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar if it is not in ~/.m2." echo "Set PYTHON or MVN to override the Python or Maven executable." echo - echo "--gpu-e2e indexes 2,000 rows through CuVS2510GPUSearchCodec and requires cuVS native support." + echo "Case groups: gpu-basic, gpu-segments, cpu-hnsw, cagra-hnsw, algorithm-matrix, all." + echo "Core cases: smoke, hnsw, cagra, hnsw-single, cagra-single." + echo "Segment cases: hnsw-1seg, cagra-1seg, hnsw-10seg, cagra-10seg," + echo " hnsw-10seg-force-1, cagra-10seg-force-1," + echo " hnsw-100seg-force-10, cagra-100seg-force-10." + echo "CPU HNSW cases: hnsw-cpu, hnsw-cpu-single, hnsw-cpu-1seg," + echo " hnsw-cpu-10seg, hnsw-cpu-10seg-force-1," + echo " hnsw-cpu-100seg-force-10." + echo "CAGRA-to-HNSW cases: cagra-hnsw-1layer, cagra-hnsw-3layer." + echo "--gpu-e2e defaults to cases=all, base rows=2000, dims=32, topk=20." + echo "High-segment cases use at least 257 rows per segment to avoid expected cuVS graph-degree clamps." + echo "--gpu-e2e requires a PyLucene environment plus cuVS native support on a GPU machine." exit 0 ;; *) @@ -134,24 +161,47 @@ for service in \ } done -grep -qx "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.class" "${entries_file}" || { - echo "Missing cuvs-lucene codec classes in ${sidecar_jar}" >&2 - exit 1 -} +for class_file in \ + "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.class" \ + "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.class" \ + "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.class" \ + "com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.class" \ + "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.class" \ + "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.class"; do + grep -qx "${class_file}" "${entries_file}" || { + echo "Missing cuvs-lucene class in ${sidecar_jar}: ${class_file}" >&2 + exit 1 + } +done -if grep -q "^org/apache/lucene/" "${entries_file}"; then +lucene_entries="$( + grep "^org/apache/lucene/" "${entries_file}" \ + | grep -v '/$' || true +)" +if [[ -n "${lucene_entries}" ]]; then echo "${sidecar_jar} contains org.apache.lucene classes; PyLucene must provide Lucene." >&2 + echo "${lucene_entries}" >&2 exit 1 fi -if grep -q "^com/nvidia/cuvs/" "${entries_file}" \ - && grep "^com/nvidia/cuvs/" "${entries_file}" | grep -vq "^com/nvidia/cuvs/lucene/"; then +flattened_cuvs_entries="$( + grep "^com/nvidia/cuvs/" "${entries_file}" \ + | grep -v '/$' \ + | grep -v "^com/nvidia/cuvs/lucene/" || true +)" +if [[ -n "${flattened_cuvs_entries}" ]]; then echo "${sidecar_jar} contains flattened cuvs-java classes; use the base cuvs-java jar separately." >&2 + echo "${flattened_cuvs_entries}" >&2 exit 1 fi -if grep -q "^META-INF/versions/.*/com/nvidia/cuvs/" "${entries_file}"; then +flattened_multi_release_entries="$( + grep "^META-INF/versions/.*/com/nvidia/cuvs/" "${entries_file}" \ + | grep -v '/$' || true +)" +if [[ -n "${flattened_multi_release_entries}" ]]; then echo "${sidecar_jar} contains flattened multi-release cuvs-java classes." >&2 + echo "${flattened_multi_release_entries}" >&2 exit 1 fi @@ -174,21 +224,52 @@ fi META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat ) -for descriptor in \ - "${services_dir}/META-INF/services/org.apache.lucene.codecs.Codec" \ - "${services_dir}/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat"; do +codec_descriptor="${services_dir}/META-INF/services/org.apache.lucene.codecs.Codec" +format_descriptor="${services_dir}/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat" + +for descriptor in "${codec_descriptor}" "${format_descriptor}"; do if grep -q "^org\\.apache\\.lucene\\." "${descriptor}"; then echo "${descriptor#"${services_dir}"/} advertises Lucene-owned providers." >&2 exit 1 fi done +for provider in \ + "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec" \ + "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWBaseLayerCodec" \ + "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWMultiLayerCodec" \ + "com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedCodec" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedCodec"; do + grep -qx "${provider}" "${codec_descriptor}" || { + echo "Codec service descriptor missing provider: ${provider}" >&2 + exit 1 + } +done + +for provider in \ + "com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat" \ + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedVectorsFormat"; do + grep -qx "${provider}" "${format_descriptor}" || { + echo "KnnVectorsFormat service descriptor missing provider: ${provider}" >&2 + exit 1 + } +done + "${PYTHON_BIN}" -c "import lucene" >/dev/null 2>&1 || { echo "Python cannot import PyLucene's lucene module." >&2 echo "Activate or install a PyLucene environment compatible with this project's Lucene version." >&2 exit 1 } +"${PYTHON_BIN}" -m pytest --version >/dev/null 2>&1 || { + echo "Python cannot run pytest." >&2 + echo "Install pytest in the active PyLucene environment." >&2 + exit 1 +} + smoke_env=( "CUVS_LUCENE_PYLUCENE_JAR=${sidecar_jar_abs}" "CUVS_LUCENE_CUVS_JAVA_JAR=${cuvs_java_jar}" @@ -196,12 +277,27 @@ smoke_env=( if [[ "${GPU_E2E}" -eq 1 ]]; then smoke_env+=( - "CUVS_LUCENE_PYLUCENE_CODEC=CuVS2510GPUSearchCodec" - "CUVS_LUCENE_PYLUCENE_ROWS=2000" - "CUVS_LUCENE_PYLUCENE_DIMS=32" - "CUVS_LUCENE_PYLUCENE_TOPK=20" - "CUVS_LUCENE_EXPECT_CUVS_FILES=1" + "CUVS_LUCENE_PYLUCENE_CASES=${PYLUCENE_CASES:-all}" + "CUVS_LUCENE_PYLUCENE_ROWS=${PYLUCENE_ROWS:-2000}" + "CUVS_LUCENE_PYLUCENE_DIMS=${PYLUCENE_DIMS:-32}" + "CUVS_LUCENE_PYLUCENE_TOPK=${PYLUCENE_TOPK:-20}" + "CUVS_LUCENE_REQUIRE_CUVS=1" + "CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}" ) +else + if [[ -n "${PYLUCENE_CASES}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_CASES=${PYLUCENE_CASES}") + fi + if [[ -n "${PYLUCENE_ROWS}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_ROWS=${PYLUCENE_ROWS}") + fi + if [[ -n "${PYLUCENE_DIMS}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_DIMS=${PYLUCENE_DIMS}") + fi + if [[ -n "${PYLUCENE_TOPK}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_TOPK=${PYLUCENE_TOPK}") + fi + smoke_env+=("CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}") fi -env "${smoke_env[@]}" "${PYTHON_BIN}" examples/pylucene_smoke.py +env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s examples/test_pylucene_smoke.py diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py index 180cf92a..cbc32e1b 100644 --- a/examples/pylucene_smoke.py +++ b/examples/pylucene_smoke.py @@ -2,17 +2,141 @@ # SPDX-License-Identifier: Apache-2.0 import os -import sys import tempfile +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass from pathlib import Path import lucene REPO_ROOT = Path(__file__).resolve().parents[1] -DEFAULT_CODEC = "Lucene101AcceleratedHNSWCodec" +HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +CAGRA_HNSW_BASE_LAYER_CODEC = "Lucene101AcceleratedHNSWBaseLayerCodec" +CAGRA_HNSW_MULTI_LAYER_CODEC = "Lucene101AcceleratedHNSWMultiLayerCodec" +CAGRA_CODEC = "CuVS2510GPUSearchCodec" +HNSW_SEARCH_CODECS = ( + HNSW_CODEC, + CAGRA_HNSW_BASE_LAYER_CODEC, + CAGRA_HNSW_MULTI_LAYER_CODEC, +) +EXPECTED_CODECS = ( + HNSW_CODEC, + CAGRA_HNSW_BASE_LAYER_CODEC, + CAGRA_HNSW_MULTI_LAYER_CODEC, + CAGRA_CODEC, + "Lucene101AcceleratedHNSWBinaryQuantizedCodec", + "Lucene101AcceleratedHNSWScalarQuantizedCodec", +) +BASIC_GPU_CASES = ("hnsw", "cagra", "hnsw-single", "cagra-single") +SEGMENT_GPU_CASES = ( + "hnsw-1seg", + "cagra-1seg", + "hnsw-10seg", + "cagra-10seg", + "hnsw-10seg-force-1", + "cagra-10seg-force-1", + "hnsw-100seg-force-10", + "cagra-100seg-force-10", +) +CPU_HNSW_CASES = ( + "hnsw-cpu", + "hnsw-cpu-single", + "hnsw-cpu-1seg", + "hnsw-cpu-10seg", + "hnsw-cpu-10seg-force-1", + "hnsw-cpu-100seg-force-10", +) +CAGRA_HNSW_CASES = ( + "cagra-hnsw-1layer", + "cagra-hnsw-3layer", +) +ALGORITHM_MATRIX_CASES = ( + "hnsw-cpu", + "cagra-hnsw-1layer", + "cagra-hnsw-3layer", + "cagra", +) +CASE_GROUPS = { + "gpu": BASIC_GPU_CASES, + "gpu-basic": BASIC_GPU_CASES, + "gpu-segments": SEGMENT_GPU_CASES, + "segments": SEGMENT_GPU_CASES, + "cpu-hnsw": CPU_HNSW_CASES, + "cagra-hnsw": CAGRA_HNSW_CASES, + "algorithm-matrix": ALGORITHM_MATRIX_CASES, + "all": BASIC_GPU_CASES + SEGMENT_GPU_CASES + CPU_HNSW_CASES + CAGRA_HNSW_CASES, +} ID_FIELD = "id" +GROUP_FIELD = "group" VECTOR_FIELD = "vector" +MIN_ROWS_PER_SEGMENT_WITH_DEFAULT_GRAPH_PARAMS = 257 +# graphDegree=32 gives M=16, and the smoke suite omits every 11th vector when +# update coverage is enabled. Keep the second upper HNSW layer above cuVS's +# derived NN_DESCENT intermediate degree of 96 to avoid native clamp warnings. +MIN_ROWS_FOR_THREE_HNSW_LAYERS_WITH_GRAPH_DEGREE_32 = 27316 +FORCE_CPU_HNSW_FALLBACK_PROPERTY = "cuvs.lucene.forceCpuHnswFallback" +HNSW_WRITER_PATH_PROPERTY = "cuvs.lucene.lastHnswWriterPath" +CAGRA_WRITER_PATH_PROPERTY = "cuvs.lucene.lastCagraWriterPath" +HNSW_LAYERS_PROPERTY = "cuvs.lucene.lastHnswLayers" +HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY = "cuvs.lucene.lastHnswCagraGraphBuildAlgo" +HNSW_CAGRA_GRAPH_DEGREE_PROPERTY = "cuvs.lucene.lastHnswCagraGraphDegree" +HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY = ( + "cuvs.lucene.lastHnswCagraIntermediateGraphDegree" +) + + +@dataclass(frozen=True) +class SmokeCase: + name: str + codec_name: str + row_count: int + dims: int + top_k: int + expected_suffixes: tuple + require_cuvs: bool = False + exercise_updates: bool = False + segment_count: int = 1 + force_merge_target: int = 0 + disable_background_merges: bool = False + assert_segments: bool = False + force_cpu_hnsw: bool = False + expected_hnsw_layers: int = 0 + expected_cagra_graph_build_algo: str = "" + expected_cagra_graph_degree: int = 0 + expected_cagra_intermediate_graph_degree: int = 0 + + +@dataclass +class PyLuceneContext: + cuvs_lucene_jar: Path + cuvs_java_jar: Path + codec_class: object + jarray: object + codec_cache: dict + + +def env_flag(name, default=False): + value = os.environ.get(name) + if value is None: + return default + return value.lower() in {"1", "true", "yes", "on"} + + +def int_env(name, default): + value = os.environ.get(name) + if value is None: + return default + return int(value) + + +def suffixes_from_env(): + configured = os.environ.get("CUVS_LUCENE_EXPECT_INDEX_SUFFIXES") + if configured: + return tuple(suffix.strip() for suffix in configured.split(",") if suffix.strip()) + if os.environ.get("CUVS_LUCENE_EXPECT_CUVS_FILES") == "1": + return (".vcag", ".vemc") + return () def find_cuvs_lucene_jar(): @@ -107,105 +231,755 @@ def init_vm(cuvs_java_jar, cuvs_lucene_jar): ) -def main(): - cuvs_lucene_jar = find_cuvs_lucene_jar() - cuvs_java_jar = find_cuvs_java_jar() - init_vm(cuvs_java_jar, cuvs_lucene_jar) +def quiet_expected_provider_fallback_logs(): + if not env_flag("CUVS_LUCENE_QUIET_EXPECTED_PROVIDER_LOGS", True): + return + try: + from java.util.logging import Level, Logger + + Logger.getLogger("com.nvidia.cuvs.lucene.LuceneProvider").setLevel(Level.OFF) + Logger.getLogger( + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat" + ).setLevel(Level.SEVERE) + Logger.getLogger( + "org.apache.lucene.internal.vectorization.PanamaVectorizationProvider" + ).setLevel(Level.WARNING) + except Exception: + pass + + +def print_threadlocal_provider_diagnostics(): + if not env_flag("CUVS_LUCENE_PRINT_THREADLOCAL_PROVIDER", False): + return + + from java.lang import System + + print("ThreadLocalCuVSResourcesProvider diagnostics:") + for property_name in ( + FORCE_CPU_HNSW_FALLBACK_PROPERTY, + HNSW_WRITER_PATH_PROPERTY, + CAGRA_WRITER_PATH_PROPERTY, + HNSW_LAYERS_PROPERTY, + HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY, + HNSW_CAGRA_GRAPH_DEGREE_PROPERTY, + HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY, + ): + value = System.getProperty(property_name) + print(f" {property_name}={value if value is not None else ''}") + + source_path = ( + REPO_ROOT + / "src" + / "main" + / "java" + / "com" + / "nvidia" + / "cuvs" + / "lucene" + / "ThreadLocalCuVSResourcesProvider.java" + ) + if not source_path.exists(): + print(f" source={source_path} (missing)") + return + + print(f" source={source_path}") + for line_number, line in enumerate(source_path.read_text().splitlines(), start=1): + print(f" {line_number:4d}: {line}") + + +def expand_case_names(names): + expanded = [] + for name in names: + group = CASE_GROUPS.get(name) + if group: + expanded.extend(expand_case_names(group)) + else: + expanded.append(name) + return expanded + + +def matrix_case( + name, + codec_name, + expected_suffixes, + require_cuvs, + segment_count=1, + force_merge_target=0, + disable_background_merges=False, + assert_segments=False, + exercise_updates=False, + row_count_floor=0, + force_cpu_hnsw=False, + expected_hnsw_layers=0, + expected_cagra_graph_build_algo="", + expected_cagra_graph_degree=0, + expected_cagra_intermediate_graph_degree=0, +): + matrix_rows = max(int_env("CUVS_LUCENE_PYLUCENE_ROWS", 2000), row_count_floor) + matrix_dims = int_env("CUVS_LUCENE_PYLUCENE_DIMS", 32) + matrix_top_k = int_env("CUVS_LUCENE_PYLUCENE_TOPK", 20) + return SmokeCase( + name=name, + codec_name=codec_name, + row_count=matrix_rows, + dims=matrix_dims, + top_k=matrix_top_k, + expected_suffixes=expected_suffixes, + require_cuvs=require_cuvs, + exercise_updates=exercise_updates, + segment_count=segment_count, + force_merge_target=force_merge_target, + disable_background_merges=disable_background_merges, + assert_segments=assert_segments, + force_cpu_hnsw=force_cpu_hnsw, + expected_hnsw_layers=expected_hnsw_layers, + expected_cagra_graph_build_algo=expected_cagra_graph_build_algo, + expected_cagra_graph_degree=expected_cagra_graph_degree, + expected_cagra_intermediate_graph_degree=expected_cagra_intermediate_graph_degree, + ) - from lucene import JArray - from java.lang import Class + +def build_segment_case( + name, + codec_name, + expected_suffixes, + require_cuvs, + segments, + force_target, + min_rows_per_segment=MIN_ROWS_PER_SEGMENT_WITH_DEFAULT_GRAPH_PARAMS, + force_cpu_hnsw=False, + expected_hnsw_layers=0, + expected_cagra_graph_build_algo="", + expected_cagra_graph_degree=0, + expected_cagra_intermediate_graph_degree=0, +): + force_merge_target = force_target or 0 + return matrix_case( + name=name, + codec_name=codec_name, + expected_suffixes=expected_suffixes, + require_cuvs=require_cuvs, + segment_count=segments, + force_merge_target=force_merge_target, + disable_background_merges=(force_merge_target == 0), + assert_segments=True, + row_count_floor=segments * min_rows_per_segment, + force_cpu_hnsw=force_cpu_hnsw, + expected_hnsw_layers=expected_hnsw_layers, + expected_cagra_graph_build_algo=expected_cagra_graph_build_algo, + expected_cagra_graph_degree=expected_cagra_graph_degree, + expected_cagra_intermediate_graph_degree=expected_cagra_intermediate_graph_degree, + ) + + +def build_named_case(name): + matrix_dims = int_env("CUVS_LUCENE_PYLUCENE_DIMS", 32) + force_merge = env_flag("CUVS_LUCENE_PYLUCENE_FORCE_MERGE", True) + exercise_updates = env_flag("CUVS_LUCENE_PYLUCENE_EXERCISE_UPDATES", True) + require_cuvs = env_flag("CUVS_LUCENE_REQUIRE_CUVS", False) + + if name == "smoke": + return SmokeCase( + name="smoke", + codec_name=HNSW_CODEC, + row_count=int_env("CUVS_LUCENE_PYLUCENE_ROWS", 2), + dims=int_env("CUVS_LUCENE_PYLUCENE_DIMS", 3), + top_k=int_env("CUVS_LUCENE_PYLUCENE_TOPK", 2), + expected_suffixes=suffixes_from_env(), + require_cuvs=require_cuvs, + ) + if name == "hnsw": + return matrix_case( + name="hnsw", + codec_name=HNSW_CODEC, + expected_suffixes=(".vex", ".vem"), + require_cuvs=require_cuvs, + exercise_updates=exercise_updates, + segment_count=3, + force_merge_target=(1 if force_merge else 0), + ) + if name == "cagra": + return matrix_case( + name="cagra", + codec_name=CAGRA_CODEC, + expected_suffixes=(".vcag", ".vemc"), + require_cuvs=True, + exercise_updates=exercise_updates, + segment_count=3, + force_merge_target=(1 if force_merge else 0), + ) + if name in {"hnsw-single", "single-hnsw"}: + return SmokeCase( + name="hnsw-single", + codec_name=HNSW_CODEC, + row_count=1, + dims=matrix_dims, + top_k=1, + expected_suffixes=(".vex", ".vem"), + require_cuvs=require_cuvs, + ) + if name in {"cagra-single", "single-cagra"}: + return SmokeCase( + name="cagra-single", + codec_name=CAGRA_CODEC, + row_count=1, + dims=matrix_dims, + top_k=1, + expected_suffixes=(".vcag", ".vemc"), + require_cuvs=True, + ) + if name == "hnsw-1seg": + return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 1, 0) + if name == "cagra-1seg": + return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 1, 0) + if name == "hnsw-10seg": + return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 10, 0) + if name == "cagra-10seg": + return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 10, 0) + if name == "hnsw-10seg-force-1": + return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 10, 1) + if name == "cagra-10seg-force-1": + return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 10, 1) + if name == "hnsw-100seg-force-10": + return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 100, 10) + if name == "cagra-100seg-force-10": + return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 100, 10) + if name in {"cagra-hnsw-1layer", "cagra-hnsw-base", "cagra-hnsw-base-layer"}: + return matrix_case( + name="cagra-hnsw-1layer", + codec_name=CAGRA_HNSW_BASE_LAYER_CODEC, + expected_suffixes=(".vex", ".vem"), + require_cuvs=True, + exercise_updates=exercise_updates, + segment_count=1, + expected_hnsw_layers=1, + expected_cagra_graph_build_algo="NN_DESCENT", + expected_cagra_graph_degree=32, + expected_cagra_intermediate_graph_degree=64, + ) + if name in { + "cagra-hnsw-3layer", + "cagra-hnsw-multilayer", + "cagra-hnsw-multi", + "cagra-hnsw-multi-layer", + }: + return matrix_case( + name="cagra-hnsw-3layer", + codec_name=CAGRA_HNSW_MULTI_LAYER_CODEC, + expected_suffixes=(".vex", ".vem"), + require_cuvs=True, + exercise_updates=exercise_updates, + segment_count=1, + row_count_floor=MIN_ROWS_FOR_THREE_HNSW_LAYERS_WITH_GRAPH_DEGREE_32, + expected_hnsw_layers=3, + expected_cagra_graph_build_algo="NN_DESCENT", + expected_cagra_graph_degree=32, + expected_cagra_intermediate_graph_degree=64, + ) + if name == "hnsw-cpu": + return matrix_case( + name="hnsw-cpu", + codec_name=HNSW_CODEC, + expected_suffixes=(".vex", ".vem"), + require_cuvs=False, + exercise_updates=exercise_updates, + segment_count=3, + force_merge_target=(1 if force_merge else 0), + force_cpu_hnsw=True, + ) + if name in {"hnsw-cpu-single", "single-hnsw-cpu"}: + return SmokeCase( + name="hnsw-cpu-single", + codec_name=HNSW_CODEC, + row_count=1, + dims=matrix_dims, + top_k=1, + expected_suffixes=(".vex", ".vem"), + force_cpu_hnsw=True, + ) + if name == "hnsw-cpu-1seg": + return build_segment_case( + name, HNSW_CODEC, (".vex", ".vem"), False, 1, 0, force_cpu_hnsw=True + ) + if name == "hnsw-cpu-10seg": + return build_segment_case( + name, HNSW_CODEC, (".vex", ".vem"), False, 10, 0, force_cpu_hnsw=True + ) + if name == "hnsw-cpu-10seg-force-1": + return build_segment_case( + name, HNSW_CODEC, (".vex", ".vem"), False, 10, 1, force_cpu_hnsw=True + ) + if name == "hnsw-cpu-100seg-force-10": + return build_segment_case( + name, HNSW_CODEC, (".vex", ".vem"), False, 100, 10, force_cpu_hnsw=True + ) + raise ValueError(f"Unknown PyLucene smoke case: {name}") + + +def cases_from_env(): + configured_cases = os.environ.get("CUVS_LUCENE_PYLUCENE_CASES") + if configured_cases: + names = [name.strip() for name in configured_cases.split(",") if name.strip()] + return [build_named_case(name) for name in expand_case_names(names)] + + configured_codec = os.environ.get("CUVS_LUCENE_PYLUCENE_CODEC") + if configured_codec: + return [ + SmokeCase( + name="custom", + codec_name=configured_codec, + row_count=int_env("CUVS_LUCENE_PYLUCENE_ROWS", 2), + dims=int_env("CUVS_LUCENE_PYLUCENE_DIMS", 3), + top_k=int_env("CUVS_LUCENE_PYLUCENE_TOPK", 2), + expected_suffixes=suffixes_from_env(), + require_cuvs=env_flag("CUVS_LUCENE_REQUIRE_CUVS", False), + exercise_updates=env_flag("CUVS_LUCENE_PYLUCENE_EXERCISE_UPDATES", False), + force_merge_target=( + 1 if env_flag("CUVS_LUCENE_PYLUCENE_FORCE_MERGE", False) else 0 + ), + ) + ] + + return [build_named_case("smoke")] + + +SELECTED_CASES = cases_from_env() + + +def verify_codecs_advertised(codec_class): + available_codecs = codec_class.availableCodecs() + codecs_to_check = (EXPECTED_CODECS if env_flag("CUVS_LUCENE_VERIFY_ALL_CODECS", True) else ()) + for codec_name in codecs_to_check: + if not available_codecs.contains(codec_name): + raise AssertionError( + f"{codec_name} was not advertised by Lucene SPI. " + f"Available codecs: {available_codecs}" + ) + + +def clear_writer_path_properties(): + from java.lang import System + + System.clearProperty(HNSW_WRITER_PATH_PROPERTY) + System.clearProperty(CAGRA_WRITER_PATH_PROPERTY) + System.clearProperty(HNSW_LAYERS_PROPERTY) + System.clearProperty(HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY) + System.clearProperty(HNSW_CAGRA_GRAPH_DEGREE_PROPERTY) + System.clearProperty(HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY) + + +def expected_writer_path(case): + if case.codec_name == CAGRA_CODEC: + return "gpu-cagra" + if case.force_cpu_hnsw: + return "cpu-hnsw-fallback" + if case.codec_name in HNSW_SEARCH_CODECS and case.require_cuvs: + return "gpu-hnsw" + return None + + +def observed_writer_path(case): + from java.lang import System + + property_name = ( + CAGRA_WRITER_PATH_PROPERTY + if case.codec_name == CAGRA_CODEC + else HNSW_WRITER_PATH_PROPERTY + ) + observed = System.getProperty(property_name) + expected = expected_writer_path(case) + if expected is not None and observed != expected: + raise AssertionError( + f"{case.name}: expected writer path {expected}, got {observed or ''}" + ) + return observed or "unknown" + + +def assert_hnsw_telemetry(case): + from java.lang import System + + if case.expected_hnsw_layers: + observed_layers = System.getProperty(HNSW_LAYERS_PROPERTY) + expected_layers = str(case.expected_hnsw_layers) + if observed_layers != expected_layers: + raise AssertionError( + f"{case.name}: expected HNSW layers {expected_layers}, " + f"got {observed_layers or ''}" + ) + + if case.expected_cagra_graph_build_algo: + observed_algo = System.getProperty(HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY) + if observed_algo != case.expected_cagra_graph_build_algo: + raise AssertionError( + f"{case.name}: expected CAGRA graph build algorithm " + f"{case.expected_cagra_graph_build_algo}, got {observed_algo or ''}" + ) + + if case.expected_cagra_graph_degree: + observed_degree = System.getProperty(HNSW_CAGRA_GRAPH_DEGREE_PROPERTY) + expected_degree = str(case.expected_cagra_graph_degree) + if observed_degree != expected_degree: + raise AssertionError( + f"{case.name}: expected CAGRA graph degree {expected_degree}, " + f"got {observed_degree or ''}" + ) + + if case.expected_cagra_intermediate_graph_degree: + observed_intermediate_degree = System.getProperty( + HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY + ) + expected_intermediate_degree = str(case.expected_cagra_intermediate_graph_degree) + if observed_intermediate_degree != expected_intermediate_degree: + raise AssertionError( + f"{case.name}: expected CAGRA intermediate graph degree " + f"{expected_intermediate_degree}, " + f"got {observed_intermediate_degree or ''}" + ) + + +def build_search_label(case, writer_path): + if case.codec_name == CAGRA_CODEC: + label = "build=cagra, search=cagra" + elif case.force_cpu_hnsw or writer_path.startswith("cpu-hnsw"): + label = "build=hnsw, search=hnsw" + elif case.codec_name in HNSW_SEARCH_CODECS: + label = "build=cagra, search=hnsw" + else: + label = "build=unknown, search=unknown" + + if case.expected_hnsw_layers: + label += f", hnswLayers={case.expected_hnsw_layers}" + if case.expected_cagra_graph_build_algo: + label += f", cagraGraphBuildAlgo={case.expected_cagra_graph_build_algo}" + if case.expected_cagra_graph_degree: + label += f", cagraGraphDegree={case.expected_cagra_graph_degree}" + if case.expected_cagra_intermediate_graph_degree: + label += ( + ", cagraIntermediateGraphDegree=" + f"{case.expected_cagra_intermediate_graph_degree}" + ) + return label + + +@contextmanager +def hnsw_cpu_fallback(case): + if not case.force_cpu_hnsw: + yield + return + + from java.lang import System + + previous_value = System.getProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY) + System.setProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY, "true") + try: + yield + finally: + if previous_value is None: + System.clearProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY) + else: + System.setProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY, previous_value) + + +def missing_doc_ids(case): + if not case.exercise_updates or case.row_count < 30: + return set() + return {doc_id for doc_id in range(case.row_count) if doc_id % 11 == 0} + + +def deleted_doc_ids(case, missing_ids): + if not case.exercise_updates or case.row_count < 30: + return set() + + deleted = set() + for target in (case.row_count // 4, case.row_count // 2, (case.row_count * 3) // 4): + for offset in range(case.row_count): + for candidate in (target + offset, target - offset): + if 0 <= candidate < case.row_count and candidate not in missing_ids: + deleted.add(candidate) + break + if len(deleted) >= 3: + break + if len(deleted) >= 3: + break + return deleted + + +def choose_query_ids(row_count, active_vector_ids): + targets = [0, row_count // 2, row_count - 1] + query_ids = [] + for target in targets: + nearest = min(active_vector_ids, key=lambda doc_id: abs(doc_id - target)) + if nearest not in query_ids: + query_ids.append(nearest) + return query_ids + + +def effective_segment_count(case): + return max(1, min(case.segment_count, case.row_count)) + + +def segment_end_doc_ids(case): + segments = effective_segment_count(case) + docs_per_segment, remainder = divmod(case.row_count, segments) + end_doc_ids = set() + end_exclusive = 0 + for segment_id in range(segments): + end_exclusive += docs_per_segment + (1 if segment_id < remainder else 0) + end_doc_ids.add(end_exclusive - 1) + return end_doc_ids + + +def assert_expected_index_files(index_path, expected_suffixes): + if not expected_suffixes: + return + index_files = sorted(path.name for path in Path(index_path).iterdir()) + for suffix in expected_suffixes: + if not any(name.endswith(suffix) for name in index_files): + raise AssertionError( + f"No index file ending with {suffix} found: {index_files}" + ) + + +def assert_vector_metadata(reader, vector_field, expected_count, expected_dims): + vector_count = 0 + for leaf_reader_context in reader.leaves(): + leaf_reader = leaf_reader_context.reader() + values = leaf_reader.getFloatVectorValues(vector_field) + if values is None: + continue + if values.dimension() != expected_dims: + raise AssertionError( + f"Vector dimension mismatch: expected {expected_dims}, got {values.dimension()}" + ) + vector_count += values.size() + if vector_count != expected_count: + raise AssertionError( + f"Vector count mismatch: expected {expected_count}, got {vector_count}" + ) + + +def assert_segment_topology(reader, case): + if not case.assert_segments: + return + + segment_count = sum(1 for _ in reader.leaves()) + if case.force_merge_target: + if segment_count > case.force_merge_target: + raise AssertionError( + f"{case.name}: expected at most {case.force_merge_target} segment(s), " + f"got {segment_count}" + ) + if case.force_merge_target == 1 and segment_count != 1: + raise AssertionError(f"{case.name}: expected one merged segment, got {segment_count}") + return + + if case.disable_background_merges: + expected_segments = effective_segment_count(case) + if segment_count != expected_segments: + raise AssertionError( + f"{case.name}: expected {expected_segments} unmerged segment(s), " + f"got {segment_count}" + ) + + +def hit_ids(stored_fields, hits): + return [stored_fields.document(hit.doc).get(ID_FIELD) for hit in hits] + + +def assert_search_results(searcher, jarray, case, query_ids, inactive_ids): + from org.apache.lucene.search import KnnFloatVectorQuery + + stored_fields = searcher.storedFields() + top_k = min(case.top_k, max(1, case.row_count - len(inactive_ids))) + for query_id in query_ids: + query = KnnFloatVectorQuery( + VECTOR_FIELD, fvec(jarray, vector_for(query_id, case.dims)), top_k + ) + ids = hit_ids(stored_fields, searcher.search(query, top_k).scoreDocs) + expected = f"doc-{query_id}" + if expected not in ids: + raise AssertionError(f"{case.name}: expected {expected} in top {top_k}, got {ids}") + bad_ids = [doc_id for doc_id in ids if doc_id in inactive_ids] + if bad_ids: + raise AssertionError(f"{case.name}: inactive docs returned: {bad_ids}") + + +def assert_filtered_search(searcher, jarray, case, query_id): + from org.apache.lucene.index import Term + from org.apache.lucene.search import KnnFloatVectorQuery, TermQuery + + expected = f"doc-{query_id}" + filter_query = TermQuery(Term(ID_FIELD, expected)) + query = KnnFloatVectorQuery( + VECTOR_FIELD, fvec(jarray, vector_for(query_id, case.dims)), 1, filter_query + ) + ids = hit_ids(searcher.storedFields(), searcher.search(query, 1).scoreDocs) + if ids != [expected]: + raise AssertionError(f"{case.name}: filtered search expected {[expected]}, got {ids}") + + +def codec_for_case(case, codec_class, codec_cache): + codec = codec_cache.get(case.codec_name) + if codec is not None: + return codec + + codec = codec_class.forName(case.codec_name) + if codec.getName() != case.codec_name: + raise AssertionError(f"Expected codec {case.codec_name}, got {codec.getName()}") + codec_cache[case.codec_name] = codec + return codec + + +def run_case(case, codec_class, codec_cache, jarray): from java.nio.file import Paths - from org.apache.lucene.codecs import Codec from org.apache.lucene.document import Document, Field, KnnFloatVectorField, StringField from org.apache.lucene.index import ( DirectoryReader, IndexWriter, IndexWriterConfig, + Term, VectorSimilarityFunction, ) - from org.apache.lucene.search import IndexSearcher, KnnFloatVectorQuery + from org.apache.lucene.search import IndexSearcher from org.apache.lucene.store import FSDirectory - # Forces an early, clear failure if cuvs-java was flattened or omitted. - Class.forName("com.nvidia.cuvs.spi.JDKProvider") - - codec_name = os.environ.get("CUVS_LUCENE_PYLUCENE_CODEC", DEFAULT_CODEC) - row_count = int(os.environ.get("CUVS_LUCENE_PYLUCENE_ROWS", "2")) - dims = int(os.environ.get("CUVS_LUCENE_PYLUCENE_DIMS", "3")) - top_k = int(os.environ.get("CUVS_LUCENE_PYLUCENE_TOPK", "2")) - expect_cuvs_files = os.environ.get("CUVS_LUCENE_EXPECT_CUVS_FILES") == "1" - - available_codecs = Codec.availableCodecs() - if not available_codecs.contains(codec_name): + available_codecs = codec_class.availableCodecs() + if not available_codecs.contains(case.codec_name): raise AssertionError( - f"{codec_name} was not advertised by Lucene SPI. " + f"{case.codec_name} was not advertised by Lucene SPI. " f"Available codecs: {available_codecs}" ) - - codec = Codec.forName(codec_name) - if codec.getName() != codec_name: - raise AssertionError(f"Expected codec {codec_name}, got {codec.getName()}") - - with tempfile.TemporaryDirectory(prefix="cuvs-lucene-pylucene-") as index_path: + codec = codec_for_case(case, codec_class, codec_cache) + + missing_ids = missing_doc_ids(case) + deleted_ids = deleted_doc_ids(case, missing_ids) + active_vector_ids = [ + doc_id + for doc_id in range(case.row_count) + if doc_id not in missing_ids and doc_id not in deleted_ids + ] + if not active_vector_ids: + raise AssertionError(f"{case.name}: no active vectors available for search") + query_ids = choose_query_ids(case.row_count, active_vector_ids) + inactive_doc_names = {f"doc-{doc_id}" for doc_id in missing_ids | deleted_ids} + writer_path = "unknown" + + with ExitStack() as stack: + stack.enter_context(hnsw_cpu_fallback(case)) + clear_writer_path_properties() + index_path = stack.enter_context( + tempfile.TemporaryDirectory(prefix=f"cuvs-lucene-pylucene-{case.name}-") + ) directory = FSDirectory.open(Paths.get(index_path)) config = IndexWriterConfig() config.setCodec(codec) config.setUseCompoundFile(False) + if case.disable_background_merges: + from org.apache.lucene.index import NoMergePolicy + + config.setMergePolicy(NoMergePolicy.INSTANCE) writer = IndexWriter(directory, config) try: - for doc_id in range(row_count): + segment_ends = segment_end_doc_ids(case) + for doc_id in range(case.row_count): doc = Document() doc.add(StringField(ID_FIELD, f"doc-{doc_id}", Field.Store.YES)) - doc.add( - KnnFloatVectorField( - VECTOR_FIELD, - fvec(JArray, vector_for(doc_id, dims)), - VectorSimilarityFunction.EUCLIDEAN, + doc.add(StringField(GROUP_FIELD, f"group-{doc_id % 3}", Field.Store.YES)) + if doc_id not in missing_ids: + doc.add( + KnnFloatVectorField( + VECTOR_FIELD, + fvec(jarray, vector_for(doc_id, case.dims)), + VectorSimilarityFunction.EUCLIDEAN, + ) ) - ) writer.addDocument(doc) + if doc_id in segment_ends: + writer.commit() + writer.commit() + for doc_id in deleted_ids: + writer.deleteDocuments(Term(ID_FIELD, f"doc-{doc_id}")) + if deleted_ids: + writer.commit() + if case.force_merge_target: + writer.forceMerge(case.force_merge_target) + writer.commit() finally: writer.close() - if expect_cuvs_files: - index_files = sorted(path.name for path in Path(index_path).iterdir()) - if not any(name.endswith(".vcag") for name in index_files): - raise AssertionError(f"No cuVS .vcag file found: {index_files}") - if not any(name.endswith(".vemc") for name in index_files): - raise AssertionError(f"No cuVS .vemc file found: {index_files}") + assert_expected_index_files(index_path, case.expected_suffixes) reader = DirectoryReader.open(directory) try: - searcher = IndexSearcher(reader) - stored_fields = searcher.storedFields() - query_ids = sorted({0, row_count // 2, row_count - 1}) - for query_id in query_ids: - query = KnnFloatVectorQuery( - VECTOR_FIELD, fvec(JArray, vector_for(query_id, dims)), top_k + expected_live_docs = case.row_count - len(deleted_ids) + if reader.numDocs() != expected_live_docs: + raise AssertionError( + f"{case.name}: expected {expected_live_docs} live docs, got {reader.numDocs()}" ) - hits = searcher.search(query, top_k).scoreDocs - ids = [stored_fields.document(hit.doc).get(ID_FIELD) for hit in hits] - expected = f"doc-{query_id}" - if expected not in ids: - raise AssertionError( - f"Expected {expected} in top {top_k}, got {ids}" - ) + expected_vector_count = ( + len(active_vector_ids) + if case.force_merge_target + else case.row_count - len(missing_ids) + ) + assert_vector_metadata(reader, VECTOR_FIELD, expected_vector_count, case.dims) + assert_segment_topology(reader, case) + searcher = IndexSearcher(reader) + assert_search_results(searcher, jarray, case, query_ids, inactive_doc_names) + assert_filtered_search(searcher, jarray, case, query_ids[0]) + writer_path = observed_writer_path(case) + assert_hnsw_telemetry(case) finally: reader.close() directory.close() + algorithm_label = build_search_label(case, writer_path) print( - "PyLucene smoke test passed: " - f"loaded {codec_name} from {cuvs_lucene_jar.name}, " - f"used {cuvs_java_jar.name}, and searched {row_count} docs" + "PASS: " + f"{case.name} loaded {case.codec_name}, " + f"indexed {case.row_count} docs x {case.dims} dims, " + f"segments={effective_segment_count(case)}" + f"{'->' + str(case.force_merge_target) if case.force_merge_target else ''}, " + f"searched topK={case.top_k}, " + f"{algorithm_label}, " + f"path={writer_path}" ) +def initialize_pylucene_context(): + cuvs_lucene_jar = find_cuvs_lucene_jar() + cuvs_java_jar = find_cuvs_java_jar() + init_vm(cuvs_java_jar, cuvs_lucene_jar) + + from lucene import JArray + from java.lang import Class + + # Forces an early, clear failure if cuvs-java was flattened or omitted. + Class.forName("com.nvidia.cuvs.spi.JDKProvider") + + quiet_expected_provider_fallback_logs() + print_threadlocal_provider_diagnostics() + + from org.apache.lucene.codecs import Codec + + verify_codecs_advertised(Codec) + + return PyLuceneContext( + cuvs_lucene_jar=cuvs_lucene_jar, + cuvs_java_jar=cuvs_java_jar, + codec_class=Codec, + jarray=JArray, + codec_cache={}, + ) + + +def main(): + import pytest + + test_path = Path(__file__).with_name("test_pylucene_smoke.py") + return pytest.main(["-q", "-s", str(test_path)]) + + if __name__ == "__main__": - try: - main() - except Exception as exc: - print(f"PyLucene smoke test failed: {exc}", file=sys.stderr) - raise + raise SystemExit(main()) diff --git a/examples/test_pylucene_smoke.py b/examples/test_pylucene_smoke.py new file mode 100644 index 00000000..fc001daa --- /dev/null +++ b/examples/test_pylucene_smoke.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from pylucene_smoke import SELECTED_CASES, initialize_pylucene_context, run_case + + +@pytest.fixture(scope="session") +def pylucene_context(request): + context = initialize_pylucene_context() + yield context + + if request.session.testsfailed == 0: + print( + "PyLucene smoke suite passed: " + f"loaded {context.cuvs_lucene_jar.name}, used {context.cuvs_java_jar.name}, " + f"ran {len(SELECTED_CASES)} case(s)" + ) + + +@pytest.mark.parametrize("case", SELECTED_CASES, ids=lambda case: case.name) +def test_pylucene_smoke_case(pylucene_context, case): + run_case( + case, + pylucene_context.codec_class, + pylucene_context.codec_cache, + pylucene_context.jarray, + ) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index ccc61eae..35e8bd3a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -34,6 +34,8 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { public static final String CUVS_INDEX_EXT = "vcag"; public static final int VERSION_START = 0; public static final int VERSION_CURRENT = VERSION_START; + public static final String LAST_WRITER_PATH_PROPERTY = "cuvs.lucene.lastCagraWriterPath"; + public static final String WRITER_PATH_GPU = "gpu-cagra"; private GPUSearchParams gpuSearchParams; @@ -73,6 +75,7 @@ public CuVS2510GPUVectorsFormat(GPUSearchParams gpuSearchParams) { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { assertIsSupported(); + System.setProperty(LAST_WRITER_PATH_PROPERTY, WRITER_PATH_GPU); var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); return new CuVS2510GPUVectorsWriter(state, gpuSearchParams, flatWriter); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.java new file mode 100644 index 00000000..9fab1971 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.java @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; + +/** + * Accelerated HNSW codec that builds an intermediate CAGRA graph and writes only the base HNSW layer. + */ +public class Lucene101AcceleratedHNSWBaseLayerCodec extends Lucene101AcceleratedHNSWCodec { + + private static final String NAME = "Lucene101AcceleratedHNSWBaseLayerCodec"; + private static final int CAGRA_GRAPH_DEGREE = 32; + private static final int CAGRA_INTERMEDIATE_GRAPH_DEGREE = 64; + + /** Default constructor used by Lucene SPI. */ + public Lucene101AcceleratedHNSWBaseLayerCodec() throws Exception { + super( + NAME, + LuceneProvider.getDefaultDelegateCodec(), + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(CAGRA_GRAPH_DEGREE) + .withIntermediateGraphDegree(CAGRA_INTERMEDIATE_GRAPH_DEGREE) + .withHNSWLayer(1) + .build()); + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index 4944b761..7ab4a154 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -52,7 +52,19 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { */ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getDefaultDelegateCodec()); + this(NAME, LuceneProvider.getDefaultDelegateCodec(), acceleratedHNSWParams); + } + + /** + * Constructor for subclasses that expose named accelerated HNSW configurations via SPI. + * + * @param name the codec's name + * @param delegate the delegate codec to filter + * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} + */ + protected Lucene101AcceleratedHNSWCodec( + String name, Codec delegate, AcceleratedHNSWParams acceleratedHNSWParams) { + super(name, delegate); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.java new file mode 100644 index 00000000..3ad87b4b --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.java @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; + +/** + * Accelerated HNSW codec that builds an intermediate CAGRA graph and writes multiple HNSW layers. + */ +public class Lucene101AcceleratedHNSWMultiLayerCodec extends Lucene101AcceleratedHNSWCodec { + + private static final String NAME = "Lucene101AcceleratedHNSWMultiLayerCodec"; + private static final int CAGRA_GRAPH_DEGREE = 32; + private static final int CAGRA_INTERMEDIATE_GRAPH_DEGREE = 64; + + /** Default constructor used by Lucene SPI. */ + public Lucene101AcceleratedHNSWMultiLayerCodec() throws Exception { + super( + NAME, + LuceneProvider.getDefaultDelegateCodec(), + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(CAGRA_GRAPH_DEGREE) + .withIntermediateGraphDegree(CAGRA_INTERMEDIATE_GRAPH_DEGREE) + .withHNSWLayer(3) + .build()); + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 3c42707c..2c9f024a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isCpuHnswFallbackForced; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import com.nvidia.cuvs.LibraryException; @@ -28,6 +29,17 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final Logger log = Logger.getLogger(Lucene99AcceleratedHNSWVectorsFormat.class.getName()); + public static final String LAST_WRITER_PATH_PROPERTY = "cuvs.lucene.lastHnswWriterPath"; + public static final String LAST_HNSW_LAYERS_PROPERTY = "cuvs.lucene.lastHnswLayers"; + public static final String LAST_CAGRA_GRAPH_BUILD_ALGO_PROPERTY = + "cuvs.lucene.lastHnswCagraGraphBuildAlgo"; + public static final String LAST_CAGRA_GRAPH_DEGREE_PROPERTY = + "cuvs.lucene.lastHnswCagraGraphDegree"; + public static final String LAST_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY = + "cuvs.lucene.lastHnswCagraIntermediateGraphDegree"; + public static final String WRITER_PATH_GPU = "gpu-hnsw"; + public static final String WRITER_PATH_CPU_FORCED = "cpu-hnsw-fallback"; + public static final String WRITER_PATH_CPU_AUTO = "cpu-hnsw-auto-fallback"; private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; @@ -75,13 +87,32 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + System.setProperty( + LAST_HNSW_LAYERS_PROPERTY, Integer.toString(acceleratedHNSWParams.getHnswLayers())); + System.setProperty( + LAST_CAGRA_GRAPH_BUILD_ALGO_PROPERTY, + acceleratedHNSWParams.getCagraGraphBuildAlgo().name()); + System.setProperty( + LAST_CAGRA_GRAPH_DEGREE_PROPERTY, + Integer.toString(acceleratedHNSWParams.getGraphdegree())); + System.setProperty( + LAST_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY, + Integer.toString(acceleratedHNSWParams.getIntermediateGraphDegree())); if (isSupported()) { + System.setProperty(LAST_WRITER_PATH_PROPERTY, WRITER_PATH_GPU); log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); } else { + boolean forcedCpuFallback = isCpuHnswFallbackForced(); + System.setProperty( + LAST_WRITER_PATH_PROPERTY, + forcedCpuFallback ? WRITER_PATH_CPU_FORCED : WRITER_PATH_CPU_AUTO); log.log( - Level.WARNING, - "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); + forcedCpuFallback ? Level.FINE : Level.WARNING, + forcedCpuFallback + ? "Forced CPU HNSW fallback, using the Lucene99HnswVectorsWriter" + : "GPU based indexing not supported, falling back to using the" + + " Lucene99HnswVectorsWriter"); try { return LUCENE_PROVIDER.getLuceneHnswVectorsWriterInstance( state, diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index df758cc5..01db5f68 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -8,6 +8,7 @@ import java.lang.invoke.VarHandle; import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -33,6 +34,7 @@ public class LuceneProvider { static final Logger log = Logger.getLogger(LuceneProvider.class.getName()); + private static final List SUPPORTED_DELEGATE_CODEC_VERSIONS = List.of("101", "99"); private static final String BASE = "org.apache.lucene."; private static String codecs = "codecs.lucene."; @@ -152,14 +154,20 @@ private static Class loadClass(String defaultClassName, String fallbackClassN throws ClassNotFoundException { try { return Class.forName(defaultClassName); - } catch (ClassNotFoundException e) { + } catch (ClassNotFoundException defaultException) { // Load class from fallback package. try { return Class.forName(fallbackClassName); - } catch (ClassNotFoundException e1) { - // Should not reach here. - log.log(Level.SEVERE, "Unable to load class: " + fallbackClassName); - throw e1; + } catch (ClassNotFoundException fallbackException) { + ClassNotFoundException missing = + new ClassNotFoundException( + "Unable to load Lucene class. Tried " + + defaultClassName + + " and " + + fallbackClassName); + missing.addSuppressed(defaultException); + missing.addSuppressed(fallbackException); + throw missing; } } } @@ -179,18 +187,23 @@ public static Codec getCodec(String version) } public static Codec getDefaultDelegateCodec() { - for (String version : List.of("101", "99")) { + List failures = new ArrayList<>(); + for (String version : SUPPORTED_DELEGATE_CODEC_VERSIONS) { try { return getCodec(version); } catch (ReflectiveOperationException | SecurityException | IllegalArgumentException | LinkageError e) { + failures.add("Lucene" + version + ": " + e.getMessage()); log.log(Level.FINE, "Unable to load Lucene" + version + "Codec", e); } } - log.log(Level.FINE, "Falling back to the runtime default codec"); - return Codec.getDefault(); + throw new IllegalStateException( + "Unable to load a supported Lucene delegate codec. Tried " + + SUPPORTED_DELEGATE_CODEC_VERSIONS + + ". Failures: " + + failures); } public FlatVectorsFormat getLuceneFlatVectorsFormatInstance(FlatVectorsScorer scorer) diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 9e259e27..a13271d4 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -18,6 +18,7 @@ public class ThreadLocalCuVSResourcesProvider { private static final Logger log = Logger.getLogger(ThreadLocalCuVSResourcesProvider.class.getName()); + static final String FORCE_CPU_HNSW_FALLBACK_PROPERTY = "cuvs.lucene.forceCpuHnswFallback"; private static final ThreadLocal cuVSResources; static { @@ -30,6 +31,9 @@ public class ThreadLocalCuVSResourcesProvider { * @return an instance of CuVSResources */ public static CuVSResources getCuVSResourcesInstance() { + if (isCpuHnswFallbackForced()) { + return null; + } return cuVSResources.get(); } @@ -75,7 +79,7 @@ public static void closeCuVSResourcesInstance() { * @throws UnsupportedOperationException */ public static void assertIsSupported() throws UnsupportedOperationException { - if (cuVSResources.get() == null) { + if (isCpuHnswFallbackForced() || cuVSResources.get() == null) { throw new UnsupportedOperationException("cuVS is not supported"); } } @@ -86,6 +90,10 @@ public static void assertIsSupported() throws UnsupportedOperationException { * @return true if cuVS is supported else false */ public static boolean isSupported() { - return cuVSResources.get() != null; + return !isCpuHnswFallbackForced() && cuVSResources.get() != null; + } + + static boolean isCpuHnswFallbackForced() { + return Boolean.getBoolean(FORCE_CPU_HNSW_FALLBACK_PROPERTY); } } 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 index faa0684c..2b9d6a7f 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec +com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWBaseLayerCodec +com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWMultiLayerCodec com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedCodec com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedCodec diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 1407e0ab..5c974301 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -7,9 +7,10 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; +import java.util.Set; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.junit.Test; @@ -42,21 +43,30 @@ public void testExistingComponents() throws Exception { } @Test - public void testProviderCachesByVersion() throws Exception { + public void testProviderCachesSupportedVersion() throws Exception { LuceneProvider provider99 = LuceneProvider.getInstance("99"); - LuceneProvider provider102 = LuceneProvider.getInstance("102"); - assertNotSame(provider99, provider102); + assertSame(provider99, LuceneProvider.getInstance("99")); + } + + @Test(expected = ClassNotFoundException.class) + public void testProviderDoesNotPretendLucene102IsACompleteProviderVersion() throws Exception { + LuceneProvider.getInstance("102"); } @Test public void testDefaultDelegateCodec() { - assertNotNull(LuceneProvider.getDefaultDelegateCodec()); + Codec delegate = LuceneProvider.getDefaultDelegateCodec(); + assertNotNull(delegate); + assertTrue(Set.of("Lucene101", "Lucene99").contains(delegate.getName())); + assertTrue(delegate.getClass().getName().startsWith("org.apache.lucene.")); } @Test public void testServiceLoadedCodecsCanBeInstantiated() { String[] codecNames = { "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", "CuVS2510GPUSearchCodec", "Lucene101AcceleratedHNSWBinaryQuantizedCodec", "Lucene101AcceleratedHNSWScalarQuantizedCodec" From 692c3db3b143cb57b5e07ff90e8d3bcdf93be6fe Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:33:22 +0000 Subject: [PATCH 05/13] Make writer telemetry on-demand --- examples/pylucene_smoke.py | 75 +++++++------------ .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 12 ++- .../Lucene99AcceleratedHNSWVectorsFormat.java | 38 +++------- .../nvidia/cuvs/lucene/WriterTelemetry.java | 37 +++++++++ .../cuvs/lucene/TestWriterTelemetry.java | 49 ++++++++++++ 5 files changed, 133 insertions(+), 78 deletions(-) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/WriterTelemetry.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py index cbc32e1b..c3ad7c5a 100644 --- a/examples/pylucene_smoke.py +++ b/examples/pylucene_smoke.py @@ -76,14 +76,6 @@ # derived NN_DESCENT intermediate degree of 96 to avoid native clamp warnings. MIN_ROWS_FOR_THREE_HNSW_LAYERS_WITH_GRAPH_DEGREE_32 = 27316 FORCE_CPU_HNSW_FALLBACK_PROPERTY = "cuvs.lucene.forceCpuHnswFallback" -HNSW_WRITER_PATH_PROPERTY = "cuvs.lucene.lastHnswWriterPath" -CAGRA_WRITER_PATH_PROPERTY = "cuvs.lucene.lastCagraWriterPath" -HNSW_LAYERS_PROPERTY = "cuvs.lucene.lastHnswLayers" -HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY = "cuvs.lucene.lastHnswCagraGraphBuildAlgo" -HNSW_CAGRA_GRAPH_DEGREE_PROPERTY = "cuvs.lucene.lastHnswCagraGraphDegree" -HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY = ( - "cuvs.lucene.lastHnswCagraIntermediateGraphDegree" -) @dataclass(frozen=True) @@ -255,17 +247,11 @@ def print_threadlocal_provider_diagnostics(): from java.lang import System print("ThreadLocalCuVSResourcesProvider diagnostics:") - for property_name in ( - FORCE_CPU_HNSW_FALLBACK_PROPERTY, - HNSW_WRITER_PATH_PROPERTY, - CAGRA_WRITER_PATH_PROPERTY, - HNSW_LAYERS_PROPERTY, - HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY, - HNSW_CAGRA_GRAPH_DEGREE_PROPERTY, - HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY, - ): - value = System.getProperty(property_name) - print(f" {property_name}={value if value is not None else ''}") + value = System.getProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY) + print( + f" {FORCE_CPU_HNSW_FALLBACK_PROPERTY}=" + f"{value if value is not None else ''}" + ) source_path = ( REPO_ROOT @@ -557,15 +543,17 @@ def verify_codecs_advertised(codec_class): ) -def clear_writer_path_properties(): - from java.lang import System - - System.clearProperty(HNSW_WRITER_PATH_PROPERTY) - System.clearProperty(CAGRA_WRITER_PATH_PROPERTY) - System.clearProperty(HNSW_LAYERS_PROPERTY) - System.clearProperty(HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY) - System.clearProperty(HNSW_CAGRA_GRAPH_DEGREE_PROPERTY) - System.clearProperty(HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY) +def writer_telemetry(codec): + vector_format = codec.knnVectorsFormat() + method = vector_format.getClass().getMethod("getWriterTelemetry", ()) + payload = str(method.invoke(vector_format, ())) + telemetry = {} + for item in payload.split(";"): + key, separator, value = item.partition("=") + if not separator or not key: + raise AssertionError(f"Malformed writer telemetry item: {item!r}") + telemetry[key] = value + return telemetry def expected_writer_path(case): @@ -578,15 +566,8 @@ def expected_writer_path(case): return None -def observed_writer_path(case): - from java.lang import System - - property_name = ( - CAGRA_WRITER_PATH_PROPERTY - if case.codec_name == CAGRA_CODEC - else HNSW_WRITER_PATH_PROPERTY - ) - observed = System.getProperty(property_name) +def observed_writer_path(case, telemetry): + observed = telemetry.get("writerPath") expected = expected_writer_path(case) if expected is not None and observed != expected: raise AssertionError( @@ -595,11 +576,9 @@ def observed_writer_path(case): return observed or "unknown" -def assert_hnsw_telemetry(case): - from java.lang import System - +def assert_hnsw_telemetry(case, telemetry): if case.expected_hnsw_layers: - observed_layers = System.getProperty(HNSW_LAYERS_PROPERTY) + observed_layers = telemetry.get("hnswLayers") expected_layers = str(case.expected_hnsw_layers) if observed_layers != expected_layers: raise AssertionError( @@ -608,7 +587,7 @@ def assert_hnsw_telemetry(case): ) if case.expected_cagra_graph_build_algo: - observed_algo = System.getProperty(HNSW_CAGRA_GRAPH_BUILD_ALGO_PROPERTY) + observed_algo = telemetry.get("cagraGraphBuildAlgo") if observed_algo != case.expected_cagra_graph_build_algo: raise AssertionError( f"{case.name}: expected CAGRA graph build algorithm " @@ -616,7 +595,7 @@ def assert_hnsw_telemetry(case): ) if case.expected_cagra_graph_degree: - observed_degree = System.getProperty(HNSW_CAGRA_GRAPH_DEGREE_PROPERTY) + observed_degree = telemetry.get("cagraGraphDegree") expected_degree = str(case.expected_cagra_graph_degree) if observed_degree != expected_degree: raise AssertionError( @@ -625,9 +604,7 @@ def assert_hnsw_telemetry(case): ) if case.expected_cagra_intermediate_graph_degree: - observed_intermediate_degree = System.getProperty( - HNSW_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY - ) + observed_intermediate_degree = telemetry.get("cagraIntermediateGraphDegree") expected_intermediate_degree = str(case.expected_cagra_intermediate_graph_degree) if observed_intermediate_degree != expected_intermediate_degree: raise AssertionError( @@ -866,7 +843,6 @@ def run_case(case, codec_class, codec_cache, jarray): with ExitStack() as stack: stack.enter_context(hnsw_cpu_fallback(case)) - clear_writer_path_properties() index_path = stack.enter_context( tempfile.TemporaryDirectory(prefix=f"cuvs-lucene-pylucene-{case.name}-") ) @@ -928,8 +904,9 @@ def run_case(case, codec_class, codec_cache, jarray): searcher = IndexSearcher(reader) assert_search_results(searcher, jarray, case, query_ids, inactive_doc_names) assert_filtered_search(searcher, jarray, case, query_ids[0]) - writer_path = observed_writer_path(case) - assert_hnsw_telemetry(case) + telemetry = writer_telemetry(codec) + writer_path = observed_writer_path(case, telemetry) + assert_hnsw_telemetry(case, telemetry) finally: reader.close() directory.close() diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 35e8bd3a..257ec501 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -34,8 +34,6 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { public static final String CUVS_INDEX_EXT = "vcag"; public static final int VERSION_START = 0; public static final int VERSION_CURRENT = VERSION_START; - public static final String LAST_WRITER_PATH_PROPERTY = "cuvs.lucene.lastCagraWriterPath"; - public static final String WRITER_PATH_GPU = "gpu-cagra"; private GPUSearchParams gpuSearchParams; @@ -75,11 +73,19 @@ public CuVS2510GPUVectorsFormat(GPUSearchParams gpuSearchParams) { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { assertIsSupported(); - System.setProperty(LAST_WRITER_PATH_PROPERTY, WRITER_PATH_GPU); var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); return new CuVS2510GPUVectorsWriter(state, gpuSearchParams, flatWriter); } + /** + * Returns an on-demand diagnostic snapshot for integration tests. + * + * @return writer diagnostics + */ + public String getWriterTelemetry() { + return WriterTelemetry.forCagra(); + } + /** * Returns a KnnVectorsReader instance to read the vectors from the index. */ diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 2c9f024a..b380eedf 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -29,17 +29,6 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final Logger log = Logger.getLogger(Lucene99AcceleratedHNSWVectorsFormat.class.getName()); - public static final String LAST_WRITER_PATH_PROPERTY = "cuvs.lucene.lastHnswWriterPath"; - public static final String LAST_HNSW_LAYERS_PROPERTY = "cuvs.lucene.lastHnswLayers"; - public static final String LAST_CAGRA_GRAPH_BUILD_ALGO_PROPERTY = - "cuvs.lucene.lastHnswCagraGraphBuildAlgo"; - public static final String LAST_CAGRA_GRAPH_DEGREE_PROPERTY = - "cuvs.lucene.lastHnswCagraGraphDegree"; - public static final String LAST_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY = - "cuvs.lucene.lastHnswCagraIntermediateGraphDegree"; - public static final String WRITER_PATH_GPU = "gpu-hnsw"; - public static final String WRITER_PATH_CPU_FORCED = "cpu-hnsw-fallback"; - public static final String WRITER_PATH_CPU_AUTO = "cpu-hnsw-auto-fallback"; private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; @@ -87,26 +76,11 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); - System.setProperty( - LAST_HNSW_LAYERS_PROPERTY, Integer.toString(acceleratedHNSWParams.getHnswLayers())); - System.setProperty( - LAST_CAGRA_GRAPH_BUILD_ALGO_PROPERTY, - acceleratedHNSWParams.getCagraGraphBuildAlgo().name()); - System.setProperty( - LAST_CAGRA_GRAPH_DEGREE_PROPERTY, - Integer.toString(acceleratedHNSWParams.getGraphdegree())); - System.setProperty( - LAST_CAGRA_INTERMEDIATE_GRAPH_DEGREE_PROPERTY, - Integer.toString(acceleratedHNSWParams.getIntermediateGraphDegree())); if (isSupported()) { - System.setProperty(LAST_WRITER_PATH_PROPERTY, WRITER_PATH_GPU); log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); } else { boolean forcedCpuFallback = isCpuHnswFallbackForced(); - System.setProperty( - LAST_WRITER_PATH_PROPERTY, - forcedCpuFallback ? WRITER_PATH_CPU_FORCED : WRITER_PATH_CPU_AUTO); log.log( forcedCpuFallback ? Level.FINE : Level.WARNING, forcedCpuFallback @@ -127,6 +101,18 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException } } + /** + * Returns an on-demand diagnostic snapshot for integration tests. + * + *

Normal indexing does not collect or publish telemetry. The PyLucene smoke suite invokes + * this method explicitly after indexing, avoiding shared mutable request state. + * + * @return semicolon-delimited writer diagnostics + */ + public String getWriterTelemetry() { + return WriterTelemetry.forHnsw(acceleratedHNSWParams); + } + /** * Returns a KnnVectorsReader to read the vectors from the index. */ diff --git a/src/main/java/com/nvidia/cuvs/lucene/WriterTelemetry.java b/src/main/java/com/nvidia/cuvs/lucene/WriterTelemetry.java new file mode 100644 index 00000000..d4aeef59 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/WriterTelemetry.java @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +/** Builds on-demand diagnostic snapshots for the PyLucene smoke suite. */ +final class WriterTelemetry { + + private WriterTelemetry() {} + + static String forCagra() { + return "writerPath=gpu-cagra"; + } + + static String forHnsw(AcceleratedHNSWParams params) { + String writerPath; + if (ThreadLocalCuVSResourcesProvider.isSupported()) { + writerPath = "gpu-hnsw"; + } else if (ThreadLocalCuVSResourcesProvider.isCpuHnswFallbackForced()) { + writerPath = "cpu-hnsw-fallback"; + } else { + writerPath = "cpu-hnsw-auto-fallback"; + } + + return "writerPath=" + + writerPath + + ";hnswLayers=" + + params.getHnswLayers() + + ";cagraGraphBuildAlgo=" + + params.getCagraGraphBuildAlgo().name() + + ";cagraGraphDegree=" + + params.getGraphdegree() + + ";cagraIntermediateGraphDegree=" + + params.getIntermediateGraphDegree(); + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java b/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java new file mode 100644 index 00000000..d726c8eb --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import org.junit.Test; + +public class TestWriterTelemetry { + + private static final String FORCE_CPU_HNSW_FALLBACK_PROPERTY = "cuvs.lucene.forceCpuHnswFallback"; + + @Test + public void testCagraTelemetryIsComputedOnDemand() { + assertEquals("writerPath=gpu-cagra", WriterTelemetry.forCagra()); + assertNull(System.getProperty("cuvs.lucene.lastCagraWriterPath")); + } + + @Test + public void testForcedCpuHnswTelemetryIsComputedOnDemand() { + String previousValue = System.getProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY); + System.setProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY, "true"); + try { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withHNSWLayer(3) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .build(); + + assertEquals( + "writerPath=cpu-hnsw-fallback;hnswLayers=3;" + + "cagraGraphBuildAlgo=NN_DESCENT;cagraGraphDegree=32;" + + "cagraIntermediateGraphDegree=64", + WriterTelemetry.forHnsw(params)); + assertNull(System.getProperty("cuvs.lucene.lastHnswWriterPath")); + assertNull(System.getProperty("cuvs.lucene.lastHnswLayers")); + } finally { + if (previousValue == null) { + System.clearProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY); + } else { + System.setProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY, previousValue); + } + } + } +} From 0e6bb43a98725c747e2095179d2951a5e27ff50e Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:34:52 +0000 Subject: [PATCH 06/13] Avoid duplicate binary format initialization --- .../cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index 6af2d3db..5e149b49 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java @@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) { public LuceneAcceleratedHNSWBinaryQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getDefaultDelegateCodec()); + super(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } From eeff876abffe8c2187186bb0cf43b9e0bfba9a49 Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:37:17 +0000 Subject: [PATCH 07/13] Cache binary quantized vector formats --- ...ratedHNSWBinaryQuantizedVectorsFormat.java | 53 ++++++++++++++----- 1 file changed, 40 insertions(+), 13 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index d70efc5e..625eaf7f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -28,8 +28,11 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector private static final Logger log = Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class.getName()); private static final int MAX_DIMENSIONS = 4096; + private static final LuceneProvider LUCENE99_PROVIDER = luceneProvider("99"); + private static volatile FlatVectorsFormat flatVectorsFormat; private final AcceleratedHNSWParams acceleratedHNSWParams; + private volatile KnnVectorsFormat fallbackFormat; private static LuceneProvider luceneProvider(String version) { try { @@ -41,13 +44,41 @@ private static LuceneProvider luceneProvider(String version) { } private static FlatVectorsFormat flatVectorsFormat() { - try { - return luceneProvider("102") - .getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); - } catch (Exception e) { - throw new UnsupportedOperationException( - "Binary quantized vectors require Lucene102 vector formats", e); + FlatVectorsFormat current = flatVectorsFormat; + if (current == null) { + synchronized (LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class) { + current = flatVectorsFormat; + if (current == null) { + try { + current = + luceneProvider("102") + .getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + flatVectorsFormat = current; + } catch (Exception e) { + throw new UnsupportedOperationException( + "Binary quantized vectors require Lucene102 vector formats", e); + } + } + } + } + return current; + } + + private KnnVectorsFormat fallbackFormat() throws Exception { + KnnVectorsFormat current = fallbackFormat; + if (current == null) { + synchronized (this) { + current = fallbackFormat; + if (current == null) { + current = + luceneProvider("102") + .getLuceneHnswBinaryQuantizedVectorsFormatInstance( + acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); + fallbackFormat = current; + } + } } + return current; } /** @@ -89,11 +120,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException Level.WARNING, "GPU based indexing not supported, falling back to using the" + " Lucene102HnswBinaryQuantizedVectorsFormat"); - KnnVectorsFormat fallbackFormat = - luceneProvider("102") - .getLuceneHnswBinaryQuantizedVectorsFormatInstance( - acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); - return fallbackFormat.fieldsWriter(state); + return fallbackFormat().fieldsWriter(state); } catch (Exception e) { throw new RuntimeException(e.getMessage()); } @@ -106,8 +133,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { try { - return luceneProvider("99") - .getLuceneHnswVectorsReaderInstance(state, flatVectorsFormat().fieldsReader(state)); + return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance( + state, flatVectorsFormat().fieldsReader(state)); } catch (Exception e) { throw new RuntimeException(e.getMessage()); } From b76e2f522a81225d70ff601146380492fa38635d Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:48:17 +0000 Subject: [PATCH 08/13] Expose writer diagnostics through format descriptions --- examples/pylucene_smoke.py | 7 +++++-- .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 10 +++------- .../Lucene99AcceleratedHNSWVectorsFormat.java | 13 +++---------- .../com/nvidia/cuvs/lucene/TestWriterTelemetry.java | 9 +++++++++ 4 files changed, 20 insertions(+), 19 deletions(-) diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py index c3ad7c5a..b25bc587 100644 --- a/examples/pylucene_smoke.py +++ b/examples/pylucene_smoke.py @@ -545,8 +545,11 @@ def verify_codecs_advertised(codec_class): def writer_telemetry(codec): vector_format = codec.knnVectorsFormat() - method = vector_format.getClass().getMethod("getWriterTelemetry", ()) - payload = str(method.invoke(vector_format, ())) + description = str(vector_format) + _, separator, payload = description.partition("(") + if not separator or not payload.endswith(")"): + raise AssertionError(f"Malformed vector format diagnostics: {description!r}") + payload = payload[:-1] telemetry = {} for item in payload.split(";"): key, separator, value = item.partition("=") diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 257ec501..014d956b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -77,13 +77,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException return new CuVS2510GPUVectorsWriter(state, gpuSearchParams, flatWriter); } - /** - * Returns an on-demand diagnostic snapshot for integration tests. - * - * @return writer diagnostics - */ - public String getWriterTelemetry() { - return WriterTelemetry.forCagra(); + @Override + public String toString() { + return getName() + "(" + WriterTelemetry.forCagra() + ")"; } /** diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index b380eedf..290a0586 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -101,16 +101,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException } } - /** - * Returns an on-demand diagnostic snapshot for integration tests. - * - *

Normal indexing does not collect or publish telemetry. The PyLucene smoke suite invokes - * this method explicitly after indexing, avoiding shared mutable request state. - * - * @return semicolon-delimited writer diagnostics - */ - public String getWriterTelemetry() { - return WriterTelemetry.forHnsw(acceleratedHNSWParams); + @Override + public String toString() { + return getName() + "(" + WriterTelemetry.forHnsw(acceleratedHNSWParams) + ")"; } /** diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java b/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java index d726c8eb..73e21be0 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java @@ -16,6 +16,9 @@ public class TestWriterTelemetry { @Test public void testCagraTelemetryIsComputedOnDemand() { assertEquals("writerPath=gpu-cagra", WriterTelemetry.forCagra()); + assertEquals( + "CuVS2510GPUVectorsFormat(writerPath=gpu-cagra)", + new CuVS2510GPUVectorsFormat().toString()); assertNull(System.getProperty("cuvs.lucene.lastCagraWriterPath")); } @@ -36,6 +39,12 @@ public void testForcedCpuHnswTelemetryIsComputedOnDemand() { + "cagraGraphBuildAlgo=NN_DESCENT;cagraGraphDegree=32;" + "cagraIntermediateGraphDegree=64", WriterTelemetry.forHnsw(params)); + assertEquals( + "Lucene99AcceleratedHNSWVectorsFormat(" + + "writerPath=cpu-hnsw-fallback;hnswLayers=3;" + + "cagraGraphBuildAlgo=NN_DESCENT;cagraGraphDegree=32;" + + "cagraIntermediateGraphDegree=64)", + new Lucene99AcceleratedHNSWVectorsFormat(params).toString()); assertNull(System.getProperty("cuvs.lucene.lastHnswWriterPath")); assertNull(System.getProperty("cuvs.lucene.lastHnswLayers")); } finally { From bf23e24607e366b0bfe44aef9ecc7d64bd25f5ce Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:53:51 +0000 Subject: [PATCH 09/13] Use Lucene 102 binary vector formats --- ...ratedHNSWBinaryQuantizedVectorsFormat.java | 5 +-- .../nvidia/cuvs/lucene/LuceneProvider.java | 31 ++++++++++--------- .../nvidia/cuvs/lucene/TestBackCompat.java | 8 +++-- 3 files changed, 22 insertions(+), 22 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 625eaf7f..5f88b3e8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -13,7 +13,6 @@ import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; -import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; @@ -50,9 +49,7 @@ private static FlatVectorsFormat flatVectorsFormat() { current = flatVectorsFormat; if (current == null) { try { - current = - luceneProvider("102") - .getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + current = luceneProvider("102").getLuceneBinaryQuantizedVectorsFormatInstance(); flatVectorsFormat = current; } catch (Exception e) { throw new UnsupportedOperationException( diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 01db5f68..bfeb20c8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -15,6 +15,7 @@ import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; @@ -107,6 +108,18 @@ public static synchronized LuceneProvider getInstance(String version) } private LuceneProvider(String version) throws ClassNotFoundException { + if ("102".equals(version)) { + binaryQuantizedVectorsFormat = + loadClass( + setVersion(luceneBinaryQuantizedVectorsFormat, version), + setVersion(luceneBinaryQuantizedVectorsFormatFallback, version)); + hnswBinaryQuantizedVectorsFormat = + loadClass( + setVersion(luceneHnswBinaryQuantizedVectorsFormat, version), + setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, version)); + return; + } + flatVectorsFormat = loadClass( setVersion(luceneFlatVectorsFormat, version), @@ -132,18 +145,6 @@ private LuceneProvider(String version) throws ClassNotFoundException { loadClass( setVersion(luceneHnswScalarQuantizedVectorsFormat, version), setVersion(luceneHnswScalarQuantizedVectorsFormatFallback, version)); - - // TODO: Find a better way if possible, but as a separate initiative. - if ("102".equals(version)) { - binaryQuantizedVectorsFormat = - loadClass( - setVersion(luceneBinaryQuantizedVectorsFormat, version), - setVersion(luceneBinaryQuantizedVectorsFormatFallback, version)); - hnswBinaryQuantizedVectorsFormat = - loadClass( - setVersion(luceneHnswBinaryQuantizedVectorsFormat, version), - setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, version)); - } } private static String setVersion(String pkg, String version) { @@ -278,7 +279,7 @@ public List getSimilarityFunctions() } } - public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws Exception { + public FlatVectorsFormat getLuceneBinaryQuantizedVectorsFormatInstance() throws Exception { try { Constructor luceneBinaryQuantizedVectorsFormatConstructor = binaryQuantizedVectorsFormat.getConstructor(); @@ -291,12 +292,12 @@ public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws } } - public FlatVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( + public KnnVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( int maxConn, int beamWidth) throws Exception { try { Constructor luceneHnswBinaryQuantizedVectorsFormatConstructor = hnswBinaryQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); - return (FlatVectorsFormat) + return (KnnVectorsFormat) luceneHnswBinaryQuantizedVectorsFormatConstructor.newInstance(maxConn, beamWidth); } catch (Exception e) { log.log( diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 5c974301..946b4a93 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -48,9 +48,11 @@ public void testProviderCachesSupportedVersion() throws Exception { assertSame(provider99, LuceneProvider.getInstance("99")); } - @Test(expected = ClassNotFoundException.class) - public void testProviderDoesNotPretendLucene102IsACompleteProviderVersion() throws Exception { - LuceneProvider.getInstance("102"); + @Test + public void testProviderSupportsLucene102BinaryFormats() throws Exception { + LuceneProvider provider102 = LuceneProvider.getInstance("102"); + assertNotNull(provider102.getLuceneBinaryQuantizedVectorsFormatInstance()); + assertNotNull(provider102.getLuceneHnswBinaryQuantizedVectorsFormatInstance(16, 100)); } @Test From f7a8c664f2506e18ee6ccd05cd5996d6c0e204aa Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:56:30 +0000 Subject: [PATCH 10/13] Use the standard jar for PyLucene --- README.md | 23 +++++----- ci/test_pylucene_smoke.sh | 43 ++++++++++--------- examples/pylucene_smoke.py | 13 +++--- pom.xml | 1 - .../jar-with-pylucene-dependencies.xml | 15 ------- 5 files changed, 42 insertions(+), 53 deletions(-) delete mode 100644 src/main/assembly/jar-with-pylucene-dependencies.xml diff --git a/README.md b/README.md index 82bc3597..7c77379b 100644 --- a/README.md +++ b/README.md @@ -27,14 +27,14 @@ Because PyLucene's generated Python module only exposes the Java classes it was to wrap, use Lucene's service provider lookup to load `cuvs-lucene` codecs from Python instead of importing `com.nvidia.cuvs.lucene` classes directly. -Build the PyLucene sidecar jar: +Build the standard cuvs-lucene jar: ```sh mvn clean package -DskipTests ``` -Then start PyLucene with the base `cuvs-java` jar, the generated PyLucene -sidecar jar, and PyLucene's own Lucene classpath: +Then start PyLucene with the base `cuvs-java` jar, the standard `cuvs-lucene` +jar, and PyLucene's own Lucene classpath: ```python import os @@ -44,7 +44,10 @@ import lucene cuvs_java_jar = Path(os.environ["CUVS_LUCENE_CUVS_JAVA_JAR"]) cuvs_lucene_jar = next( - Path("target").glob("cuvs-lucene-*-jar-with-pylucene-dependencies.jar") + jar + for jar in Path("target").glob("cuvs-lucene-*.jar") + if "-jar-with-" not in jar.name + and not jar.name.endswith(("-sources.jar", "-javadoc.jar")) ) lucene.initVM( classpath=os.pathsep.join( @@ -62,11 +65,11 @@ codec = Codec.forName("Lucene101AcceleratedHNSWCodec") ``` Use the returned `codec` with `IndexWriterConfig.setCodec(codec)`. The -`jar-with-pylucene-dependencies` artifact includes only `cuvs-lucene` classes and -service descriptors. PyLucene must provide Lucene classes, and the base -multi-release `cuvs-java` jar must be present separately on the JVM classpath. Do -not use a native classifier `cuvs-java` jar here unless you also want to rely on -its embedded native libraries; the base jar uses native libraries from +The standard artifact includes `cuvs-lucene` classes and service descriptors. +PyLucene must provide Lucene classes, and the base multi-release `cuvs-java` jar +must be present separately on the JVM classpath. Do not use a native classifier +`cuvs-java` jar here unless you also want to rely on its embedded native +libraries; the base jar uses native libraries from `LD_LIBRARY_PATH`/`java.library.path`. To run the PyLucene pytest smoke suite against a local PyLucene environment: @@ -93,7 +96,7 @@ cover one-layer and three-layer HNSW graphs built from CAGRA with NN_DESCENT, `graphDegree=32`, and `intermediateGraphDegree=64`. The base matrix uses 2,000 documents and 32 dimensions; high-segment cases use at least 257 rows per segment to avoid expected cuVS graph-degree clamps on tiny per-segment datasets. -The suite checks Lucene SPI discovery, sidecar packaging, index file suffixes +The suite checks Lucene SPI discovery, jar packaging, index file suffixes (`.vex`/`.vem` for HNSW and `.vcag`/`.vemc` for CAGRA), indexed vector metadata, unfiltered KNN, filtered KNN, missing-vector documents, deletions, and force merge behavior. To run a subset or resize the test: diff --git a/ci/test_pylucene_smoke.sh b/ci/test_pylucene_smoke.sh index b6953b97..c7e3fb5f 100755 --- a/ci/test_pylucene_smoke.sh +++ b/ci/test_pylucene_smoke.sh @@ -16,6 +16,7 @@ PYLUCENE_CASES="${CUVS_LUCENE_PYLUCENE_CASES:-}" PYLUCENE_ROWS="${CUVS_LUCENE_PYLUCENE_ROWS:-}" PYLUCENE_DIMS="${CUVS_LUCENE_PYLUCENE_DIMS:-}" PYLUCENE_TOPK="${CUVS_LUCENE_PYLUCENE_TOPK:-}" +CUVS_LUCENE_JAR_PATH="${CUVS_LUCENE_JAR:-}" for arg in "$@"; do case "${arg}" in @@ -40,8 +41,8 @@ for arg in "$@"; do -h|--help) echo "Usage: $0 [--no-build] [--gpu-e2e] [--cases=CASE[,CASE...]] [--rows=N] [--dims=N] [--topk=N]" echo - echo "Builds and checks the PyLucene sidecar jar, then runs the PyLucene pytest suite." - echo "Set CUVS_LUCENE_PYLUCENE_JAR to test an existing sidecar jar." + echo "Builds and checks the standard cuvs-lucene jar, then runs the PyLucene pytest suite." + echo "Set CUVS_LUCENE_JAR to test an existing cuvs-lucene jar." echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar if it is not in ~/.m2." echo "Set PYTHON or MVN to override the Python or Maven executable." echo @@ -76,7 +77,7 @@ require_command() { require_command "${PYTHON_BIN}" require_command jar -if [[ "${SKIP_BUILD}" -eq 0 && -z "${CUVS_LUCENE_PYLUCENE_JAR:-}" ]]; then +if [[ "${SKIP_BUILD}" -eq 0 && -z "${CUVS_LUCENE_JAR_PATH}" ]]; then require_command "${MVN_BIN}" "${MVN_BIN}" clean package -DskipTests fi @@ -90,24 +91,24 @@ if [[ -z "${project_version}" ]]; then exit 1 fi -if [[ -n "${CUVS_LUCENE_PYLUCENE_JAR:-}" ]]; then - sidecar_jar="${CUVS_LUCENE_PYLUCENE_JAR}" +if [[ -n "${CUVS_LUCENE_JAR_PATH}" ]]; then + cuvs_lucene_jar="${CUVS_LUCENE_JAR_PATH}" else - sidecar_jar="target/cuvs-lucene-${project_version}-jar-with-pylucene-dependencies.jar" + cuvs_lucene_jar="target/cuvs-lucene-${project_version}.jar" fi -if [[ ! -f "${sidecar_jar}" ]]; then - echo "PyLucene sidecar jar not found: ${sidecar_jar}" >&2 - echo "Run without --no-build, or set CUVS_LUCENE_PYLUCENE_JAR to an existing jar." >&2 +if [[ ! -f "${cuvs_lucene_jar}" ]]; then + echo "cuvs-lucene jar not found: ${cuvs_lucene_jar}" >&2 + echo "Run without --no-build, or set CUVS_LUCENE_JAR to an existing jar." >&2 exit 1 fi -case "${sidecar_jar}" in +case "${cuvs_lucene_jar}" in /*) - sidecar_jar_abs="${sidecar_jar}" + cuvs_lucene_jar_abs="${cuvs_lucene_jar}" ;; *) - sidecar_jar_abs="${REPO_ROOT}/${sidecar_jar}" + cuvs_lucene_jar_abs="${REPO_ROOT}/${cuvs_lucene_jar}" ;; esac @@ -150,13 +151,13 @@ entries_file="$(mktemp)" services_dir="$(mktemp -d)" trap 'rm -f "${entries_file}"; rm -rf "${services_dir}"' EXIT -jar tf "${sidecar_jar_abs}" >"${entries_file}" +jar tf "${cuvs_lucene_jar_abs}" >"${entries_file}" for service in \ "META-INF/services/org.apache.lucene.codecs.Codec" \ "META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat"; do grep -qx "${service}" "${entries_file}" || { - echo "Missing service descriptor in ${sidecar_jar}: ${service}" >&2 + echo "Missing service descriptor in ${cuvs_lucene_jar}: ${service}" >&2 exit 1 } done @@ -169,7 +170,7 @@ for class_file in \ "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.class" \ "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.class"; do grep -qx "${class_file}" "${entries_file}" || { - echo "Missing cuvs-lucene class in ${sidecar_jar}: ${class_file}" >&2 + echo "Missing cuvs-lucene class in ${cuvs_lucene_jar}: ${class_file}" >&2 exit 1 } done @@ -179,7 +180,7 @@ lucene_entries="$( | grep -v '/$' || true )" if [[ -n "${lucene_entries}" ]]; then - echo "${sidecar_jar} contains org.apache.lucene classes; PyLucene must provide Lucene." >&2 + echo "${cuvs_lucene_jar} contains org.apache.lucene classes; PyLucene must provide Lucene." >&2 echo "${lucene_entries}" >&2 exit 1 fi @@ -190,7 +191,7 @@ flattened_cuvs_entries="$( | grep -v "^com/nvidia/cuvs/lucene/" || true )" if [[ -n "${flattened_cuvs_entries}" ]]; then - echo "${sidecar_jar} contains flattened cuvs-java classes; use the base cuvs-java jar separately." >&2 + echo "${cuvs_lucene_jar} contains flattened cuvs-java classes; use the base cuvs-java jar separately." >&2 echo "${flattened_cuvs_entries}" >&2 exit 1 fi @@ -200,7 +201,7 @@ flattened_multi_release_entries="$( | grep -v '/$' || true )" if [[ -n "${flattened_multi_release_entries}" ]]; then - echo "${sidecar_jar} contains flattened multi-release cuvs-java classes." >&2 + echo "${cuvs_lucene_jar} contains flattened multi-release cuvs-java classes." >&2 echo "${flattened_multi_release_entries}" >&2 exit 1 fi @@ -211,7 +212,7 @@ extra_lucene_services="$( | grep -v "^META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat$" || true )" if [[ -n "${extra_lucene_services}" ]]; then - echo "${sidecar_jar} contains unexpected Lucene service descriptors:" >&2 + echo "${cuvs_lucene_jar} contains unexpected Lucene service descriptors:" >&2 echo "${extra_lucene_services}" >&2 exit 1 fi @@ -219,7 +220,7 @@ fi ( cd "${services_dir}" jar xf \ - "${sidecar_jar_abs}" \ + "${cuvs_lucene_jar_abs}" \ META-INF/services/org.apache.lucene.codecs.Codec \ META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat ) @@ -271,7 +272,7 @@ done } smoke_env=( - "CUVS_LUCENE_PYLUCENE_JAR=${sidecar_jar_abs}" + "CUVS_LUCENE_JAR=${cuvs_lucene_jar_abs}" "CUVS_LUCENE_CUVS_JAVA_JAR=${cuvs_java_jar}" ) diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py index b25bc587..3117141b 100644 --- a/examples/pylucene_smoke.py +++ b/examples/pylucene_smoke.py @@ -132,9 +132,7 @@ def suffixes_from_env(): def find_cuvs_lucene_jar(): - configured = os.environ.get("CUVS_LUCENE_PYLUCENE_JAR") or os.environ.get( - "CUVS_LUCENE_JAR" - ) + configured = os.environ.get("CUVS_LUCENE_JAR") if configured: jar = Path(configured) if not jar.exists(): @@ -142,13 +140,16 @@ def find_cuvs_lucene_jar(): return jar jars = sorted( - (REPO_ROOT / "target").glob( - "cuvs-lucene-*-jar-with-pylucene-dependencies.jar" + jar + for jar in (REPO_ROOT / "target").glob("cuvs-lucene-*.jar") + if not any( + marker in jar.name + for marker in ("-jar-with-", "-sources.jar", "-javadoc.jar") ) ) if not jars: raise FileNotFoundError( - "No PyLucene sidecar jar found under target/. " + "No cuvs-lucene jar found under target/. " "Run `mvn clean package -DskipTests` first." ) return jars[-1] diff --git a/pom.xml b/pom.xml index ad86452f..fc9add8b 100644 --- a/pom.xml +++ b/pom.xml @@ -168,7 +168,6 @@ src/main/assembly/jar-with-dependencies.xml - src/main/assembly/jar-with-pylucene-dependencies.xml diff --git a/src/main/assembly/jar-with-pylucene-dependencies.xml b/src/main/assembly/jar-with-pylucene-dependencies.xml deleted file mode 100644 index de01eda5..00000000 --- a/src/main/assembly/jar-with-pylucene-dependencies.xml +++ /dev/null @@ -1,15 +0,0 @@ - - jar-with-pylucene-dependencies - - jar - - false - - - ${project.build.outputDirectory} - / - - - From 7084157d170e7219a0039dbc6a71c0a87a199eaa Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:58:17 +0000 Subject: [PATCH 11/13] Add a public PyLucene test entrypoint --- README.md | 19 ++- ci/test_pylucene_smoke.sh | 297 +------------------------------------ test_pylucene.sh | 304 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 319 insertions(+), 301 deletions(-) create mode 100755 test_pylucene.sh diff --git a/README.md b/README.md index 7c77379b..36a541e1 100644 --- a/README.md +++ b/README.md @@ -64,8 +64,8 @@ from org.apache.lucene.codecs import Codec codec = Codec.forName("Lucene101AcceleratedHNSWCodec") ``` -Use the returned `codec` with `IndexWriterConfig.setCodec(codec)`. The -The standard artifact includes `cuvs-lucene` classes and service descriptors. +Use the returned `codec` with `IndexWriterConfig.setCodec(codec)`. The standard +artifact includes `cuvs-lucene` classes and service descriptors. PyLucene must provide Lucene classes, and the base multi-release `cuvs-java` jar must be present separately on the JVM classpath. Do not use a native classifier `cuvs-java` jar here unless you also want to rely on its embedded native @@ -75,14 +75,23 @@ libraries; the base jar uses native libraries from To run the PyLucene pytest smoke suite against a local PyLucene environment: ```sh -./ci/test_pylucene_smoke.sh +./test_pylucene.sh +``` + +The script builds and validates the jar before invoking pytest. To invoke pytest +directly against existing artifacts instead: + +```sh +CUVS_LUCENE_JAR=target/cuvs-lucene-26.08.0.jar \ +CUVS_LUCENE_CUVS_JAVA_JAR=/path/to/cuvs-java-26.08.0.jar \ +python3 -m pytest -q -s examples/test_pylucene_smoke.py ``` To run an expanded GPU end-to-end pytest suite through CPU HNSW, CAGRA-to-HNSW, and CAGRA search paths: ```sh -./ci/test_pylucene_smoke.sh --gpu-e2e +./test_pylucene.sh --gpu-e2e ``` The expanded suite runs the `gpu-basic`, `gpu-segments`, `cpu-hnsw`, and @@ -102,7 +111,7 @@ unfiltered KNN, filtered KNN, missing-vector documents, deletions, and force merge behavior. To run a subset or resize the test: ```sh -./ci/test_pylucene_smoke.sh --gpu-e2e --cases=gpu-segments --rows=5000 --dims=64 --topk=20 +./test_pylucene.sh --gpu-e2e --cases=gpu-segments --rows=5000 --dims=64 --topk=20 ``` ### Running Tests diff --git a/ci/test_pylucene_smoke.sh b/ci/test_pylucene_smoke.sh index c7e3fb5f..79222be2 100755 --- a/ci/test_pylucene_smoke.sh +++ b/ci/test_pylucene_smoke.sh @@ -6,299 +6,4 @@ set -euo pipefail REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -cd "${REPO_ROOT}" - -MVN_BIN="${MVN:-mvn}" -PYTHON_BIN="${PYTHON:-python3}" -SKIP_BUILD=0 -GPU_E2E=0 -PYLUCENE_CASES="${CUVS_LUCENE_PYLUCENE_CASES:-}" -PYLUCENE_ROWS="${CUVS_LUCENE_PYLUCENE_ROWS:-}" -PYLUCENE_DIMS="${CUVS_LUCENE_PYLUCENE_DIMS:-}" -PYLUCENE_TOPK="${CUVS_LUCENE_PYLUCENE_TOPK:-}" -CUVS_LUCENE_JAR_PATH="${CUVS_LUCENE_JAR:-}" - -for arg in "$@"; do - case "${arg}" in - --gpu-e2e) - GPU_E2E=1 - ;; - --no-build) - SKIP_BUILD=1 - ;; - --cases=*) - PYLUCENE_CASES="${arg#--cases=}" - ;; - --rows=*) - PYLUCENE_ROWS="${arg#--rows=}" - ;; - --dims=*) - PYLUCENE_DIMS="${arg#--dims=}" - ;; - --topk=*) - PYLUCENE_TOPK="${arg#--topk=}" - ;; - -h|--help) - echo "Usage: $0 [--no-build] [--gpu-e2e] [--cases=CASE[,CASE...]] [--rows=N] [--dims=N] [--topk=N]" - echo - echo "Builds and checks the standard cuvs-lucene jar, then runs the PyLucene pytest suite." - echo "Set CUVS_LUCENE_JAR to test an existing cuvs-lucene jar." - echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar if it is not in ~/.m2." - echo "Set PYTHON or MVN to override the Python or Maven executable." - echo - echo "Case groups: gpu-basic, gpu-segments, cpu-hnsw, cagra-hnsw, algorithm-matrix, all." - echo "Core cases: smoke, hnsw, cagra, hnsw-single, cagra-single." - echo "Segment cases: hnsw-1seg, cagra-1seg, hnsw-10seg, cagra-10seg," - echo " hnsw-10seg-force-1, cagra-10seg-force-1," - echo " hnsw-100seg-force-10, cagra-100seg-force-10." - echo "CPU HNSW cases: hnsw-cpu, hnsw-cpu-single, hnsw-cpu-1seg," - echo " hnsw-cpu-10seg, hnsw-cpu-10seg-force-1," - echo " hnsw-cpu-100seg-force-10." - echo "CAGRA-to-HNSW cases: cagra-hnsw-1layer, cagra-hnsw-3layer." - echo "--gpu-e2e defaults to cases=all, base rows=2000, dims=32, topk=20." - echo "High-segment cases use at least 257 rows per segment to avoid expected cuVS graph-degree clamps." - echo "--gpu-e2e requires a PyLucene environment plus cuVS native support on a GPU machine." - exit 0 - ;; - *) - echo "Unknown argument: ${arg}" >&2 - exit 2 - ;; - esac -done - -require_command() { - if ! command -v "$1" >/dev/null 2>&1; then - echo "Required command not found: $1" >&2 - exit 127 - fi -} - -require_command "${PYTHON_BIN}" -require_command jar - -if [[ "${SKIP_BUILD}" -eq 0 && -z "${CUVS_LUCENE_JAR_PATH}" ]]; then - require_command "${MVN_BIN}" - "${MVN_BIN}" clean package -DskipTests -fi - -project_version="$( - sed -n 's/.*CUVS_LUCENE#VERSION_UPDATE_MARKER_START-->\([^<]*\)<\/version>.*/\1/p' pom.xml \ - | head -n 1 -)" -if [[ -z "${project_version}" ]]; then - echo "Unable to determine project version from pom.xml" >&2 - exit 1 -fi - -if [[ -n "${CUVS_LUCENE_JAR_PATH}" ]]; then - cuvs_lucene_jar="${CUVS_LUCENE_JAR_PATH}" -else - cuvs_lucene_jar="target/cuvs-lucene-${project_version}.jar" -fi - -if [[ ! -f "${cuvs_lucene_jar}" ]]; then - echo "cuvs-lucene jar not found: ${cuvs_lucene_jar}" >&2 - echo "Run without --no-build, or set CUVS_LUCENE_JAR to an existing jar." >&2 - exit 1 -fi - -case "${cuvs_lucene_jar}" in - /*) - cuvs_lucene_jar_abs="${cuvs_lucene_jar}" - ;; - *) - cuvs_lucene_jar_abs="${REPO_ROOT}/${cuvs_lucene_jar}" - ;; -esac - -find_cuvs_java_jar() { - if [[ -n "${CUVS_LUCENE_CUVS_JAVA_JAR:-}" ]]; then - printf '%s\n' "${CUVS_LUCENE_CUVS_JAVA_JAR}" - return - fi - - local m2_base="${HOME}/.m2/repository/com/nvidia/cuvs/cuvs-java/${project_version}" - local jar="${m2_base}/cuvs-java-${project_version}.jar" - if [[ -f "${jar}" ]]; then - printf '%s\n' "${jar}" - return - fi - - local m2_repo="${HOME}/.m2/repository/com/nvidia/cuvs/cuvs-java" - if [[ ! -d "${m2_repo}" ]]; then - return - fi - - find "${m2_repo}" \ - -type f \ - -name 'cuvs-java-*.jar' \ - ! -name '*sources*' \ - ! -name '*javadoc*' \ - ! -name '*x86_64*' \ - | sort -V \ - | tail -n 1 -} - -cuvs_java_jar="$(find_cuvs_java_jar)" -if [[ -z "${cuvs_java_jar}" || ! -f "${cuvs_java_jar}" ]]; then - echo "Base cuvs-java jar not found." >&2 - echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar, not a native classifier jar." >&2 - exit 1 -fi - -entries_file="$(mktemp)" -services_dir="$(mktemp -d)" -trap 'rm -f "${entries_file}"; rm -rf "${services_dir}"' EXIT - -jar tf "${cuvs_lucene_jar_abs}" >"${entries_file}" - -for service in \ - "META-INF/services/org.apache.lucene.codecs.Codec" \ - "META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat"; do - grep -qx "${service}" "${entries_file}" || { - echo "Missing service descriptor in ${cuvs_lucene_jar}: ${service}" >&2 - exit 1 - } -done - -for class_file in \ - "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.class" \ - "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.class" \ - "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.class" \ - "com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.class" \ - "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.class" \ - "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.class"; do - grep -qx "${class_file}" "${entries_file}" || { - echo "Missing cuvs-lucene class in ${cuvs_lucene_jar}: ${class_file}" >&2 - exit 1 - } -done - -lucene_entries="$( - grep "^org/apache/lucene/" "${entries_file}" \ - | grep -v '/$' || true -)" -if [[ -n "${lucene_entries}" ]]; then - echo "${cuvs_lucene_jar} contains org.apache.lucene classes; PyLucene must provide Lucene." >&2 - echo "${lucene_entries}" >&2 - exit 1 -fi - -flattened_cuvs_entries="$( - grep "^com/nvidia/cuvs/" "${entries_file}" \ - | grep -v '/$' \ - | grep -v "^com/nvidia/cuvs/lucene/" || true -)" -if [[ -n "${flattened_cuvs_entries}" ]]; then - echo "${cuvs_lucene_jar} contains flattened cuvs-java classes; use the base cuvs-java jar separately." >&2 - echo "${flattened_cuvs_entries}" >&2 - exit 1 -fi - -flattened_multi_release_entries="$( - grep "^META-INF/versions/.*/com/nvidia/cuvs/" "${entries_file}" \ - | grep -v '/$' || true -)" -if [[ -n "${flattened_multi_release_entries}" ]]; then - echo "${cuvs_lucene_jar} contains flattened multi-release cuvs-java classes." >&2 - echo "${flattened_multi_release_entries}" >&2 - exit 1 -fi - -extra_lucene_services="$( - grep "^META-INF/services/org.apache.lucene." "${entries_file}" \ - | grep -v "^META-INF/services/org.apache.lucene.codecs.Codec$" \ - | grep -v "^META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat$" || true -)" -if [[ -n "${extra_lucene_services}" ]]; then - echo "${cuvs_lucene_jar} contains unexpected Lucene service descriptors:" >&2 - echo "${extra_lucene_services}" >&2 - exit 1 -fi - -( - cd "${services_dir}" - jar xf \ - "${cuvs_lucene_jar_abs}" \ - META-INF/services/org.apache.lucene.codecs.Codec \ - META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat -) - -codec_descriptor="${services_dir}/META-INF/services/org.apache.lucene.codecs.Codec" -format_descriptor="${services_dir}/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat" - -for descriptor in "${codec_descriptor}" "${format_descriptor}"; do - if grep -q "^org\\.apache\\.lucene\\." "${descriptor}"; then - echo "${descriptor#"${services_dir}"/} advertises Lucene-owned providers." >&2 - exit 1 - fi -done - -for provider in \ - "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec" \ - "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWBaseLayerCodec" \ - "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWMultiLayerCodec" \ - "com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec" \ - "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedCodec" \ - "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedCodec"; do - grep -qx "${provider}" "${codec_descriptor}" || { - echo "Codec service descriptor missing provider: ${provider}" >&2 - exit 1 - } -done - -for provider in \ - "com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat" \ - "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat" \ - "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat" \ - "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedVectorsFormat"; do - grep -qx "${provider}" "${format_descriptor}" || { - echo "KnnVectorsFormat service descriptor missing provider: ${provider}" >&2 - exit 1 - } -done - -"${PYTHON_BIN}" -c "import lucene" >/dev/null 2>&1 || { - echo "Python cannot import PyLucene's lucene module." >&2 - echo "Activate or install a PyLucene environment compatible with this project's Lucene version." >&2 - exit 1 -} - -"${PYTHON_BIN}" -m pytest --version >/dev/null 2>&1 || { - echo "Python cannot run pytest." >&2 - echo "Install pytest in the active PyLucene environment." >&2 - exit 1 -} - -smoke_env=( - "CUVS_LUCENE_JAR=${cuvs_lucene_jar_abs}" - "CUVS_LUCENE_CUVS_JAVA_JAR=${cuvs_java_jar}" -) - -if [[ "${GPU_E2E}" -eq 1 ]]; then - smoke_env+=( - "CUVS_LUCENE_PYLUCENE_CASES=${PYLUCENE_CASES:-all}" - "CUVS_LUCENE_PYLUCENE_ROWS=${PYLUCENE_ROWS:-2000}" - "CUVS_LUCENE_PYLUCENE_DIMS=${PYLUCENE_DIMS:-32}" - "CUVS_LUCENE_PYLUCENE_TOPK=${PYLUCENE_TOPK:-20}" - "CUVS_LUCENE_REQUIRE_CUVS=1" - "CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}" - ) -else - if [[ -n "${PYLUCENE_CASES}" ]]; then - smoke_env+=("CUVS_LUCENE_PYLUCENE_CASES=${PYLUCENE_CASES}") - fi - if [[ -n "${PYLUCENE_ROWS}" ]]; then - smoke_env+=("CUVS_LUCENE_PYLUCENE_ROWS=${PYLUCENE_ROWS}") - fi - if [[ -n "${PYLUCENE_DIMS}" ]]; then - smoke_env+=("CUVS_LUCENE_PYLUCENE_DIMS=${PYLUCENE_DIMS}") - fi - if [[ -n "${PYLUCENE_TOPK}" ]]; then - smoke_env+=("CUVS_LUCENE_PYLUCENE_TOPK=${PYLUCENE_TOPK}") - fi - smoke_env+=("CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}") -fi - -env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s examples/test_pylucene_smoke.py +exec "${REPO_ROOT}/test_pylucene.sh" "$@" diff --git a/test_pylucene.sh b/test_pylucene.sh new file mode 100755 index 00000000..7fb38e51 --- /dev/null +++ b/test_pylucene.sh @@ -0,0 +1,304 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "${REPO_ROOT}" + +MVN_BIN="${MVN:-mvn}" +PYTHON_BIN="${PYTHON:-python3}" +SKIP_BUILD=0 +GPU_E2E=0 +PYLUCENE_CASES="${CUVS_LUCENE_PYLUCENE_CASES:-}" +PYLUCENE_ROWS="${CUVS_LUCENE_PYLUCENE_ROWS:-}" +PYLUCENE_DIMS="${CUVS_LUCENE_PYLUCENE_DIMS:-}" +PYLUCENE_TOPK="${CUVS_LUCENE_PYLUCENE_TOPK:-}" +CUVS_LUCENE_JAR_PATH="${CUVS_LUCENE_JAR:-}" + +for arg in "$@"; do + case "${arg}" in + --gpu-e2e) + GPU_E2E=1 + ;; + --no-build) + SKIP_BUILD=1 + ;; + --cases=*) + PYLUCENE_CASES="${arg#--cases=}" + ;; + --rows=*) + PYLUCENE_ROWS="${arg#--rows=}" + ;; + --dims=*) + PYLUCENE_DIMS="${arg#--dims=}" + ;; + --topk=*) + PYLUCENE_TOPK="${arg#--topk=}" + ;; + -h|--help) + echo "Usage: ./test_pylucene.sh [--no-build] [--gpu-e2e] [--cases=CASE[,CASE...]] [--rows=N] [--dims=N] [--topk=N]" + echo + echo "Builds and checks the standard cuvs-lucene jar, then runs the PyLucene pytest suite." + echo "Set CUVS_LUCENE_JAR to test an existing cuvs-lucene jar." + echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar if it is not in ~/.m2." + echo "Set PYTHON or MVN to override the Python or Maven executable." + echo + echo "Case groups: gpu-basic, gpu-segments, cpu-hnsw, cagra-hnsw, algorithm-matrix, all." + echo "Core cases: smoke, hnsw, cagra, hnsw-single, cagra-single." + echo "Segment cases: hnsw-1seg, cagra-1seg, hnsw-10seg, cagra-10seg," + echo " hnsw-10seg-force-1, cagra-10seg-force-1," + echo " hnsw-100seg-force-10, cagra-100seg-force-10." + echo "CPU HNSW cases: hnsw-cpu, hnsw-cpu-single, hnsw-cpu-1seg," + echo " hnsw-cpu-10seg, hnsw-cpu-10seg-force-1," + echo " hnsw-cpu-100seg-force-10." + echo "CAGRA-to-HNSW cases: cagra-hnsw-1layer, cagra-hnsw-3layer." + echo "--gpu-e2e defaults to cases=all, base rows=2000, dims=32, topk=20." + echo "High-segment cases use at least 257 rows per segment to avoid expected cuVS graph-degree clamps." + echo "--gpu-e2e requires a PyLucene environment plus cuVS native support on a GPU machine." + exit 0 + ;; + *) + echo "Unknown argument: ${arg}" >&2 + exit 2 + ;; + esac +done + +require_command() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "Required command not found: $1" >&2 + exit 127 + fi +} + +require_command "${PYTHON_BIN}" +require_command jar + +if [[ "${SKIP_BUILD}" -eq 0 && -z "${CUVS_LUCENE_JAR_PATH}" ]]; then + require_command "${MVN_BIN}" + "${MVN_BIN}" clean package -DskipTests +fi + +project_version="$( + sed -n 's/.*CUVS_LUCENE#VERSION_UPDATE_MARKER_START-->\([^<]*\)<\/version>.*/\1/p' pom.xml \ + | head -n 1 +)" +if [[ -z "${project_version}" ]]; then + echo "Unable to determine project version from pom.xml" >&2 + exit 1 +fi + +if [[ -n "${CUVS_LUCENE_JAR_PATH}" ]]; then + cuvs_lucene_jar="${CUVS_LUCENE_JAR_PATH}" +else + cuvs_lucene_jar="target/cuvs-lucene-${project_version}.jar" +fi + +if [[ ! -f "${cuvs_lucene_jar}" ]]; then + echo "cuvs-lucene jar not found: ${cuvs_lucene_jar}" >&2 + echo "Run without --no-build, or set CUVS_LUCENE_JAR to an existing jar." >&2 + exit 1 +fi + +case "${cuvs_lucene_jar}" in + /*) + cuvs_lucene_jar_abs="${cuvs_lucene_jar}" + ;; + *) + cuvs_lucene_jar_abs="${REPO_ROOT}/${cuvs_lucene_jar}" + ;; +esac + +find_cuvs_java_jar() { + if [[ -n "${CUVS_LUCENE_CUVS_JAVA_JAR:-}" ]]; then + printf '%s\n' "${CUVS_LUCENE_CUVS_JAVA_JAR}" + return + fi + + local m2_base="${HOME}/.m2/repository/com/nvidia/cuvs/cuvs-java/${project_version}" + local jar="${m2_base}/cuvs-java-${project_version}.jar" + if [[ -f "${jar}" ]]; then + printf '%s\n' "${jar}" + return + fi + + local m2_repo="${HOME}/.m2/repository/com/nvidia/cuvs/cuvs-java" + if [[ ! -d "${m2_repo}" ]]; then + return + fi + + find "${m2_repo}" \ + -type f \ + -name 'cuvs-java-*.jar' \ + ! -name '*sources*' \ + ! -name '*javadoc*' \ + ! -name '*x86_64*' \ + | sort -V \ + | tail -n 1 +} + +cuvs_java_jar="$(find_cuvs_java_jar)" +if [[ -z "${cuvs_java_jar}" || ! -f "${cuvs_java_jar}" ]]; then + echo "Base cuvs-java jar not found." >&2 + echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar, not a native classifier jar." >&2 + exit 1 +fi + +entries_file="$(mktemp)" +services_dir="$(mktemp -d)" +trap 'rm -f "${entries_file}"; rm -rf "${services_dir}"' EXIT + +jar tf "${cuvs_lucene_jar_abs}" >"${entries_file}" + +for service in \ + "META-INF/services/org.apache.lucene.codecs.Codec" \ + "META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat"; do + grep -qx "${service}" "${entries_file}" || { + echo "Missing service descriptor in ${cuvs_lucene_jar}: ${service}" >&2 + exit 1 + } +done + +for class_file in \ + "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.class" \ + "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWBaseLayerCodec.class" \ + "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWMultiLayerCodec.class" \ + "com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.class" \ + "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.class" \ + "com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.class"; do + grep -qx "${class_file}" "${entries_file}" || { + echo "Missing cuvs-lucene class in ${cuvs_lucene_jar}: ${class_file}" >&2 + exit 1 + } +done + +lucene_entries="$( + grep "^org/apache/lucene/" "${entries_file}" \ + | grep -v '/$' || true +)" +if [[ -n "${lucene_entries}" ]]; then + echo "${cuvs_lucene_jar} contains org.apache.lucene classes; PyLucene must provide Lucene." >&2 + echo "${lucene_entries}" >&2 + exit 1 +fi + +flattened_cuvs_entries="$( + grep "^com/nvidia/cuvs/" "${entries_file}" \ + | grep -v '/$' \ + | grep -v "^com/nvidia/cuvs/lucene/" || true +)" +if [[ -n "${flattened_cuvs_entries}" ]]; then + echo "${cuvs_lucene_jar} contains flattened cuvs-java classes; use the base cuvs-java jar separately." >&2 + echo "${flattened_cuvs_entries}" >&2 + exit 1 +fi + +flattened_multi_release_entries="$( + grep "^META-INF/versions/.*/com/nvidia/cuvs/" "${entries_file}" \ + | grep -v '/$' || true +)" +if [[ -n "${flattened_multi_release_entries}" ]]; then + echo "${cuvs_lucene_jar} contains flattened multi-release cuvs-java classes." >&2 + echo "${flattened_multi_release_entries}" >&2 + exit 1 +fi + +extra_lucene_services="$( + grep "^META-INF/services/org.apache.lucene." "${entries_file}" \ + | grep -v "^META-INF/services/org.apache.lucene.codecs.Codec$" \ + | grep -v "^META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat$" || true +)" +if [[ -n "${extra_lucene_services}" ]]; then + echo "${cuvs_lucene_jar} contains unexpected Lucene service descriptors:" >&2 + echo "${extra_lucene_services}" >&2 + exit 1 +fi + +( + cd "${services_dir}" + jar xf \ + "${cuvs_lucene_jar_abs}" \ + META-INF/services/org.apache.lucene.codecs.Codec \ + META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +) + +codec_descriptor="${services_dir}/META-INF/services/org.apache.lucene.codecs.Codec" +format_descriptor="${services_dir}/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat" + +for descriptor in "${codec_descriptor}" "${format_descriptor}"; do + if grep -q "^org\\.apache\\.lucene\\." "${descriptor}"; then + echo "${descriptor#"${services_dir}"/} advertises Lucene-owned providers." >&2 + exit 1 + fi +done + +for provider in \ + "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec" \ + "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWBaseLayerCodec" \ + "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWMultiLayerCodec" \ + "com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedCodec" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedCodec"; do + grep -qx "${provider}" "${codec_descriptor}" || { + echo "Codec service descriptor missing provider: ${provider}" >&2 + exit 1 + } +done + +for provider in \ + "com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat" \ + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat" \ + "com.nvidia.cuvs.lucene.LuceneAcceleratedHNSWScalarQuantizedVectorsFormat"; do + grep -qx "${provider}" "${format_descriptor}" || { + echo "KnnVectorsFormat service descriptor missing provider: ${provider}" >&2 + exit 1 + } +done + +"${PYTHON_BIN}" -c "import lucene" >/dev/null 2>&1 || { + echo "Python cannot import PyLucene's lucene module." >&2 + echo "Activate or install a PyLucene environment compatible with this project's Lucene version." >&2 + exit 1 +} + +"${PYTHON_BIN}" -m pytest --version >/dev/null 2>&1 || { + echo "Python cannot run pytest." >&2 + echo "Install pytest in the active PyLucene environment." >&2 + exit 1 +} + +smoke_env=( + "CUVS_LUCENE_JAR=${cuvs_lucene_jar_abs}" + "CUVS_LUCENE_CUVS_JAVA_JAR=${cuvs_java_jar}" +) + +if [[ "${GPU_E2E}" -eq 1 ]]; then + smoke_env+=( + "CUVS_LUCENE_PYLUCENE_CASES=${PYLUCENE_CASES:-all}" + "CUVS_LUCENE_PYLUCENE_ROWS=${PYLUCENE_ROWS:-2000}" + "CUVS_LUCENE_PYLUCENE_DIMS=${PYLUCENE_DIMS:-32}" + "CUVS_LUCENE_PYLUCENE_TOPK=${PYLUCENE_TOPK:-20}" + "CUVS_LUCENE_REQUIRE_CUVS=1" + "CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}" + ) +else + if [[ -n "${PYLUCENE_CASES}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_CASES=${PYLUCENE_CASES}") + fi + if [[ -n "${PYLUCENE_ROWS}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_ROWS=${PYLUCENE_ROWS}") + fi + if [[ -n "${PYLUCENE_DIMS}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_DIMS=${PYLUCENE_DIMS}") + fi + if [[ -n "${PYLUCENE_TOPK}" ]]; then + smoke_env+=("CUVS_LUCENE_PYLUCENE_TOPK=${PYLUCENE_TOPK}") + fi + smoke_env+=("CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}") +fi + +env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s examples/test_pylucene_smoke.py From d13184dc25f242d7f54976226fe925651674af00 Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 02:59:20 +0000 Subject: [PATCH 12/13] Move PyLucene tests under examples Python --- README.md | 2 +- examples/{ => Python}/pylucene_smoke.py | 2 +- examples/{ => Python}/test_pylucene_smoke.py | 0 test_pylucene.sh | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename examples/{ => Python}/pylucene_smoke.py (99%) rename examples/{ => Python}/test_pylucene_smoke.py (100%) diff --git a/README.md b/README.md index 36a541e1..04251424 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,7 @@ directly against existing artifacts instead: ```sh CUVS_LUCENE_JAR=target/cuvs-lucene-26.08.0.jar \ CUVS_LUCENE_CUVS_JAVA_JAR=/path/to/cuvs-java-26.08.0.jar \ -python3 -m pytest -q -s examples/test_pylucene_smoke.py +python3 -m pytest -q -s examples/Python/test_pylucene_smoke.py ``` To run an expanded GPU end-to-end pytest suite through CPU HNSW, diff --git a/examples/pylucene_smoke.py b/examples/Python/pylucene_smoke.py similarity index 99% rename from examples/pylucene_smoke.py rename to examples/Python/pylucene_smoke.py index 3117141b..48fc36ec 100644 --- a/examples/pylucene_smoke.py +++ b/examples/Python/pylucene_smoke.py @@ -10,7 +10,7 @@ import lucene -REPO_ROOT = Path(__file__).resolve().parents[1] +REPO_ROOT = Path(__file__).resolve().parents[2] HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" CAGRA_HNSW_BASE_LAYER_CODEC = "Lucene101AcceleratedHNSWBaseLayerCodec" CAGRA_HNSW_MULTI_LAYER_CODEC = "Lucene101AcceleratedHNSWMultiLayerCodec" diff --git a/examples/test_pylucene_smoke.py b/examples/Python/test_pylucene_smoke.py similarity index 100% rename from examples/test_pylucene_smoke.py rename to examples/Python/test_pylucene_smoke.py diff --git a/test_pylucene.sh b/test_pylucene.sh index 7fb38e51..92880689 100755 --- a/test_pylucene.sh +++ b/test_pylucene.sh @@ -301,4 +301,4 @@ else smoke_env+=("CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}") fi -env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s examples/test_pylucene_smoke.py +env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s examples/Python/test_pylucene_smoke.py From df836ccedf8dd87c18dd351ea0ce6c97e02f7ed6 Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Tue, 21 Jul 2026 04:20:57 +0000 Subject: [PATCH 13/13] Clarify binary format version handling --- README.md | 4 +- ...ratedHNSWBinaryQuantizedVectorsFormat.java | 57 ++++++++++--------- .../nvidia/cuvs/lucene/LuceneProvider.java | 12 ++-- .../nvidia/cuvs/lucene/TestBackCompat.java | 16 ++++-- 4 files changed, 50 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 04251424..78c8a683 100644 --- a/README.md +++ b/README.md @@ -82,8 +82,8 @@ The script builds and validates the jar before invoking pytest. To invoke pytest directly against existing artifacts instead: ```sh -CUVS_LUCENE_JAR=target/cuvs-lucene-26.08.0.jar \ -CUVS_LUCENE_CUVS_JAVA_JAR=/path/to/cuvs-java-26.08.0.jar \ +CUVS_LUCENE_JAR=/path/to/cuvs-lucene.jar \ +CUVS_LUCENE_CUVS_JAVA_JAR=/path/to/cuvs-java.jar \ python3 -m pytest -q -s examples/Python/test_pylucene_smoke.py ``` diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 5f88b3e8..a59cecb1 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -27,13 +27,14 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector private static final Logger log = Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class.getName()); private static final int MAX_DIMENSIONS = 4096; - private static final LuceneProvider LUCENE99_PROVIDER = luceneProvider("99"); - private static volatile FlatVectorsFormat flatVectorsFormat; + private static final LuceneProvider LUCENE_99_PROVIDER = + getLuceneProvider(LuceneProvider.LUCENE_99_FORMAT_VERSION); + private static volatile FlatVectorsFormat cachedFlatVectorsFormat; private final AcceleratedHNSWParams acceleratedHNSWParams; - private volatile KnnVectorsFormat fallbackFormat; + private volatile KnnVectorsFormat cachedFallbackFormat; - private static LuceneProvider luceneProvider(String version) { + private static LuceneProvider getLuceneProvider(String version) { try { return LuceneProvider.getInstance(version); } catch (Exception e) { @@ -42,15 +43,17 @@ private static LuceneProvider luceneProvider(String version) { } } - private static FlatVectorsFormat flatVectorsFormat() { - FlatVectorsFormat current = flatVectorsFormat; - if (current == null) { + private static FlatVectorsFormat getOrCreateFlatVectorsFormat() { + FlatVectorsFormat format = cachedFlatVectorsFormat; + if (format == null) { synchronized (LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class) { - current = flatVectorsFormat; - if (current == null) { + format = cachedFlatVectorsFormat; + if (format == null) { try { - current = luceneProvider("102").getLuceneBinaryQuantizedVectorsFormatInstance(); - flatVectorsFormat = current; + format = + getLuceneProvider(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION) + .getLuceneBinaryQuantizedVectorsFormatInstance(); + cachedFlatVectorsFormat = format; } catch (Exception e) { throw new UnsupportedOperationException( "Binary quantized vectors require Lucene102 vector formats", e); @@ -58,24 +61,24 @@ private static FlatVectorsFormat flatVectorsFormat() { } } } - return current; + return format; } - private KnnVectorsFormat fallbackFormat() throws Exception { - KnnVectorsFormat current = fallbackFormat; - if (current == null) { + private KnnVectorsFormat getOrCreateFallbackFormat() throws Exception { + KnnVectorsFormat format = cachedFallbackFormat; + if (format == null) { synchronized (this) { - current = fallbackFormat; - if (current == null) { - current = - luceneProvider("102") + format = cachedFallbackFormat; + if (format == null) { + format = + getLuceneProvider(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION) .getLuceneHnswBinaryQuantizedVectorsFormatInstance( acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); - fallbackFormat = current; + cachedFallbackFormat = format; } } } - return current; + return format; } /** @@ -103,7 +106,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - var flatWriter = flatVectorsFormat().fieldsWriter(state); + var flatWriter = getOrCreateFlatVectorsFormat().fieldsWriter(state); if (isSupported()) { log.log( Level.FINE, @@ -117,9 +120,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException Level.WARNING, "GPU based indexing not supported, falling back to using the" + " Lucene102HnswBinaryQuantizedVectorsFormat"); - return fallbackFormat().fieldsWriter(state); + return getOrCreateFallbackFormat().fieldsWriter(state); } catch (Exception e) { - throw new RuntimeException(e.getMessage()); + throw new IOException("Unable to initialize the binary quantized fallback writer", e); } } } @@ -130,10 +133,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { try { - return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance( - state, flatVectorsFormat().fieldsReader(state)); + return LUCENE_99_PROVIDER.getLuceneHnswVectorsReaderInstance( + state, getOrCreateFlatVectorsFormat().fieldsReader(state)); } catch (Exception e) { - throw new RuntimeException(e.getMessage()); + throw new IOException("Unable to initialize the binary quantized vectors reader", e); } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index bfeb20c8..8dfaf52d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -35,7 +35,11 @@ public class LuceneProvider { static final Logger log = Logger.getLogger(LuceneProvider.class.getName()); - private static final List SUPPORTED_DELEGATE_CODEC_VERSIONS = List.of("101", "99"); + static final String LUCENE_99_FORMAT_VERSION = "99"; + static final String LUCENE_102_BINARY_FORMAT_VERSION = "102"; + + private static final List SUPPORTED_DELEGATE_CODEC_VERSIONS = + List.of("101", LUCENE_99_FORMAT_VERSION); private static final String BASE = "org.apache.lucene."; private static String codecs = "codecs.lucene."; @@ -108,7 +112,7 @@ public static synchronized LuceneProvider getInstance(String version) } private LuceneProvider(String version) throws ClassNotFoundException { - if ("102".equals(version)) { + if (LUCENE_102_BINARY_FORMAT_VERSION.equals(version)) { binaryQuantizedVectorsFormat = loadClass( setVersion(luceneBinaryQuantizedVectorsFormat, version), @@ -296,13 +300,13 @@ public KnnVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( int maxConn, int beamWidth) throws Exception { try { Constructor luceneHnswBinaryQuantizedVectorsFormatConstructor = - hnswBinaryQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); + hnswBinaryQuantizedVectorsFormat.getConstructor(int.class, int.class); return (KnnVectorsFormat) luceneHnswBinaryQuantizedVectorsFormatConstructor.newInstance(maxConn, beamWidth); } catch (Exception e) { log.log( Level.SEVERE, - "Unable to initialize LuceneBinaryQuantizedVectorsFormat: " + e.getMessage()); + "Unable to initialize LuceneHnswBinaryQuantizedVectorsFormat: " + e.getMessage()); throw e; } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 946b4a93..4b152ab0 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -36,7 +36,7 @@ public void testNonexistentCodec() throws Exception { @Test public void testExistingComponents() throws Exception { - LuceneProvider provider = LuceneProvider.getInstance("99"); + LuceneProvider provider = LuceneProvider.getInstance(LuceneProvider.LUCENE_99_FORMAT_VERSION); assertTrue(provider.getLuceneFlatVectorsFormatInstance(null) instanceof FlatVectorsFormat); assertEquals(provider.getStaticIntParam("VERSION_CURRENT"), 0); assertNotEquals(provider.getSimilarityFunctions().size(), 0); @@ -44,15 +44,19 @@ public void testExistingComponents() throws Exception { @Test public void testProviderCachesSupportedVersion() throws Exception { - LuceneProvider provider99 = LuceneProvider.getInstance("99"); - assertSame(provider99, LuceneProvider.getInstance("99")); + LuceneProvider lucene99Provider = + LuceneProvider.getInstance(LuceneProvider.LUCENE_99_FORMAT_VERSION); + assertSame( + lucene99Provider, LuceneProvider.getInstance(LuceneProvider.LUCENE_99_FORMAT_VERSION)); } @Test public void testProviderSupportsLucene102BinaryFormats() throws Exception { - LuceneProvider provider102 = LuceneProvider.getInstance("102"); - assertNotNull(provider102.getLuceneBinaryQuantizedVectorsFormatInstance()); - assertNotNull(provider102.getLuceneHnswBinaryQuantizedVectorsFormatInstance(16, 100)); + LuceneProvider lucene102BinaryFormatProvider = + LuceneProvider.getInstance(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION); + assertNotNull(lucene102BinaryFormatProvider.getLuceneBinaryQuantizedVectorsFormatInstance()); + assertNotNull( + lucene102BinaryFormatProvider.getLuceneHnswBinaryQuantizedVectorsFormatInstance(16, 100)); } @Test