diff --git a/README.md b/README.md index 7e177c45..78c8a683 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,100 @@ 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 standard cuvs-lucene jar: + +```sh +mvn clean package -DskipTests +``` + +Then start PyLucene with the base `cuvs-java` jar, the standard `cuvs-lucene` +jar, and PyLucene's own Lucene classpath: + +```python +import os +from pathlib import Path + +import lucene + +cuvs_java_jar = Path(os.environ["CUVS_LUCENE_CUVS_JAVA_JAR"]) +cuvs_lucene_jar = next( + 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( + [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 + +codec = Codec.forName("Lucene101AcceleratedHNSWCodec") +``` + +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 +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: + +```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=/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 +``` + +To run an expanded GPU end-to-end pytest suite through CPU HNSW, +CAGRA-to-HNSW, and CAGRA search paths: + +```sh +./test_pylucene.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, 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: + +```sh +./test_pylucene.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 new file mode 100755 index 00000000..79222be2 --- /dev/null +++ b/ci/test_pylucene_smoke.sh @@ -0,0 +1,9 @@ +#!/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)" +exec "${REPO_ROOT}/test_pylucene.sh" "$@" diff --git a/examples/Python/pylucene_smoke.py b/examples/Python/pylucene_smoke.py new file mode 100644 index 00000000..48fc36ec --- /dev/null +++ b/examples/Python/pylucene_smoke.py @@ -0,0 +1,966 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import tempfile +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path + +import lucene + + +REPO_ROOT = Path(__file__).resolve().parents[2] +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" + + +@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(): + configured = 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( + 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 cuvs-lucene 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 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:") + 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 + / "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, + ) + + +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 writer_telemetry(codec): + vector_format = codec.knnVectorsFormat() + 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("=") + 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): + 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, telemetry): + observed = telemetry.get("writerPath") + 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, telemetry): + if case.expected_hnsw_layers: + observed_layers = telemetry.get("hnswLayers") + 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 = telemetry.get("cagraGraphBuildAlgo") + 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 = telemetry.get("cagraGraphDegree") + 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 = telemetry.get("cagraIntermediateGraphDegree") + 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.document import Document, Field, KnnFloatVectorField, StringField + from org.apache.lucene.index import ( + DirectoryReader, + IndexWriter, + IndexWriterConfig, + Term, + VectorSimilarityFunction, + ) + from org.apache.lucene.search import IndexSearcher + from org.apache.lucene.store import FSDirectory + + available_codecs = codec_class.availableCodecs() + if not available_codecs.contains(case.codec_name): + raise AssertionError( + f"{case.codec_name} was not advertised by Lucene SPI. " + f"Available codecs: {available_codecs}" + ) + 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)) + 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: + 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(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() + + assert_expected_index_files(index_path, case.expected_suffixes) + + reader = DirectoryReader.open(directory) + try: + 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()}" + ) + 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]) + telemetry = writer_telemetry(codec) + writer_path = observed_writer_path(case, telemetry) + assert_hnsw_telemetry(case, telemetry) + finally: + reader.close() + directory.close() + + algorithm_label = build_search_label(case, writer_path) + print( + "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__": + raise SystemExit(main()) diff --git a/examples/Python/test_pylucene_smoke.py b/examples/Python/test_pylucene_smoke.py new file mode 100644 index 00000000..fc001daa --- /dev/null +++ b/examples/Python/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/pom.xml b/pom.xml index 14390391..fc9add8b 100644 --- a/pom.xml +++ b/pom.xml @@ -166,9 +166,9 @@ maven-assembly-plugin 3.6.0 - - jar-with-dependencies - + + src/main/assembly/jar-with-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/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/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index ccc61eae..014d956b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -77,6 +77,11 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException return new CuVS2510GPUVectorsWriter(state, gpuSearchParams, flatWriter); } + @Override + public String toString() { + return getName() + "(" + WriterTelemetry.forCagra() + ")"; + } + /** * Returns a KnnVectorsReader instance to read the vectors from the index. */ 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 b4c5a33d..7ab4a154 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,19 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { */ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + 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..290a0586 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; @@ -79,9 +80,13 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); } else { + boolean forcedCpuFallback = isCpuHnswFallbackForced(); 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, @@ -96,6 +101,11 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException } } + @Override + public String toString() { + return getName() + "(" + WriterTelemetry.forHnsw(acceleratedHNSWParams) + ")"; + } + /** * Returns a KnnVectorsReader to read the vectors from the index. */ diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index f2c1aa37..5e149b49 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")); + super(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..a59cecb1 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; @@ -27,24 +26,61 @@ 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 static final LuceneProvider LUCENE_99_PROVIDER = + getLuceneProvider(LuceneProvider.LUCENE_99_FORMAT_VERSION); + private static volatile FlatVectorsFormat cachedFlatVectorsFormat; private final AcceleratedHNSWParams acceleratedHNSWParams; + private volatile KnnVectorsFormat cachedFallbackFormat; - static { + private static LuceneProvider getLuceneProvider(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 getOrCreateFlatVectorsFormat() { + FlatVectorsFormat format = cachedFlatVectorsFormat; + if (format == null) { + synchronized (LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class) { + format = cachedFlatVectorsFormat; + if (format == null) { + try { + 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); + } + } + } + } + return format; + } + + private KnnVectorsFormat getOrCreateFallbackFormat() throws Exception { + KnnVectorsFormat format = cachedFallbackFormat; + if (format == null) { + synchronized (this) { + format = cachedFallbackFormat; + if (format == null) { + format = + getLuceneProvider(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION) + .getLuceneHnswBinaryQuantizedVectorsFormatInstance( + acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); + cachedFallbackFormat = format; + } + } + } + return format; + } + /** * Initializes {@link LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat} with default values. * @@ -70,7 +106,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + var flatWriter = getOrCreateFlatVectorsFormat().fieldsWriter(state); if (isSupported()) { log.log( Level.FINE, @@ -84,12 +120,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException Level.WARNING, "GPU based indexing not supported, falling back to using the" + " Lucene102HnswBinaryQuantizedVectorsFormat"); - KnnVectorsFormat fallbackFormat = - LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance( - acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); - 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); } } } @@ -100,10 +133,10 @@ 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 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/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..8dfaf52d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -8,10 +8,14 @@ import java.lang.invoke.VarHandle; import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; +import java.util.ArrayList; +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; +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; @@ -31,6 +35,11 @@ public class LuceneProvider { static final Logger log = Logger.getLogger(LuceneProvider.class.getName()); + 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."; @@ -79,7 +88,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,14 +101,29 @@ 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; } private LuceneProvider(String version) throws ClassNotFoundException { + if (LUCENE_102_BINARY_FORMAT_VERSION.equals(version)) { + binaryQuantizedVectorsFormat = + loadClass( + setVersion(luceneBinaryQuantizedVectorsFormat, version), + setVersion(luceneBinaryQuantizedVectorsFormatFallback, version)); + hnswBinaryQuantizedVectorsFormat = + loadClass( + setVersion(luceneHnswBinaryQuantizedVectorsFormat, version), + setVersion(luceneHnswBinaryQuantizedVectorsFormatFallback, version)); + return; + } + flatVectorsFormat = loadClass( setVersion(luceneFlatVectorsFormat, version), @@ -125,18 +149,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) { @@ -147,14 +159,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; } } } @@ -173,6 +191,26 @@ public static Codec getCodec(String version) return (Codec) codecClassConstructor.newInstance(); } + public static Codec getDefaultDelegateCodec() { + 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); + } + } + throw new IllegalStateException( + "Unable to load a supported Lucene delegate codec. Tried " + + SUPPORTED_DELEGATE_CODEC_VERSIONS + + ". Failures: " + + failures); + } + public FlatVectorsFormat getLuceneFlatVectorsFormatInstance(FlatVectorsScorer scorer) throws Exception { try { @@ -245,7 +283,7 @@ public List getSimilarityFunctions() } } - public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws Exception { + public FlatVectorsFormat getLuceneBinaryQuantizedVectorsFormatInstance() throws Exception { try { Constructor luceneBinaryQuantizedVectorsFormatConstructor = binaryQuantizedVectorsFormat.getConstructor(); @@ -258,17 +296,17 @@ public FlatVectorsFormat getluceneBinaryQuantizedVectorsFormatInstance() throws } } - public FlatVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( + public KnnVectorsFormat getLuceneHnswBinaryQuantizedVectorsFormatInstance( int maxConn, int beamWidth) throws Exception { try { Constructor luceneHnswBinaryQuantizedVectorsFormatConstructor = - hnswBinaryQuantizedVectorsFormat.getConstructor(Integer.TYPE, Integer.TYPE); - return (FlatVectorsFormat) + 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/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/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/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/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 diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e660..4b152ab0 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -6,8 +6,11 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +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; @@ -33,9 +36,50 @@ 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); } + + @Test + public void testProviderCachesSupportedVersion() throws Exception { + 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 lucene102BinaryFormatProvider = + LuceneProvider.getInstance(LuceneProvider.LUCENE_102_BINARY_FORMAT_VERSION); + assertNotNull(lucene102BinaryFormatProvider.getLuceneBinaryQuantizedVectorsFormatInstance()); + assertNotNull( + lucene102BinaryFormatProvider.getLuceneHnswBinaryQuantizedVectorsFormatInstance(16, 100)); + } + + @Test + public void testDefaultDelegateCodec() { + 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" + }; + for (String codecName : codecNames) { + assertTrue(Codec.availableCodecs().contains(codecName)); + assertEquals(codecName, Codec.forName(codecName).getName()); + } + } } 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..73e21be0 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestWriterTelemetry.java @@ -0,0 +1,58 @@ +/* + * 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()); + assertEquals( + "CuVS2510GPUVectorsFormat(writerPath=gpu-cagra)", + new CuVS2510GPUVectorsFormat().toString()); + 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)); + 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 { + if (previousValue == null) { + System.clearProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY); + } else { + System.setProperty(FORCE_CPU_HNSW_FALLBACK_PROPERTY, previousValue); + } + } + } +} diff --git a/test_pylucene.sh b/test_pylucene.sh new file mode 100755 index 00000000..92880689 --- /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/Python/test_pylucene_smoke.py