diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3e35c05767..878b1fbfd9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -97,7 +97,7 @@ repos: name: spotless-fmt entry: ci/checks/run_spotless.sh pass_filenames: false - files: ^java/(cuvs-java|cuvs-lucene)/([^/]+/)?src/.*\.java$ + files: ^(java/(cuvs-java|cuvs-lucene)/([^/]+/)?src/.*|python/cuvs_bench/(cuvs_bench/backends/_java|tests/java)/.*)\.java$ exclude: .*/panama/.* language: script verbose: true diff --git a/ci/checks/run_spotless.sh b/ci/checks/run_spotless.sh index 08bf023d20..a9e01f2350 100755 --- a/ci/checks/run_spotless.sh +++ b/ci/checks/run_spotless.sh @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# pre-commit hook wrapper that runs 'spotless:apply' to format the Java sources of every Maven -# project under java/. +# pre-commit hook wrapper that runs 'spotless:apply' for the Java sources owned by the Maven +# projects under java/ and the Java adapters owned by cuvs-bench. # # Most cuvs contributors do not work on the Java client and do not have Maven installed. For them, # running 'pre-commit run --all-files' matches every Java source file in the repo regardless of @@ -15,11 +15,15 @@ set -euo pipefail # Keep these in sync with the spotless-fmt hook's 'files'/'exclude' entries in # .pre-commit-config.yaml. -JAVA_SRC_PATTERN='^java/(cuvs-java|cuvs-lucene)/([^/]+/)?src/.*\.java$' +JAVA_SRC_PATTERN='^(java/(cuvs-java|cuvs-lucene)/([^/]+/)?src/.*|python/cuvs_bench/(cuvs_bench/backends/_java|tests/java)/.*)\.java$' JAVA_SRC_EXCLUDE='.*/panama/.*' java_sources_modified() { - git status --porcelain --untracked-files=all -- java/cuvs-java java/cuvs-lucene | + git status --porcelain --untracked-files=all -- \ + java/cuvs-java \ + java/cuvs-lucene \ + python/cuvs_bench/cuvs_bench/backends/_java \ + python/cuvs_bench/tests/java | cut -c4- | grep -Ev "${JAVA_SRC_EXCLUDE}" | grep -Eq "${JAVA_SRC_PATTERN}" @@ -43,6 +47,7 @@ POMS=( java/cuvs-lucene/pom.xml java/cuvs-lucene/bench/pom.xml java/cuvs-lucene/examples/pom.xml + python/cuvs_bench/java-format-pom.xml ) for pom in "${POMS[@]}"; do diff --git a/conda/recipes/cuvs-bench/recipe.yaml b/conda/recipes/cuvs-bench/recipe.yaml index 655164f962..b8c037f2d3 100644 --- a/conda/recipes/cuvs-bench/recipe.yaml +++ b/conda/recipes/cuvs-bench/recipe.yaml @@ -63,7 +63,12 @@ requirements: tests: - package_contents: files: + - site-packages/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java + - site-packages/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml + - site-packages/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml - site-packages/cuvs_bench/config/datasets/datasets.yaml + - script: + - python -c "import importlib.metadata as metadata; paths = {path.as_posix() for path in (metadata.files('cuvs-bench') or ())}; required = {'cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java', 'cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml', 'cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml'}; missing = required - paths; forbidden = sorted(path for path in paths if '/tests/java/' in '/' + path or 'PyLuceneTestSupport' in path or path.endswith('.class')); assert not missing, f'missing cuvs-bench package resources: {sorted(missing)}'; assert not forbidden, f'forbidden cuvs-bench package payload: {forbidden}'" - python: imports: - cuvs_bench diff --git a/dependencies.yaml b/dependencies.yaml index 86ccd33d06..3f33feb81d 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -640,6 +640,7 @@ dependencies: packages: - click - cuvs==26.10.*,>=0.0.0a0 + - h5py>=3.8.0 - pandas - pyyaml - requests diff --git a/fern/docs.yml b/fern/docs.yml index 2aafaea69e..f372b30b7a 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -221,6 +221,8 @@ navigation: path: "./pages/cuvs_bench/synthesize_dataset.md" - page: "Backends" path: "./pages/cuvs_bench/pluggable_backend.md" + - page: "PyLucene Backend" + path: "./pages/cuvs_bench/pylucene_backend.md" - page: "cuVS Bench Parameter Tuning Guide" hidden: true path: "./pages/cuvs_bench/param_tuning.md" diff --git a/fern/pages/cuvs_bench/index.md b/fern/pages/cuvs_bench/index.md index ba026f47d7..9a91c7acc4 100644 --- a/fern/pages/cuvs_bench/index.md +++ b/fern/pages/cuvs_bench/index.md @@ -11,3 +11,4 @@ Use these guides to install cuVS Bench, run benchmark workloads, prepare dataset - [Usage](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage): configure algorithms, run benchmarks, and inspect results. - [Datasets](/user-guide/benchmarking-guide/cu-vs-bench-tool/datasets): prepare benchmark datasets, ground truth, and dataset descriptors. - [Backends](/user-guide/benchmarking-guide/cu-vs-bench-tool/backends): understand and extend cuVS Bench backend integrations. +- [PyLucene Backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/pylucene-backend): prepare the PyLucene runtime and run cuVS-Lucene HNSW and CAGRA benchmarks. diff --git a/fern/pages/cuvs_bench/install.md b/fern/pages/cuvs_bench/install.md index 9c525b4973..4f0b610d81 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -49,6 +49,12 @@ Exact tags are listed on Docker Hub: **Note:** GPU containers use the CUDA toolkit inside the container. The host only needs a compatible driver, so CUDA 12 containers can run on systems with CUDA 13.x-capable drivers. GPU access also requires the NVIDIA Docker runtime from the [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker). +## PyLucene backend prerequisites + +The optional `pylucene` backend requires runtime components that the standard cuVS Bench packages do not install automatically: JDK 22, a source-built PyLucene wrapper for Lucene 10.2.0, and version-matched cuVS Java, cuVS-Lucene, and native artifacts. + +See [PyLucene Backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/pylucene-backend) for the pinned dependency build, validation procedure, runtime configuration, and benchmark workflows. + ## Build from Source Build cuVS Bench from source when you need local benchmark executables that match a development checkout, include custom algorithm targets, or use dependencies that are not available in the pre-built packages. diff --git a/fern/pages/cuvs_bench/param_tuning.md b/fern/pages/cuvs_bench/param_tuning.md index a25b66378e..6911ff0192 100644 --- a/fern/pages/cuvs_bench/param_tuning.md +++ b/fern/pages/cuvs_bench/param_tuning.md @@ -210,6 +210,8 @@ Use FAISS IVF-PQ index on CPU ## HNSW +The PyLucene accelerated HNSW path has backend-specific parameter mappings and runtime constraints. See [PyLucene Backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/pylucene-backend#configure-hnsw-benchmarks) for its `m`, `ef_construction`, `direct_single_segment`, and `num_candidates` settings. + ### cuvs_hnsw NVIDIA cuVS HNSW constructs a CAGRA graph on the GPU and converts it to an HNSW index for CPU-based search. The build process uses in-memory CAGRA if the estimated memory requirements fit within available host and GPU memory. If not, it falls back to ACE (Augmented Core Extraction), which enables construction for larger datasets by partitioning them. Explicitly setting ACE parameters also triggers the ACE build path. The ACE-specific parameters listed below are relevant only when ACE is being used. diff --git a/fern/pages/cuvs_bench/pluggable_backend.md b/fern/pages/cuvs_bench/pluggable_backend.md index ce5eb2ec51..bcdd3880f9 100644 --- a/fern/pages/cuvs_bench/pluggable_backend.md +++ b/fern/pages/cuvs_bench/pluggable_backend.md @@ -194,7 +194,7 @@ class ElasticsearchBackend(BenchmarkBackend): dry_run=False, ): n_queries = dataset.n_queries - return SearchResult( + return [SearchResult( neighbors=np.zeros((n_queries, k), dtype=np.int64), distances=np.zeros((n_queries, k), dtype=np.float32), search_time_ms=0.0, @@ -203,7 +203,7 @@ class ElasticsearchBackend(BenchmarkBackend): algorithm=self.algo, search_params=indexes[0].search_params if indexes else [], success=True, - ) + )] ``` ```python @@ -219,9 +219,13 @@ get_registry().register("elasticsearch", ElasticsearchBackend) | Component | Description | | --- | --- | | `ConfigLoader` | Abstract class whose `load(**kwargs)` method returns `(DatasetConfig, List[BenchmarkConfig])`. Register with `register_config_loader(backend_type, loader_class)`. | -| `BenchmarkBackend` | Abstract class whose `build(...)` method returns `BuildResult` and whose `search(...)` method returns `SearchResult`. Register with `BackendRegistry.register(name, backend_class)`. | +| `BenchmarkBackend` | Abstract class whose `build(...)` method returns `BuildResult` and whose `search(...)` method returns `List[SearchResult]`. Register with `BackendRegistry.register(name, backend_class)`. | | `BackendRegistry` | Singleton registry returned by `get_registry()`. It maps backend type names to backend classes. | +## PyLucene Backend + +The built-in `pylucene` loader expands algorithm YAML groups into one Lucene index per selected codec. The backend initializes PyLucene's process-global JVM and resolves the production `Lucene101AcceleratedHNSWCodec` and `CuVS2510GPUSearchCodec` through Lucene's service-provider interface. See [PyLucene Backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/pylucene-backend) for codec behavior, dependency setup, configuration, benchmark parameters, and runtime limits. + ## C++ Backend The built-in `CppGoogleBenchmarkBackend` uses `backend_type="cpp_gbench"`. Its config loader reads YAML under `config/datasets` and `config/algos`, expands parameter combinations, and validates constraints. Its backend runs the C++ benchmark executables and merges their results. diff --git a/fern/pages/cuvs_bench/pylucene_backend.md b/fern/pages/cuvs_bench/pylucene_backend.md new file mode 100644 index 0000000000..5e106015a0 --- /dev/null +++ b/fern/pages/cuvs_bench/pylucene_backend.md @@ -0,0 +1,221 @@ +--- +slug: user-guide/benchmarking-guide/cu-vs-bench-tool/pylucene-backend +--- + +# PyLucene Backend + +The optional `pylucene` backend runs cuVS Bench through PyLucene's embedded JVM and cuVS-Lucene. It does not use the `*_ANN_BENCH` executables or `--executable-dir`. Use this guide to prepare its external dependencies, configure the runtime, and run representative HNSW and CAGRA benchmarks. + +## Supported benchmark paths + +| Algorithm | Codec | Index construction | Search | +| --- | --- | --- | --- | +| `pylucene_cuvs_hnsw` | `Lucene101AcceleratedHNSWCodec` | cuVS-accelerated HNSW construction when GPU support is available, with an intentional Lucene CPU-writer fallback | Lucene CPU HNSW | +| `pylucene_cuvs_cagra` | `CuVS2510GPUSearchCodec` | GPU CAGRA | GPU CAGRA | + +The backend accepts only FLOAT32 datasets with Euclidean distance, at most 4096 dimensions, and at least two indexed vectors. CAGRA requires a supported NVIDIA GPU for construction and search. HNSW can fall back to Lucene's CPU writer, with a warning, when cuVS GPU support is unavailable. Select the CAGRA path with `--algorithms pylucene_cuvs_cagra`. + +## Prerequisites + +cuVS Bench does not install the PyLucene runtime automatically. The backend requires: + +- JDK 22, including `javac`; +- a custom PyLucene wrapper generated against Lucene 10.2.0; +- [Maven 3.9.6 or newer](https://maven.apache.org/download.cgi) to build cuVS-Lucene; and +- the base `cuvs-java` JAR, the standard `cuvs-lucene` JAR, and matching native cuVS libraries for GPU execution. + +Apache does not publish PyLucene 10.2.0, and the official PyLucene 10.0.0 distribution is incompatible with the Lucene 10.2 APIs used here. The [PyLucene source-build instructions](https://lucene.apache.org/pylucene/install.html) describe the general build mechanics, but neither Apache nor cuVS currently provides a ready-to-use 10.2.0 wrapper. + +### Build matched cuVS dependencies + +cuVS-Lucene lives in this repository under `java/cuvs-lucene`. [NVIDIA/cuvs#2475](https://github.com/NVIDIA/cuvs/pull/2475) supplies the Lucene 10.2 compatibility and codec changes required by this backend, including HNSW heuristic delegation for the `m` and `ef_construction` benchmark parameters. This PR owns the Bench backend and its PyLucene end-to-end tests. Until #2475 is merged, use the pinned revision below rather than a moving pull-request head. + +Build the dependencies in a separate checkout so this does not change your cuVS Bench working tree. The pinned monorepo revision keeps cuVS Java, cuVS-Lucene, and the native libraries aligned. + +```bash +git clone https://github.com/NVIDIA/cuvs.git cuvs-pylucene-deps +cd cuvs-pylucene-deps +git fetch origin pull/2475/head +git switch --detach d6fcab0946837d7d3997cec4ed18189d3faa12e6 +./build.sh libcuvs java lucene +cd .. +``` + +If matching native cuVS libraries are already built and installed, `./build.sh java lucene` is sufficient. The Java build installs the base and native-classifier JARs into the local Maven repository; see the [cuVS Java build guide](https://github.com/NVIDIA/cuvs/blob/main/java/README.md). + +The conventional JAR paths are: + +```text +~/.m2/repository/com/nvidia/cuvs/cuvs-java/26.10.0/cuvs-java-26.10.0.jar +/java/cuvs-lucene/target/cuvs-lucene-26.10.0.jar +``` + +Use the base `cuvs-java` JAR, not a native-classifier JAR. Use the standard cuVS-Lucene JAR, not its `-jar-with-dependencies`, sources, or Javadoc variants. Native-library paths must resolve `libcuvs.so`, `libcuvs_c.so`, their dependencies, and the CUDA runtime libraries from the matching cuVS build. + +Use a clean environment without another cuVS native installation on its library path; otherwise, the JVM can load the other `libcuvs_c.so` first and reject the Java/native version mismatch. + +### Validate the dependency build + +The backend checks that `lucene.VERSION` is exactly `10.2.0` before starting the process-wide JVM. It also compiles its configured-codec adapter against the selected JAR, which fails early when the HNSW heuristic API is missing. + +Validate the pinned cuVS artifacts with the opt-in Bench PyLucene suite. Pytest owns the scenarios, assertions, parameterization, and reporting under `cuvs_bench/tests/pylucene`. Test-only Java adapters live under `python/cuvs_bench/tests/java`; the shared session fixture compiles all of them with `javac` into one temporary classes directory before the process-wide JVM starts. They are excluded from the cuVS Bench wheel and the production cuVS-Lucene JAR. + +```bash +python -m pip install pytest + +CUVS_DEPS_ROOT="$(cd cuvs-pylucene-deps && pwd)" +CUVS_NATIVE_BUILD="$CUVS_DEPS_ROOT/cpp/build" +export JAVA_LIBRARY_PATH="$CUVS_NATIVE_BUILD:$CUVS_NATIVE_BUILD/c:/usr/local/cuda/lib64" +export LD_LIBRARY_PATH="$JAVA_LIBRARY_PATH${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" +export CUVS_LUCENE_CUVS_JAVA_JAR="$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/26.10.0/cuvs-java-26.10.0.jar" +export CUVS_LUCENE_JAR="$CUVS_DEPS_ROOT/java/cuvs-lucene/target/cuvs-lucene-26.10.0.jar" +export CUVS_BENCH_PYLUCENE_INTEGRATION=1 + +cd /python/cuvs_bench + +# Non-extended live suite, including one case for each execution path. +python -m pytest -q -s cuvs_bench/tests/pylucene \ + -m "pylucene and not pylucene_extended" + +# Full GPU E2E suite, including the extended execution-path matrix. +python -m pytest -q -s cuvs_bench/tests/pylucene -m pylucene +``` + +Run the codec-path matrix directly with: + +```bash +python -m pytest -q -s \ + cuvs_bench/tests/pylucene/test_pylucene_execution_paths.py +``` + +The matrix distinguishes CPU HNSW, GPU/CAGRA-built HNSW, and GPU CAGRA search. Extended cases cover segment and force-merge layouts, one- and three-layer HNSW, CAGRA `searchWidth` values 16 and 32, deletion, filtering, deterministic brute-force recall, duplicate-hit checks, and rank-one self matches. `searchWidth` is the cuVS CAGRA query setting exercised through a test-only codec bridge; it is not the backend's HNSW `num_candidates` setting or cuVS Bench `--batch-size`. + +### Install cuVS Bench + +Install the cuVS Bench development checkout that contains this backend, separate from `cuvs-pylucene-deps`, into the activated PyLucene environment: + +```bash +cd +python -m pip install -e ./python/cuvs_bench +``` + +## Configure the backend + +Create `pylucene-backend.yaml` with absolute paths to the base `cuvs-java` JAR, the standard `cuvs-lucene` JAR, and a path-separated list of directories containing the matching cuVS and CUDA native libraries: + +```yaml +backend: pylucene +cuvs_java_jar: /home/user/.m2/repository/com/nvidia/cuvs/cuvs-java/26.10.0/cuvs-java-26.10.0.jar +cuvs_lucene_jar: /work/cuvs-pylucene-deps/java/cuvs-lucene/target/cuvs-lucene-26.10.0.jar +java_library_path: /work/cuvs-pylucene-deps/cpp/build:/work/cuvs-pylucene-deps/cpp/build/c:/usr/local/cuda/lib64 +``` + +Configuration sources are: + +| Backend configuration key | Environment variable | +| --- | --- | +| `cuvs_java_jar` | `CUVS_LUCENE_CUVS_JAVA_JAR` | +| `cuvs_lucene_jar` | `CUVS_LUCENE_JAR` | +| `java_library_path` | `JAVA_LIBRARY_PATH` or `LD_LIBRARY_PATH` | +| `jvm_args` | No environment alias; YAML list of additional JVM arguments. | + +PyLucene's JVM is process-global and can be initialized only once. Set the JAR paths, native-library locations, and `jvm_args` before the first PyLucene benchmark, and start a new Python process to change any of them. + +## Run a functional smoke test + +Use one explicit dataset root for both preparation and execution: + +```bash +export DATASET_ROOT=/absolute/path/to/cuvs-bench-data + +python -m cuvs_bench.get_dataset \ + --dataset test-data \ + --dataset-path "$DATASET_ROOT" \ + --test-data-n-train 512 \ + --test-data-n-test 4 \ + --test-data-k 5 + +python -m cuvs_bench.run \ + --backend-config pylucene-backend.yaml \ + --dataset test-data \ + --dataset-path "$DATASET_ROOT" \ + --algorithms pylucene_cuvs_hnsw \ + --groups test \ + --batch-size 2 -k 5 \ + -m latency \ + --build --search --force +``` + +This small synthetic run verifies the workflow; do not use it as performance or quality evidence. For a representative recall run, prepare `deep-image-96-angular` with `--normalize` into the same `DATASET_ROOT`, then run the backend against `deep-image-96-inner` with the same `--dataset-path`. + +## Configure HNSW benchmarks + +The HNSW algorithm maps benchmark parameters to cuVS-Lucene as follows: + +| Parameter | Scope | Default | Accepted values | Effect | +| --- | --- | --- | --- | --- | +| `codec` | Build | `Lucene101AcceleratedHNSWCodec` | That codec name | Selects the cuVS-accelerated HNSW codec. | +| `m` | Build | `32` | Integer from 1 through 512 | Sets `AcceleratedHNSWParams.maxConn`. | +| `ef_construction` | Build | `32` | Integer from 1 through 512 | Sets `AcceleratedHNSWParams.beamWidth`. | +| `direct_single_segment` | Build | `false` | Boolean | Requests one direct segment, as described below. | +| `num_candidates` | Search | `top_k` | Integer greater than or equal to `top_k` | Sets the candidate count passed to `KnnFloatVectorQuery`, capped at the index size; the backend returns `top_k` neighbors. | + +`m` and `ef_construction` are HNSW-equivalent inputs to cuVS-Lucene's `SAME_GRAPH_FOOTPRINT` heuristic. cuVS-Lucene derives the underlying CAGRA build parameters from them. `num_candidates` is Lucene's candidate budget, not a direct cuVS `ef_search` setting. + +Automatic tune mode samples `m` and `ef_construction` from 1 through 512 and `num_candidates` from `top_k` through 500. Explicit YAML sweeps may use larger candidate counts. + +Elasticsearch shard, replica, and field-name settings do not apply to this local Lucene index. Its HNSW index type maps to `codec`, and the backend validates Euclidean similarity from the dataset rather than accepting Elasticsearch's type, quantization, or similarity options. + +### Run a representative sweep + +The following configuration creates three indexes and searches each one with three candidate counts. Save it as `pylucene-deep1m.yaml`: + +```yaml +name: pylucene_cuvs_hnsw +groups: + deep1m: + build: + codec: ["Lucene101AcceleratedHNSWCodec"] + m: [16, 24, 32] + ef_construction: [32] + direct_single_segment: [true] + search: + num_candidates: [150, 200, 300] +``` + +Run it with `top_k=150`, substituting the dataset name and paths from its dataset configuration: + +```bash +python -m cuvs_bench.run \ + --backend-config pylucene-backend.yaml \ + --configuration pylucene-deep1m.yaml \ + --dataset deep1b-1M \ + --dataset-configuration /absolute/path/to/deep1b-1M.yaml \ + --dataset-path /absolute/path/to/dataset-root \ + --algorithms pylucene_cuvs_hnsw \ + --groups deep1m \ + --batch-size 100 -k 150 \ + -m latency \ + --build --search --force +``` + +## Index topology, provenance, and reuse + +When `direct_single_segment` is true, the backend disables ordinary RAM-triggered flushes and merging, buffers the requested vectors for one flush, and fails unless the committed index has exactly one segment. Lucene's per-indexing-thread hard RAM limit remains in force: 1945 MiB by default and less than 2048 MiB through the public API. If that limit forces an earlier flush, the build fails instead of merging the segments. + +`direct_single_segment` does not call Lucene `forceMerge`. The cuVS Bench `--force` option requests a rebuild; it does not control Lucene's segment merging. A larger JVM heap does not disable the per-thread hard limit. + +New HNSW and CAGRA builds atomically write commit-bound provenance manifests named `.cuvs-bench-pylucene-hnsw.json` and `.cuvs-bench-pylucene-cagra.json`, respectively. Reuse and search fail if the applicable manifest is missing, malformed, stale, or names different build parameters, writer policy, or compound-file policy. Indexes created outside this backend without the applicable manifest must be rebuilt with `--force`. + +For HNSW, the writer policy permits cuVS-Lucene's production CPU fallback. For CAGRA, the backend fails closed: it verifies the committed vector segments and checksums and rejects an index that is not CAGRA-only. + +For CAGRA, the backend disables compound files for both flushed and merged segments so the verifier can inspect the codec's `.vemc` and `.vcag` files directly. HNSW retains Lucene's default compound-file policy and its default index-writer scheduling unless `direct_single_segment` is enabled. + +## Runtime and measurement limits + +- The backend currently supports latency mode with one search thread. +- `--batch-size` groups queries for measurement; reported latency percentiles are milliseconds per batch. +- `num_candidates` must be greater than or equal to `top_k` for HNSW. +- Although cuVS-Lucene uses an effective `lucene_k` of `min(k, document_count)`, the backend conservatively requires `k <= 1024` for CAGRA to avoid paths that can use brute-force search above that limit. +- Throughput mode and multiple search threads are not implemented. +- PyLucene requires a wrapper generated against Lucene 10.2.0, and its process-wide JVM configuration cannot change after initialization. diff --git a/fern/pages/cuvs_bench/running.md b/fern/pages/cuvs_bench/running.md index 94606bb214..6e15762743 100644 --- a/fern/pages/cuvs_bench/running.md +++ b/fern/pages/cuvs_bench/running.md @@ -64,6 +64,7 @@ Create a custom YAML file with a `base` group to override the default benchmark | HNSWLIB | `hnswlib` | | DiskANN | `diskann_memory`, `diskann_ssd` | | NVIDIA cuVS | `cuvs_brute_force`, `cuvs_cagra`, `cuvs_ivf_flat`, `cuvs_ivf_pq`, `cuvs_cagra_hnswlib`, `cuvs_vamana` | +| PyLucene/cuVS | `pylucene_cuvs_hnsw`, `pylucene_cuvs_cagra` | ### Multi-GPU algorithms @@ -75,6 +76,10 @@ cuVS Bench includes single-node multi-GPU versions of IVF-Flat, IVF-PQ, and CAGR | IVF-PQ | `cuvs_mg_ivf_pq` | | CAGRA | `cuvs_mg_cagra` | +## Running the PyLucene backend + +PyLucene uses an embedded JVM and external cuVS-Lucene artifacts rather than the default `cpp_gbench` executables. See [PyLucene Backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/pylucene-backend) for its dependency setup, backend configuration, parameter semantics, and benchmark workflows. + ## Smaller-scale benchmarks (<1M to 10M vectors) Use `cuvs_bench.get_dataset` to prepare a built-in dataset. By default, datasets are stored under `RAPIDS_DATASET_ROOT_DIR` when that environment variable is set, or under a local `datasets` directory otherwise. @@ -192,7 +197,9 @@ Containers can also run in detached mode. ## Evaluating results -Build benchmarks report: +The tables below describe fields emitted by the default C++ Google Benchmark backend. Other backends report the fields that apply to their execution model, so not every field is present for every backend. + +C++ build benchmarks report: | Name | Description | | --- | --- | @@ -203,7 +210,7 @@ Build benchmarks report: | GPU | GPU time spent building. | | index_size | Number of vectors used to train the index. | -Search benchmarks report: +C++ search benchmarks report: | Name | Description | | --- | --- | diff --git a/python/cuvs_bench/cuvs_bench/_validation.py b/python/cuvs_bench/cuvs_bench/_validation.py new file mode 100644 index 0000000000..0e71cd25f3 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/_validation.py @@ -0,0 +1,61 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Validation helpers for names used in benchmark artifacts.""" + +from __future__ import annotations + +import os +import re +from typing import Any + + +_RESULT_COMPONENT = re.compile(r"[A-Za-z0-9_.-]+") + + +def validate_path_component(value: Any, description: str) -> str: + """Return a safe, non-empty filesystem path component.""" + if ( + not isinstance(value, str) + or not value + or value in {".", ".."} + or os.path.basename(value) != value + or any( + ord(character) < 32 or ord(character) == 127 for character in value + ) + ): + raise ValueError(f"Invalid {description}: {value!r}") + return value + + +def validate_result_component(value: Any, description: str) -> str: + """Return a path component safe for comma-delimited result identities.""" + component = validate_path_component(value, description) + if _RESULT_COMPONENT.fullmatch(component) is None: + raise ValueError(f"Invalid {description}: {value!r}") + return component + + +def format_artifact_identity( + algorithm: Any, + group: Any, + scope: Any = None, + *, + description: str = "benchmark", +) -> str: + """Format an injective identity from validated artifact components.""" + algorithm = validate_result_component( + algorithm, f"{description} algorithm name" + ) + group = validate_result_component(group, f"{description} group name") + if scope is not None: + scope = validate_result_component(scope, f"{description} result scope") + + identity = algorithm + if group != "base": + identity += f"[group={group}]" + if scope is not None: + identity += f"[scope={scope}]" + return identity diff --git a/python/cuvs_bench/cuvs_bench/backends/__init__.py b/python/cuvs_bench/cuvs_bench/backends/__init__.py index a09e6fd1fe..e2e7845953 100644 --- a/python/cuvs_bench/cuvs_bench/backends/__init__.py +++ b/python/cuvs_bench/cuvs_bench/backends/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -26,11 +26,13 @@ from .cpp_gbench import CppGoogleBenchmarkBackend from .opensearch import OpenSearchBackend +from .pylucene import PyLuceneBackend # Auto-register built-in backends _registry = get_registry() _registry.register("cpp_gbench", CppGoogleBenchmarkBackend) _registry.register("opensearch", OpenSearchBackend) +_registry.register("pylucene", PyLuceneBackend) __all__ = [ # Base classes and data structures @@ -46,4 +48,5 @@ # Built-in backends "CppGoogleBenchmarkBackend", "OpenSearchBackend", + "PyLuceneBackend", ] diff --git a/python/cuvs_bench/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java b/python/cuvs_bench/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java new file mode 100644 index 0000000000..1799bb1def --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.bench; + +import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType; +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; + +/** + * PyLucene-compatible adapter for configuring the accelerated HNSW codec. + * + *

PyLucene instantiates codecs through a no-argument constructor. This adapter reads the two + * Lucene-equivalent HNSW build parameters from system properties and delegates all codec behavior + * to the production cuVS-Lucene implementation. + */ +public final class PyLuceneConfiguredHnswCodec extends Lucene101AcceleratedHNSWCodec { + + public static final String M_PROPERTY = "com.nvidia.cuvs.bench.pylucene.hnsw.m"; + public static final String EF_CONSTRUCTION_PROPERTY = + "com.nvidia.cuvs.bench.pylucene.hnsw.efConstruction"; + + private final int m; + private final int efConstruction; + + /** Constructs the production codec with parameters supplied through system properties. */ + public PyLuceneConfiguredHnswCodec() throws Exception { + this(configuredValues()); + } + + private PyLuceneConfiguredHnswCodec(ConfiguredValues values) throws Exception { + super(values.parameters()); + this.m = values.m(); + this.efConstruction = values.efConstruction(); + } + + private static ConfiguredValues configuredValues() { + int m = requiredIntegerProperty(M_PROPERTY); + int efConstruction = requiredIntegerProperty(EF_CONSTRUCTION_PROPERTY); + AcceleratedHNSWParams parameters = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + .withHnswHeuristicType(HnswHeuristicType.SAME_GRAPH_FOOTPRINT) + .withMaxConn(m) + .withBeamWidth(efConstruction) + .build(); + return new ConfiguredValues(m, efConstruction, parameters); + } + + private static int requiredIntegerProperty(String name) { + String value = System.getProperty(name); + if (value == null) { + throw new IllegalStateException("Required system property is not set: " + name); + } + + try { + return Integer.parseInt(value); + } catch (NumberFormatException error) { + throw new IllegalArgumentException( + "System property " + name + " must be an integer, found: " + value, error); + } + } + + @Override + public String toString() { + return getClass().getSimpleName() + "(m=" + m + ", efConstruction=" + efConstruction + ")"; + } + + private record ConfiguredValues(int m, int efConstruction, AcceleratedHNSWParams parameters) {} +} diff --git a/python/cuvs_bench/cuvs_bench/backends/_java/__init__.py b/python/cuvs_bench/cuvs_bench/backends/_java/__init__.py new file mode 100644 index 0000000000..5d56d9d52c --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/_java/__init__.py @@ -0,0 +1,6 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Packaged Java sources used by the PyLucene backend.""" diff --git a/python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py b/python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py new file mode 100644 index 0000000000..d49218fd54 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py @@ -0,0 +1,154 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Compile the no-argument codec adapter required by stock PyLucene.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import tempfile +import threading +from pathlib import Path +from typing import Optional, Union + +CONFIGURED_HNSW_CODEC_CLASS = ( + "com.nvidia.cuvs.bench.PyLuceneConfiguredHnswCodec" +) +M_PROPERTY = "com.nvidia.cuvs.bench.pylucene.hnsw.m" +EF_CONSTRUCTION_PROPERTY = "com.nvidia.cuvs.bench.pylucene.hnsw.efConstruction" + +_SOURCE_FILE = ( + Path(__file__).with_name("_java") / "PyLuceneConfiguredHnswCodec.java" +) +_COMPILE_LOCK = threading.Lock() +_COMPILED_CLASSPATH: Optional[str] = None +_CLASSES_DIRECTORY: Optional[Path] = None +_TEMPORARY_DIRECTORY: Optional[tempfile.TemporaryDirectory] = None + + +def configured_codec_classes_path( + cuvs_java_jar: Path, + cuvs_lucene_jar: Path, + pylucene_classpath: str, +) -> Path: + """Compile the codec adapter once and return its stable classes path. + + This function must be called before initializing PyLucene's process-wide + JVM so that the returned directory can be included in the JVM classpath. + """ + + compile_classpath = _compile_classpath( + cuvs_java_jar, cuvs_lucene_jar, pylucene_classpath + ) + with _COMPILE_LOCK: + if _CLASSES_DIRECTORY is not None: + _require_same_classpath(compile_classpath) + return _CLASSES_DIRECTORY + return _compile(compile_classpath) + + +def _compile_classpath( + cuvs_java_jar: Union[str, os.PathLike[str]], + cuvs_lucene_jar: Union[str, os.PathLike[str]], + pylucene_classpath: str, +) -> str: + if not isinstance(pylucene_classpath, str) or not pylucene_classpath: + raise ValueError("pylucene_classpath must be a non-empty string") + + jar_paths = ( + Path(cuvs_java_jar).resolve(), + Path(cuvs_lucene_jar).resolve(), + ) + for jar_path in jar_paths: + if not jar_path.is_file(): + raise FileNotFoundError( + f"Cannot compile the PyLucene codec adapter; JAR not found: {jar_path}" + ) + return os.pathsep.join((*map(str, jar_paths), pylucene_classpath)) + + +def _require_same_classpath(compile_classpath: str) -> None: + if compile_classpath != _COMPILED_CLASSPATH: + raise RuntimeError( + "The PyLucene codec adapter was already compiled with a different " + "dependency classpath. PyLucene JVM dependencies are process-wide." + ) + + +def _compile(compile_classpath: str) -> Path: + global _CLASSES_DIRECTORY, _COMPILED_CLASSPATH, _TEMPORARY_DIRECTORY + + if not _SOURCE_FILE.is_file(): + raise RuntimeError( + f"The packaged PyLucene codec adapter source is missing: {_SOURCE_FILE}" + ) + + javac = _find_javac() + temporary_directory = tempfile.TemporaryDirectory( + prefix="cuvs-bench-pylucene-codec-" + ) + classes_directory = Path(temporary_directory.name) / "classes" + classes_directory.mkdir() + try: + completed = subprocess.run( + [ + javac, + "--release", + "22", + "-classpath", + compile_classpath, + "-d", + str(classes_directory), + str(_SOURCE_FILE), + ], + capture_output=True, + check=False, + text=True, + ) + except OSError as error: + temporary_directory.cleanup() + raise RuntimeError( + f"Could not run JDK 22 javac at {javac}: {error}" + ) from error + + if completed.returncode != 0: + temporary_directory.cleanup() + details = completed.stderr.strip() or completed.stdout.strip() + raise RuntimeError( + "Could not compile the PyLucene codec adapter with JDK 22 javac. " + "Verify that cuvs-java, PyLucene 10.2, and the thin cuvs-lucene " + "JAR are compatible. The cuvs-lucene JAR must include the " + "in-tree PyLucene 10.2 codec support and current HNSW heuristic " + "API. " + f"javac output:\n{details}" + ) + + _TEMPORARY_DIRECTORY = temporary_directory + _CLASSES_DIRECTORY = classes_directory + _COMPILED_CLASSPATH = compile_classpath + return classes_directory + + +def _find_javac() -> str: + candidates = [] + java_home = os.environ.get("JAVA_HOME") + if java_home: + candidates.append(Path(java_home) / "bin" / "javac") + candidates.append(Path(sys.prefix) / "lib" / "jvm" / "bin" / "javac") + + for candidate in candidates: + if candidate.is_file(): + return str(candidate) + + javac = shutil.which("javac") + if javac is not None: + return javac + raise RuntimeError( + "Configurable PyLucene HNSW builds require JDK 22 javac. Set " + "JAVA_HOME to a JDK 22 installation or put its javac on PATH." + ) diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index 15209decb6..1c9ebfeb55 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -112,6 +112,13 @@ def training_vectors(self) -> np.ndarray: ) return self._training_vectors + @property + def loaded_training_vectors(self) -> Optional[np.ndarray]: + """Training vectors already in memory, without loading ``base_file``.""" + if self._training_vectors.size == 0: + return None + return self._training_vectors + @training_vectors.setter def training_vectors(self, value: Optional[np.ndarray]) -> None: """Set training vectors directly.""" diff --git a/python/cuvs_bench/cuvs_bench/backends/pylucene.py b/python/cuvs_bench/cuvs_bench/backends/pylucene.py new file mode 100644 index 0000000000..e62a478c20 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -0,0 +1,3430 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyLucene backend for cuVS-accelerated Lucene codecs.""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import os +import shutil +import tempfile +import threading +import time +import zipfile +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import ( + Any, + Callable, + Dict, + List, + NamedTuple, + Optional, + Tuple, + Union, +) + +import numpy as np + +from .._bin_format import read_bin_header +from .._validation import format_artifact_identity, validate_path_component +from ..orchestrator.config_loaders import ( + BenchmarkConfig, + ConfigLoader, + DatasetConfig, + IndexConfig, +) +from ._utils import dtype_from_filename +from ._pylucene_java import ( + CONFIGURED_HNSW_CODEC_CLASS, + EF_CONSTRUCTION_PROPERTY, + M_PROPERTY, + configured_codec_classes_path, +) +from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult + +_ID_FIELD = "id" +_VECTOR_FIELD = "vector" +_MAX_DIMENSIONS = 4096 +_REQUIRED_PYLUCENE_VERSION = "10.2.0" + +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_SUPPORTED_CODECS = frozenset({_HNSW_CODEC, _CAGRA_CODEC}) +_HNSW_BUILD_KEYS = frozenset( + {"codec", "m", "ef_construction", "direct_single_segment"} +) +_SUPPORTED_BUILD_KEYS = _HNSW_BUILD_KEYS +_SUPPORTED_SEARCH_KEYS = frozenset({"num_candidates"}) +_DEFAULT_M = 32 +_DEFAULT_EF_CONSTRUCTION = 32 +_MIN_HNSW_BUILD_PARAMETER = 1 +_MAX_HNSW_BUILD_PARAMETER = 512 +_EXPECTED_WRITER_POLICY = { + _HNSW_CODEC: "gpu-with-cpu-fallback", + _CAGRA_CODEC: "gpu-cagra", +} +_COMPOUND_FILE_POLICY = { + _HNSW_CODEC: "lucene-default", + _CAGRA_CODEC: "disabled", +} +_CAGRA_META_EXTENSION = ".vemc" +_CAGRA_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta" +_CAGRA_INDEX_EXTENSION = ".vcag" +_CAGRA_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex" +_CAGRA_META_VERSION = 0 +_CAGRA_INDEX_VERSION = 0 +# Avoid caching checksums while coarse filesystem timestamps can still make a +# same-size rewrite look unchanged. +_CAGRA_CACHE_MIN_FILE_AGE_NS = 2_000_000_000 +_FLOAT32_ENCODING_ORDINAL = 1 +_EUCLIDEAN_SIMILARITY_ORDINAL = 0 + +_HNSW_PROVENANCE_FILE = ".cuvs-bench-pylucene-hnsw.json" +_CAGRA_PROVENANCE_FILE = ".cuvs-bench-pylucene-cagra.json" +_PROVENANCE_SCHEMA_VERSION = 4 +_PROVENANCE_KEYS = frozenset( + { + "schema_version", + "codec", + "build_parameters", + "writer_policy", + "compound_file_policy", + "vector_count", + "dimensions", + "segment_count", + "commit_fingerprints", + } +) +_SHA256_HEX_DIGITS = frozenset("0123456789abcdef") +_LUCENE_CORE_CLASS = "org/apache/lucene/index/IndexWriter.class" + +_JVM_INIT_LOCK = threading.Lock() +_CONFIGURED_CODEC_LOCK = threading.Lock() +_INITIALIZED_CLASSPATH: Optional[str] = None +_INITIALIZED_VMARGS: Optional[Tuple[str, ...]] = None + + +def _attempt_cleanup( + cleanup: Callable[[], None], + description: str, + primary_error: Optional[BaseException], +) -> Optional[BaseException]: + try: + cleanup() + except BaseException as cleanup_error: + if primary_error is None: + return cleanup_error + if isinstance(primary_error, Exception) and not isinstance( + cleanup_error, Exception + ): + cleanup_error.add_note( + f"Raised while attempting to {description}; prior failure: " + f"{type(primary_error).__name__}: {primary_error}" + ) + return cleanup_error + primary_error.add_note( + f"Failed to {description}: {type(cleanup_error).__name__}: {cleanup_error}" + ) + return primary_error + + +class _CleanupStack: + """Close registered resources without losing the operation's failure.""" + + def __init__(self) -> None: + self._cleanups: List[Tuple[str, Callable[[], None]]] = [] + + def __enter__(self) -> _CleanupStack: + return self + + def add(self, description: str, cleanup: Callable[[], None]) -> None: + self._cleanups.append((description, cleanup)) + + def __exit__( + self, + _error_type: Any, + primary_error: Optional[BaseException], + _traceback: Any, + ) -> bool: + final_error = primary_error + for description, cleanup in reversed(self._cleanups): + final_error = _attempt_cleanup(cleanup, description, final_error) + if final_error is not None and final_error is not primary_error: + raise final_error + return False + + +def _exception_summary(error: Exception) -> str: + details = [f"{type(error).__name__}: {error}"] + if error.__cause__ is not None: + cause = error.__cause__ + details.append(f"caused by {type(cause).__name__}: {cause}") + details.extend(getattr(error, "__notes__", ())) + return "; ".join(details) + + +def _restore_java_property(system: Any, name: str, previous: Any) -> None: + if previous is None: + system.clearProperty(name) + else: + system.setProperty(name, str(previous)) + + +def _configured_jar( + config: Dict[str, Any], config_key: str, environment_key: str +) -> Path: + value = config.get(config_key) or os.environ.get(environment_key) + if not value: + raise RuntimeError( + f"PyLucene backend requires '{config_key}' or " + f"the {environment_key} environment variable" + ) + + path = Path(os.fspath(value)).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"{config_key} does not exist: {path}") + return path + + +def _reject_bundled_lucene_classes(cuvs_lucene_jar: Path) -> None: + """Reject fat cuVS-Lucene jars before they can poison the process JVM.""" + if not zipfile.is_zipfile(cuvs_lucene_jar): + return + + with zipfile.ZipFile(cuvs_lucene_jar) as archive: + try: + archive.getinfo(_LUCENE_CORE_CLASS) + except KeyError: + return + + raise RuntimeError( + "cuvs_lucene_jar bundles Lucene classes and is incompatible with " + "PyLucene's process-wide JVM. Use the standard thin cuvs-lucene JAR, " + "not a '-jar-with-dependencies' artifact." + ) + + +def _load_pylucene() -> Any: + try: + return importlib.import_module("lucene") + except ImportError as exc: + raise ImportError( + "PyLucene's `lucene` module is required for the pylucene " + "cuvs-bench backend. PyLucene must be built and installed " + "separately; see the cuVS Bench installation documentation." + ) from exc + + +def _validate_pylucene_version(lucene: Any) -> None: + actual_version = str(getattr(lucene, "VERSION", "")) + if actual_version != _REQUIRED_PYLUCENE_VERSION: + raise RuntimeError( + "PyLucene must be generated against the same Lucene version as " + "cuVS-Lucene: expected " + f"{_REQUIRED_PYLUCENE_VERSION}, found {actual_version}. " + "Activate a matching PyLucene build before initializing the JVM." + ) + + +def _pylucene_classpath(config: Dict[str, Any], lucene: Any) -> str: + cuvs_java_jar = _configured_jar( + config, "cuvs_java_jar", "CUVS_LUCENE_CUVS_JAVA_JAR" + ) + cuvs_lucene_jar = _configured_jar( + config, "cuvs_lucene_jar", "CUVS_LUCENE_JAR" + ) + _reject_bundled_lucene_classes(cuvs_lucene_jar) + adapter_classes = configured_codec_classes_path( + cuvs_java_jar, + cuvs_lucene_jar, + str(lucene.CLASSPATH), + ) + return os.pathsep.join( + [ + str(adapter_classes), + str(cuvs_java_jar), + str(cuvs_lucene_jar), + str(lucene.CLASSPATH), + ] + ) + + +def _pylucene_vmargs(config: Dict[str, Any]) -> List[str]: + java_library_path = ( + config.get("java_library_path") + or 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}") + + extra_vmargs = config.get("jvm_args", []) + if isinstance(extra_vmargs, (str, bytes)) or not isinstance( + extra_vmargs, (list, tuple) + ): + raise TypeError("jvm_args must be a list or tuple of strings") + if not all(isinstance(arg, str) for arg in extra_vmargs): + raise TypeError("every jvm_args entry must be a string") + vmargs.extend(extra_vmargs) + return vmargs + + +def _attach_pylucene_jvm( + lucene: Any, classpath: str, vmargs: List[str] +) -> None: + with _JVM_INIT_LOCK: + vm_environment = lucene.getVMEnv() + if vm_environment is None: + vm_environment = _start_pylucene_jvm(lucene, classpath, vmargs) + else: + _validate_pylucene_jvm_configuration(classpath, vmargs) + + vm_environment.attachCurrentThread() + + +def _start_pylucene_jvm(lucene: Any, classpath: str, vmargs: List[str]) -> Any: + global _INITIALIZED_CLASSPATH, _INITIALIZED_VMARGS + vm_environment = lucene.initVM(classpath=classpath, vmargs=vmargs) + if vm_environment is None: + vm_environment = lucene.getVMEnv() + if vm_environment is None: + raise RuntimeError("PyLucene did not return a JVM environment") + + _INITIALIZED_CLASSPATH = classpath + _INITIALIZED_VMARGS = tuple(vmargs) + return vm_environment + + +def _validate_pylucene_jvm_configuration( + classpath: str, vmargs: List[str] +) -> None: + if _INITIALIZED_CLASSPATH is None or _INITIALIZED_VMARGS is None: + raise RuntimeError( + "PyLucene's process-wide JVM was initialized before the " + "pylucene backend, so the required cuVS classpath and JVM " + "arguments cannot be verified. Start a new Python process and " + "let cuVS Bench initialize PyLucene." + ) + if _INITIALIZED_CLASSPATH != classpath: + raise RuntimeError( + "PyLucene's process-wide JVM is already initialized with " + "different cuVS Java or cuVS-Lucene jars" + ) + if _INITIALIZED_VMARGS != tuple(vmargs): + raise RuntimeError( + "PyLucene's process-wide JVM is already initialized with " + "different JVM arguments or native-library paths" + ) + + +def _initialize_pylucene(config: Dict[str, Any]) -> Any: + """Initialize PyLucene once and attach the current Python thread.""" + lucene = _load_pylucene() + _validate_pylucene_version(lucene) + classpath = _pylucene_classpath(config, lucene) + vmargs = _pylucene_vmargs(config) + _attach_pylucene_jvm(lucene, classpath, vmargs) + return lucene + + +def _score_to_squared_euclidean(score: float) -> float: + """Convert Lucene's Euclidean score to squared Euclidean distance.""" + if score <= 0.0: + return float("inf") + return max(0.0, (1.0 / score) - 1.0) + + +def _validate_float32_matrix(vectors: np.ndarray, name: str) -> np.ndarray: + array = np.asarray(vectors) + if array.ndim != 2: + raise ValueError(f"{name} must be a two-dimensional array") + if array.shape[0] == 0 or array.shape[1] == 0: + raise ValueError(f"{name} must contain at least one vector") + if array.dtype != np.float32: + raise TypeError(f"{name} must use float32 values, got {array.dtype}") + if array.shape[1] > _MAX_DIMENSIONS: + raise ValueError( + f"{name} dimensions must not exceed {_MAX_DIMENSIONS}, got {array.shape[1]}" + ) + if not np.isfinite(array).all(): + raise ValueError(f"{name} must contain only finite values") + return np.ascontiguousarray(array) + + +def _validate_metric(dataset: Dataset) -> None: + if dataset.distance_metric.lower() not in {"euclidean", "l2"}: + raise ValueError( + "PyLucene cuVS codecs currently support only Euclidean/L2 " + f"datasets, got {dataset.distance_metric!r}" + ) + + +def _validate_codec(codec_name: Any) -> str: + if not isinstance(codec_name, str) or codec_name not in _SUPPORTED_CODECS: + available = ", ".join(sorted(_SUPPORTED_CODECS)) + raise ValueError( + f"Unsupported PyLucene codec {codec_name!r}. Supported codecs: {available}" + ) + return codec_name + + +def _configured_codec( + build_params: Dict[str, Any], backend_config: Dict[str, Any] +) -> str: + if "codec" in build_params: + codec_name = build_params["codec"] + else: + codec_name = backend_config.get("codec") + return _validate_codec(codec_name) + + +def _bounded_hnsw_build_parameter(value: Any, name: str) -> int: + if type(value) is not int or not ( + _MIN_HNSW_BUILD_PARAMETER <= value <= _MAX_HNSW_BUILD_PARAMETER + ): + raise ValueError( + f"PyLucene {name} must be an integer in " + f"[{_MIN_HNSW_BUILD_PARAMETER}, {_MAX_HNSW_BUILD_PARAMETER}], " + f"got {value!r}" + ) + return value + + +def _direct_single_segment(value: Any) -> bool: + if type(value) is not bool: + raise TypeError( + f"PyLucene direct_single_segment must be a boolean, got {value!r}" + ) + return value + + +def _normalize_build_params( + build_params: Any, backend_config: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + if not isinstance(build_params, dict): + raise TypeError("PyLucene build parameters must be a mapping") + unsupported = set(build_params) - _SUPPORTED_BUILD_KEYS + if unsupported: + names = ", ".join(sorted(str(name) for name in unsupported)) + raise ValueError(f"Unsupported PyLucene build parameter(s): {names}") + codec_name = _configured_codec(build_params, backend_config or {}) + hnsw_parameters = set(build_params) - {"codec"} + if codec_name != _HNSW_CODEC and hnsw_parameters: + names = ", ".join(sorted(hnsw_parameters)) + raise ValueError( + f"PyLucene build parameter(s) {names} apply only to {_HNSW_CODEC}" + ) + if codec_name == _CAGRA_CODEC: + return {"codec": codec_name} + + return { + "codec": codec_name, + "m": _bounded_hnsw_build_parameter( + build_params.get("m", _DEFAULT_M), "m" + ), + "ef_construction": _bounded_hnsw_build_parameter( + build_params.get("ef_construction", _DEFAULT_EF_CONSTRUCTION), + "ef_construction", + ), + "direct_single_segment": _direct_single_segment( + build_params.get("direct_single_segment", False) + ), + } + + +def _normalize_search_params( + search_params: Any, + *, + codec_name: str, + k: Optional[int] = None, +) -> List[Dict[str, Any]]: + if not isinstance(search_params, list) or not search_params: + raise ValueError( + "PyLucene search parameters must be a non-empty list of mappings" + ) + + normalized = [] + for parameters in search_params: + if not isinstance(parameters, dict): + raise TypeError( + "Every PyLucene search parameter must be a mapping" + ) + unsupported = set(parameters) - _SUPPORTED_SEARCH_KEYS + if unsupported: + names = ", ".join(sorted(str(name) for name in unsupported)) + raise ValueError( + f"Unsupported PyLucene search parameter(s): {names}" + ) + if codec_name != _HNSW_CODEC and parameters: + raise ValueError( + f"PyLucene num_candidates applies only to {_HNSW_CODEC}" + ) + if not parameters: + if codec_name == _HNSW_CODEC and k is not None: + normalized.append({"num_candidates": k}) + else: + normalized.append({}) + continue + + num_candidates = parameters["num_candidates"] + if type(num_candidates) is not int or num_candidates < 1: + raise ValueError( + "PyLucene num_candidates must be a positive integer, " + f"got {num_candidates!r}" + ) + if k is not None and num_candidates < k: + raise ValueError( + "PyLucene num_candidates must be greater than or equal to " + f"k ({k}), got {num_candidates}" + ) + normalized.append({"num_candidates": num_candidates}) + return normalized + + +def _safe_remove_index(index_path: Path, trusted_index_root: Path) -> None: + resolved = index_path.resolve() + forbidden = { + Path(resolved.anchor), + Path.cwd().resolve(), + Path.home().resolve(), + } + if resolved in forbidden: + raise ValueError(f"Refusing to remove unsafe index path: {resolved}") + if resolved.parent != trusted_index_root.resolve(): + raise ValueError( + f"Refusing to remove PyLucene index outside its configured root: {resolved}" + ) + if index_path.is_symlink() or not index_path.is_dir(): + raise ValueError( + f"PyLucene index path must be a directory: {index_path}" + ) + shutil.rmtree(index_path) + + +def _index_size(index_path: Path) -> int: + return sum( + path.stat().st_size for path in index_path.rglob("*") if path.is_file() + ) + + +@dataclass(frozen=True) +class _FileSignature: + resolved_path: str + device: int + inode: int + size: int + modified_at_ns: int + changed_at_ns: int + + +def _file_signature(path: Path) -> _FileSignature: + file_stat = path.stat() + return _FileSignature( + resolved_path=str(path.resolve()), + device=file_stat.st_dev, + inode=file_stat.st_ino, + size=file_stat.st_size, + modified_at_ns=file_stat.st_mtime_ns, + changed_at_ns=file_stat.st_ctime_ns, + ) + + +def _has_lucene_segments(index_path: Path) -> bool: + return any( + path.is_file() and path.name.startswith("segments_") + for path in index_path.iterdir() + ) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for block in iter(lambda: file.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: + commit_files = sorted( + path + for path in index_path.iterdir() + if path.name.startswith("segments_") + ) + if not commit_files: + raise RuntimeError( + "Lucene index has no segments_* commit file to fingerprint" + ) + + fingerprints = [] + for path in commit_files: + if path.is_symlink() or not path.is_file(): + raise RuntimeError( + f"Lucene commit fingerprint target must be a regular file: {path}" + ) + fingerprints.append({"name": path.name, "sha256": _sha256_file(path)}) + return fingerprints + + +@dataclass(frozen=True) +class _IndexProvenanceVerification: + codec: str + build_parameters: Dict[str, Any] + writer_policy: str + compound_file_policy: str + vector_count: int + dimensions: int + segment_count: int + commit_file_count: int + + def to_metadata(self) -> Dict[str, Any]: + return { + "status": f"{self.writer_policy}-provenance", + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": self.codec, + "build_parameters": self.build_parameters, + "writer_policy": self.writer_policy, + "compound_file_policy": self.compound_file_policy, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + "segment_count": self.segment_count, + "commit_file_count": self.commit_file_count, + } + + +@dataclass(frozen=True) +class _ProvenanceExpectation: + codec: str + build_parameters: Dict[str, Any] + manifest_name: str + label: str + vector_count: Optional[int] + dimensions: Optional[int] + + +class _IndexProvenanceError(RuntimeError): + """Raised when an index cannot be proven to be backend-built.""" + + +def _write_index_provenance( + index_path: Path, + codec: str, + build_parameters: Dict[str, Any], + vector_count: int, + dimensions: int, + segment_count: int, + manifest_name: str, +) -> None: + payload = { + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": codec, + "build_parameters": build_parameters, + "writer_policy": _EXPECTED_WRITER_POLICY[codec], + "compound_file_policy": _COMPOUND_FILE_POLICY[codec], + "vector_count": int(vector_count), + "dimensions": int(dimensions), + "segment_count": int(segment_count), + "commit_fingerprints": _commit_fingerprints(index_path), + } + manifest_path = index_path / manifest_name + with _CleanupStack() as cleanups: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=index_path, + prefix=f"{manifest_name}.", + suffix=".tmp", + delete=False, + ) as file: + temporary_path = Path(file.name) + cleanups.add( + "remove temporary provenance file", + lambda: temporary_path.unlink(missing_ok=True), + ) + json.dump( + payload, + file, + allow_nan=False, + indent=2, + sort_keys=True, + ) + file.write("\n") + file.flush() + os.fsync(file.fileno()) + temporary_path.chmod(0o644) + temporary_path.replace(manifest_path) + + +def _write_hnsw_provenance( + index_path: Path, + codec: str, + vector_count: int, + dimensions: int, + build_parameters: Optional[Dict[str, Any]] = None, + segment_count: int = 1, +) -> None: + effective_parameters = _normalize_build_params( + build_parameters or {"codec": codec} + ) + _write_index_provenance( + index_path, + codec, + effective_parameters, + vector_count, + dimensions, + segment_count, + _HNSW_PROVENANCE_FILE, + ) + + +def _write_cagra_provenance( + index_path: Path, + vector_count: int, + dimensions: int, + build_parameters: Optional[Dict[str, Any]] = None, + segment_count: int = 1, +) -> None: + effective_parameters = _normalize_build_params( + build_parameters or {"codec": _CAGRA_CODEC} + ) + _write_index_provenance( + index_path, + _CAGRA_CODEC, + effective_parameters, + vector_count, + dimensions, + segment_count, + _CAGRA_PROVENANCE_FILE, + ) + + +def _require_positive_int( + value: Any, provenance_label: str, field_name: str +) -> int: + if type(value) is not int or value < 1: + raise _IndexProvenanceError( + f"{provenance_label} field {field_name!r} must be a positive integer" + ) + return value + + +def _read_index_provenance( + manifest_path: Path, provenance_label: str +) -> Dict[str, Any]: + try: + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise _IndexProvenanceError( + f"{provenance_label} manifest is missing: {manifest_path}" + ) from exc + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise _IndexProvenanceError( + f"{provenance_label} manifest cannot be read: {manifest_path}: {exc}" + ) from exc + + if not isinstance(payload, dict) or set(payload) != _PROVENANCE_KEYS: + raise _IndexProvenanceError( + f"{provenance_label} manifest has an unsupported schema" + ) + return payload + + +def _validate_provenance_identity( + payload: Dict[str, Any], + expected_build_parameters: Dict[str, Any], + provenance_label: str, +) -> Tuple[Dict[str, Any], str, str]: + if ( + type(payload["schema_version"]) is not int + or payload["schema_version"] != _PROVENANCE_SCHEMA_VERSION + ): + raise _IndexProvenanceError( + f"{provenance_label} manifest has an unsupported schema version" + ) + expected_codec = expected_build_parameters["codec"] + if payload["codec"] != expected_codec: + raise _IndexProvenanceError( + f"{provenance_label} codec does not match the requested codec: " + f"{payload['codec']!r} != {expected_codec!r}" + ) + try: + stored_build_parameters = _normalize_build_params( + payload["build_parameters"] + ) + except (TypeError, ValueError) as exc: + raise _IndexProvenanceError( + f"{provenance_label} has invalid build parameters: {exc}" + ) from exc + if stored_build_parameters != payload["build_parameters"]: + raise _IndexProvenanceError( + f"{provenance_label} build parameters are not canonical" + ) + if stored_build_parameters != expected_build_parameters: + raise _IndexProvenanceError( + f"{provenance_label} build parameters do not match the " + "requested index configuration" + ) + + writer_policy = payload["writer_policy"] + if writer_policy != _EXPECTED_WRITER_POLICY[expected_codec]: + raise _IndexProvenanceError( + f"{provenance_label} does not record the expected writer policy" + ) + compound_file_policy = payload["compound_file_policy"] + if compound_file_policy != _COMPOUND_FILE_POLICY[expected_codec]: + raise _IndexProvenanceError( + f"{provenance_label} does not record the expected compound-file policy" + ) + return stored_build_parameters, writer_policy, compound_file_policy + + +def _validate_provenance_shape( + payload: Dict[str, Any], + provenance_label: str, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> Tuple[int, int]: + vector_count = _require_positive_int( + payload["vector_count"], provenance_label, "vector_count" + ) + if vector_count < 2: + raise _IndexProvenanceError( + f"{provenance_label} vector count must be at least two" + ) + dimensions = _require_positive_int( + payload["dimensions"], provenance_label, "dimensions" + ) + if dimensions > _MAX_DIMENSIONS: + raise _IndexProvenanceError( + f"{provenance_label} dimensions exceed the supported maximum: " + f"{dimensions} > {_MAX_DIMENSIONS}" + ) + if ( + expected_vector_count is not None + and vector_count != expected_vector_count + ): + raise _IndexProvenanceError( + f"{provenance_label} vector count does not match the dataset: " + f"{vector_count} != {expected_vector_count}" + ) + if expected_dimensions is not None and dimensions != expected_dimensions: + raise _IndexProvenanceError( + f"{provenance_label} dimensions do not match the dataset: " + f"{dimensions} != {expected_dimensions}" + ) + return vector_count, dimensions + + +def _validate_provenance_segment_count( + payload: Dict[str, Any], provenance_label: str +) -> int: + return _require_positive_int( + payload["segment_count"], provenance_label, "segment_count" + ) + + +def _validate_commit_fingerprints( + stored_fingerprints: Any, provenance_label: str +) -> List[Dict[str, str]]: + if not isinstance(stored_fingerprints, list) or not stored_fingerprints: + raise _IndexProvenanceError( + f"{provenance_label} has no Lucene commit fingerprints" + ) + names = [] + for fingerprint in stored_fingerprints: + name, _ = _validate_commit_fingerprint(fingerprint, provenance_label) + names.append(name) + if names != sorted(set(names)): + raise _IndexProvenanceError( + f"{provenance_label} commit fingerprints must be unique and sorted" + ) + return stored_fingerprints + + +def _validate_commit_fingerprint( + fingerprint: Any, provenance_label: str +) -> Tuple[str, str]: + if not isinstance(fingerprint, dict) or set(fingerprint) != { + "name", + "sha256", + }: + raise _IndexProvenanceError( + f"{provenance_label} has a malformed Lucene commit fingerprint" + ) + + name = _validate_commit_filename(fingerprint["name"], provenance_label) + digest = _validate_sha256_digest(fingerprint["sha256"], provenance_label) + return name, digest + + +def _validate_commit_filename(value: Any, provenance_label: str) -> str: + if not isinstance(value, str): + raise _IndexProvenanceError( + f"{provenance_label} commit filename must be a string" + ) + if not value.startswith("segments_"): + raise _IndexProvenanceError( + f"{provenance_label} commit filename must start with 'segments_'" + ) + if Path(value).name != value: + raise _IndexProvenanceError( + f"{provenance_label} commit filename must not contain a path" + ) + return value + + +def _validate_sha256_digest(value: Any, provenance_label: str) -> str: + if not isinstance(value, str): + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must be a string" + ) + if len(value) != 64: + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must contain 64 characters" + ) + if not set(value).issubset(_SHA256_HEX_DIGITS): + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must be lowercase hexadecimal" + ) + return value + + +def _verify_index_provenance( + index_path: Path, + expectation: _ProvenanceExpectation, +) -> _IndexProvenanceVerification: + payload = _read_index_provenance( + index_path / expectation.manifest_name, expectation.label + ) + ( + build_parameters, + writer_policy, + compound_file_policy, + ) = _validate_provenance_identity( + payload, expectation.build_parameters, expectation.label + ) + vector_count, dimensions = _validate_provenance_shape( + payload, + expectation.label, + expectation.vector_count, + expectation.dimensions, + ) + segment_count = _validate_provenance_segment_count( + payload, expectation.label + ) + if build_parameters.get("direct_single_segment") and segment_count != 1: + raise _IndexProvenanceError( + f"{expectation.label} direct_single_segment provenance must " + "record exactly one segment" + ) + stored_fingerprints = _validate_commit_fingerprints( + payload["commit_fingerprints"], expectation.label + ) + try: + current_fingerprints = _commit_fingerprints(index_path) + except (OSError, RuntimeError) as exc: + raise _IndexProvenanceError( + f"{expectation.label} cannot fingerprint the Lucene commit: {exc}" + ) from exc + if stored_fingerprints != current_fingerprints: + raise _IndexProvenanceError( + f"{expectation.label} does not match the current Lucene commit" + ) + + return _IndexProvenanceVerification( + codec=expectation.codec, + build_parameters=build_parameters, + writer_policy=writer_policy, + compound_file_policy=compound_file_policy, + vector_count=vector_count, + dimensions=dimensions, + segment_count=segment_count, + commit_file_count=len(current_fingerprints), + ) + + +def _verify_hnsw_provenance( + index_path: Path, + expected_codec: str, + *, + expected_build_parameters: Optional[Dict[str, Any]] = None, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + build_parameters = _normalize_build_params( + expected_build_parameters or {"codec": expected_codec} + ) + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=expected_codec, + build_parameters=build_parameters, + manifest_name=_HNSW_PROVENANCE_FILE, + label="HNSW provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _verify_cagra_provenance( + index_path: Path, + *, + expected_build_parameters: Optional[Dict[str, Any]] = None, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + build_parameters = _normalize_build_params( + expected_build_parameters or {"codec": _CAGRA_CODEC} + ) + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=_CAGRA_CODEC, + build_parameters=build_parameters, + manifest_name=_CAGRA_PROVENANCE_FILE, + label="CAGRA provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _validate_subset_size(value: Any) -> Optional[int]: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError( + f"subset_size must be a positive integer, got {value!r}" + ) + return value + + +def _expected_training_shape(dataset: Dataset) -> Optional[Tuple[int, int]]: + vectors = dataset.loaded_training_vectors + if vectors is not None: + vectors = np.asarray(vectors) + if vectors.ndim != 2: + raise ValueError( + "training_vectors must be a two-dimensional array" + ) + if vectors.dtype != np.float32: + raise TypeError( + f"training_vectors must use float32 values, got {vectors.dtype}" + ) + rows, dimensions = vectors.shape + elif dataset.base_file: + dtype = np.dtype(dtype_from_filename(dataset.base_file)) + if dtype != np.float32: + raise TypeError( + f"training_vectors must use float32 values, got {dtype}" + ) + rows, dimensions, _ = read_bin_header( + dataset.base_file, itemsize=dtype.itemsize + ) + subset_size = _validate_subset_size( + dataset.metadata.get("subset_size") + ) + if subset_size is not None: + rows = min(rows, subset_size) + else: + return None + + if rows < 1 or dimensions < 1: + raise ValueError("training_vectors must contain at least one vector") + if dimensions > _MAX_DIMENSIONS: + raise ValueError( + f"training_vectors dimensions must not exceed {_MAX_DIMENSIONS}, " + f"got {dimensions}" + ) + return int(rows), int(dimensions) + + +@dataclass(frozen=True) +class _SearchHit: + document_id: int + score: float + + +@dataclass(frozen=True) +class _RuntimeSearchResult: + hits: List[List[_SearchHit]] + batch_latencies_ms: List[float] + index_dimensions: int + document_count: int + + +@dataclass(frozen=True) +class _ProcessedSearchResult: + neighbors: np.ndarray + distances: np.ndarray + search_time_ms: float + queries_per_second: float + latency_seconds: float + latency_percentiles: Dict[str, float] + num_batches: int + + +@dataclass(frozen=True) +class _SearchInputs: + query_vectors: np.ndarray + expected_training_shape: Optional[Tuple[int, int]] + + +@dataclass(frozen=True) +class _SearchPlan: + index_path: Path + codec_name: str + build_parameters: Dict[str, Any] + search_parameters: Dict[str, Any] + num_candidates: int + k: int + batch_size: int + mode: str + + +@dataclass(frozen=True) +class _BuildCodec: + codec_name: str + java_codec: Any + writer_policy: str + build_parameters: Dict[str, Any] + + @property + def direct_single_segment(self) -> bool: + return bool(self.build_parameters.get("direct_single_segment", False)) + + +@dataclass(frozen=True) +class _IndexTopology: + segment_document_counts: Tuple[int, ...] + segment_vector_counts: Tuple[int, ...] + + @property + def segment_count(self) -> int: + return len(self.segment_document_counts) + + def validate(self, expected_vector_count: int) -> None: + if not self.segment_document_counts: + raise RuntimeError("Lucene index contains no committed segments") + if sum(self.segment_document_counts) != expected_vector_count: + raise RuntimeError( + "Lucene segment document counts do not match the indexed " + f"vectors: {sum(self.segment_document_counts)} != " + f"{expected_vector_count}" + ) + if self.segment_document_counts != self.segment_vector_counts: + raise RuntimeError( + "Lucene segment document and vector counts do not match: " + f"{self.segment_document_counts} != " + f"{self.segment_vector_counts}" + ) + + def require_direct_single_segment( + self, expected_vector_count: int + ) -> None: + self.validate(expected_vector_count) + if self.segment_count != 1: + raise RuntimeError( + "direct_single_segment requested one committed Lucene " + f"segment, found {self.segment_count}" + ) + + def to_metadata(self) -> Dict[str, Any]: + return { + "segment_count": self.segment_count, + "document_count": sum(self.segment_document_counts), + "vector_count": sum(self.segment_vector_counts), + "segment_document_counts": list(self.segment_document_counts), + "segment_vector_counts": list(self.segment_vector_counts), + } + + +class _ExistingIndexAction(Enum): + BUILD = "build" + REUSE = "reuse" + REJECT = "reject" + + +@dataclass(frozen=True) +class _ExistingIndexDecision: + action: _ExistingIndexAction + result: Optional[BuildResult] = None + + @classmethod + def build(cls) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.BUILD) + + @classmethod + def reuse(cls, result: BuildResult) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.REUSE, result=result) + + @classmethod + def reject(cls, result: BuildResult) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.REJECT, result=result) + + def completed_result(self) -> BuildResult: + if self.result is None: + raise RuntimeError( + "Existing-index decision is missing its build result" + ) + return self.result + + +def _validate_runtime_search_result( + runtime_result: _RuntimeSearchResult, + provenance: _IndexProvenanceVerification, + *, + query_count: int, +) -> None: + if runtime_result.document_count != provenance.vector_count: + raise RuntimeError( + "Lucene document count does not match index provenance: " + f"{runtime_result.document_count} != {provenance.vector_count}" + ) + if runtime_result.index_dimensions != provenance.dimensions: + raise RuntimeError( + "Lucene index dimensions do not match index provenance: " + f"{runtime_result.index_dimensions} != {provenance.dimensions}" + ) + if len(runtime_result.hits) != query_count: + raise RuntimeError( + "PyLucene returned an unexpected number of query results: " + f"{len(runtime_result.hits)} != {query_count}" + ) + + +def _validate_search_hit( + hit: _SearchHit, + *, + query_id: int, + document_count: int, + seen_document_ids: set[int], +) -> None: + if not 0 <= hit.document_id < document_count: + raise RuntimeError( + "PyLucene returned an out-of-range stored ID for " + f"query {query_id}: {hit.document_id}" + ) + if hit.document_id in seen_document_ids: + raise RuntimeError( + "PyLucene returned a duplicate stored ID for " + f"query {query_id}: {hit.document_id}" + ) + if not np.isfinite(hit.score): + raise RuntimeError( + f"PyLucene returned a non-finite score for query {query_id}: {hit.score}" + ) + if not 0.0 <= hit.score <= 1.0: + raise RuntimeError( + "PyLucene returned a score outside the Euclidean range " + f"[0, 1] for query {query_id}: {hit.score}" + ) + seen_document_ids.add(hit.document_id) + + +def _convert_search_hits( + runtime_result: _RuntimeSearchResult, *, query_count: int, k: int +) -> Tuple[np.ndarray, np.ndarray]: + neighbors = np.full((query_count, k), -1, dtype=np.int64) + distances = np.full((query_count, k), np.inf, dtype=np.float32) + for query_id, hits in enumerate(runtime_result.hits): + query_neighbors, query_distances = _convert_query_hits( + hits, + query_id=query_id, + document_count=runtime_result.document_count, + k=k, + ) + hit_count = len(query_neighbors) + neighbors[query_id, :hit_count] = query_neighbors + distances[query_id, :hit_count] = query_distances + return neighbors, distances + + +def _convert_query_hits( + hits: List[_SearchHit], + *, + query_id: int, + document_count: int, + k: int, +) -> Tuple[List[int], List[float]]: + if len(hits) > min(k, document_count): + raise RuntimeError( + f"PyLucene returned too many hits for query {query_id}: {len(hits)}" + ) + + neighbors = [] + distances = [] + seen_document_ids = set() + for hit in hits: + _validate_search_hit( + hit, + query_id=query_id, + document_count=document_count, + seen_document_ids=seen_document_ids, + ) + neighbors.append(hit.document_id) + distances.append(_score_to_squared_euclidean(hit.score)) + return neighbors, distances + + +def _validate_batch_latencies( + runtime_result: _RuntimeSearchResult, + *, + query_count: int, + batch_size: int, +) -> Tuple[np.ndarray, int]: + latencies = np.asarray(runtime_result.batch_latencies_ms, dtype=np.float64) + num_batches = (query_count + batch_size - 1) // batch_size + if latencies.shape != (num_batches,): + raise RuntimeError( + "PyLucene returned an unexpected number of batch latencies: " + f"{latencies.size} != {num_batches}" + ) + if not np.isfinite(latencies).all() or np.any(latencies < 0.0): + raise RuntimeError( + "PyLucene returned invalid batch latency measurements" + ) + return latencies, num_batches + + +def _process_search_result( + runtime_result: _RuntimeSearchResult, + provenance: _IndexProvenanceVerification, + *, + query_count: int, + k: int, + batch_size: int, +) -> _ProcessedSearchResult: + """Validate the Java boundary result and convert it to benchmark arrays.""" + _validate_runtime_search_result( + runtime_result, provenance, query_count=query_count + ) + neighbors, distances = _convert_search_hits( + runtime_result, query_count=query_count, k=k + ) + latencies, num_batches = _validate_batch_latencies( + runtime_result, query_count=query_count, batch_size=batch_size + ) + + search_time_ms = float(latencies.sum()) + return _ProcessedSearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=search_time_ms, + queries_per_second=( + query_count / (search_time_ms / 1000.0) + if search_time_ms > 0.0 + else 0.0 + ), + latency_seconds=float(latencies.mean()) / 1000.0, + latency_percentiles={ + "p50": float(np.percentile(latencies, 50)), + "p95": float(np.percentile(latencies, 95)), + "p99": float(np.percentile(latencies, 99)), + }, + num_batches=num_batches, + ) + + +@dataclass(frozen=True) +class _CagraIndexVerification: + segment_count: int + field_count: int + vector_count: int + dimensions: int + + def to_metadata(self) -> Dict[str, Union[str, int]]: + return { + "status": "cagra-only", + "segment_count": self.segment_count, + "field_count": self.field_count, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + } + + +class _CagraVerificationError(RuntimeError): + """Raised when persisted metadata cannot prove a CAGRA-only index.""" + + +@dataclass(frozen=True) +class _CagraFieldMetadata: + field_number: int + dimensions: int + vector_count: int + cagra_offset: int + cagra_length: int + + +@dataclass(frozen=True) +class _RawCagraFieldMetadata: + encoding: int + similarity: int + dimensions: int + vector_count: int + cagra_offset: int + cagra_length: int + brute_force_offset: int + brute_force_length: int + + +@dataclass(frozen=True) +class _CagraFieldSource: + metadata: _CagraFieldMetadata + metadata_file: str + + +@dataclass(frozen=True) +class _CagraDataFileContext: + index_path: Path + directory: Any + segment_info: Any + suffix: str + metadata_file: str + + @property + def data_file(self) -> str: + stem = self.metadata_file[: -len(_CAGRA_META_EXTENSION)] + return stem + _CAGRA_INDEX_EXTENSION + + @property + def data_path(self) -> Path: + return self.index_path / self.data_file + + +@dataclass(frozen=True) +class _CagraSegmentVerification: + field_count: int + vector_count: int + dimensions: frozenset[int] + + +class _CagraIndexVerifier: + """Own persisted CAGRA inspection and its verified-file cache.""" + + def __init__( + self, + *, + attach_current_thread: Callable[[], None], + paths: Any, + codec_util: Any, + field_info: Any, + segment_commit_info: Any, + segment_infos: Any, + vector_encoding: Any, + vector_similarity_function: Any, + fs_directory: Any, + io_context: Any, + ) -> None: + self._attach_current_thread = attach_current_thread + self.Paths = paths + self.CodecUtil = codec_util + self.FieldInfo = field_info + self.SegmentCommitInfo = segment_commit_info + self.SegmentInfos = segment_infos + self.VectorEncoding = vector_encoding + self.VectorSimilarityFunction = vector_similarity_function + self.FSDirectory = fs_directory + self.IOContext = io_context + self._verified_data_files: set[_FileSignature] = set() + + @staticmethod + def _segment_suffix(segment_name: str, metadata_file: str) -> str: + stem = metadata_file[: -len(_CAGRA_META_EXTENSION)] + if stem == segment_name: + return "" + + prefix = f"{segment_name}_" + if not stem.startswith(prefix) or len(stem) == len(prefix): + raise _CagraVerificationError( + "CAGRA-only verification found a metadata filename that " + f"does not match segment {segment_name!r}: {metadata_file!r}" + ) + return stem[len(prefix) :] + + @classmethod + def _read_cagra_field( + cls, + metadata_input: Any, + metadata_file: str, + field_number: int, + ) -> Optional[_CagraFieldMetadata]: + raw_field = cls._decode_cagra_field(metadata_input) + return cls._validate_cagra_field_metadata( + raw_field, metadata_file, field_number + ) + + @staticmethod + def _decode_cagra_field(metadata_input: Any) -> _RawCagraFieldMetadata: + return _RawCagraFieldMetadata( + encoding=int(metadata_input.readInt()), + similarity=int(metadata_input.readInt()), + dimensions=int(metadata_input.readInt()), + vector_count=int(metadata_input.readInt()), + cagra_offset=int(metadata_input.readVLong()), + cagra_length=int(metadata_input.readVLong()), + brute_force_offset=int(metadata_input.readVLong()), + brute_force_length=int(metadata_input.readVLong()), + ) + + @classmethod + def _validate_cagra_field_metadata( + cls, + field: _RawCagraFieldMetadata, + metadata_file: str, + field_number: int, + ) -> Optional[_CagraFieldMetadata]: + if field.encoding != _FLOAT32_ENCODING_ORDINAL: + raise _CagraVerificationError( + "CAGRA-only verification found unsupported vector " + f"encoding ordinal {field.encoding} in {metadata_file!r}" + ) + if field.similarity != _EUCLIDEAN_SIMILARITY_ORDINAL: + raise _CagraVerificationError( + "CAGRA-only verification found unsupported similarity " + f"ordinal {field.similarity} in {metadata_file!r}" + ) + if not 1 <= field.dimensions <= _MAX_DIMENSIONS: + raise _CagraVerificationError( + "CAGRA-only verification found invalid field metadata " + f"in {metadata_file!r}" + ) + if field.vector_count < 0: + raise _CagraVerificationError( + "CAGRA-only verification found invalid field metadata " + f"in {metadata_file!r}" + ) + if field.vector_count == 0: + if cls._empty_cagra_field_has_data(field): + raise _CagraVerificationError( + "CAGRA-only verification found index data for an " + f"empty field in {metadata_file!r}" + ) + return None + if field.brute_force_length != 0: + raise _CagraVerificationError( + "CAGRA-only verification found a persisted brute-force " + f"index for field {field_number} in {metadata_file!r}" + ) + if field.cagra_length <= 0: + raise _CagraVerificationError( + "CAGRA-only verification found no persisted CAGRA index " + f"for field {field_number} in {metadata_file!r}" + ) + return _CagraFieldMetadata( + field_number=field_number, + dimensions=field.dimensions, + vector_count=field.vector_count, + cagra_offset=field.cagra_offset, + cagra_length=field.cagra_length, + ) + + @staticmethod + def _empty_cagra_field_has_data(field: _RawCagraFieldMetadata) -> bool: + return any( + ( + field.cagra_offset, + field.cagra_length, + field.brute_force_offset, + field.brute_force_length, + ) + ) + + @classmethod + def _read_cagra_fields( + cls, metadata_input: Any, metadata_file: str + ) -> List[_CagraFieldMetadata]: + field_numbers = set() + fields = [] + while True: + field_number = int(metadata_input.readInt()) + if field_number == -1: + return fields + if field_number < 0 or field_number in field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found an invalid field number " + f"{field_number} in {metadata_file!r}" + ) + + field_numbers.add(field_number) + field = cls._read_cagra_field( + metadata_input, metadata_file, field_number + ) + if field is not None: + fields.append(field) + + def _read_segment_field_infos( + self, directory: Any, segment_info: Any + ) -> Any: + if segment_info.hasFieldUpdates(): + raise _CagraVerificationError( + "CAGRA-only verification does not accept segments with " + "field-info updates" + ) + try: + return ( + segment_info.info.getCodec() + .fieldInfosFormat() + .read( + directory, + segment_info.info, + "", + self.IOContext.READONCE, + ) + ) + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read Lucene field " + f"metadata for segment {segment_info.info.name!r}: {exc}" + ) from exc + + @classmethod + def _verify_cagra_payload_coverage( + cls, + fields: List[_CagraFieldMetadata], + payload_start: int, + payload_end: int, + data_file: str, + ) -> None: + intervals = [] + for field in fields: + intervals.append( + ( + field.cagra_offset, + field.cagra_offset + field.cagra_length, + ) + ) + intervals.sort() + + if not intervals: + if payload_start != payload_end: + raise _CagraVerificationError( + f"CAGRA data file contains unreferenced payload: {data_file!r}" + ) + return + + coverage_error = f"CAGRA metadata ranges do not exactly cover data file {data_file!r}" + expected_offset = payload_start + for interval_start, interval_end in intervals: + if interval_start != expected_offset: + raise _CagraVerificationError(coverage_error) + expected_offset = interval_end + if expected_offset != payload_end: + raise _CagraVerificationError(coverage_error) + + def _verify_checksum_with_signature_cache( + self, + data_path: Path, + data_input: Any, + signature_before: Optional[_FileSignature], + data_file: str, + ) -> None: + cache_hit = ( + signature_before is not None + and signature_before in self._verified_data_files + ) + if cache_hit: + self.CodecUtil.retrieveChecksum(data_input) + else: + self.CodecUtil.checksumEntireFile(data_input) + + if signature_before is None: + return + signature_after = _file_signature(data_path) + if signature_after != signature_before: + raise _CagraVerificationError( + f"CAGRA data file changed during verification: {data_file!r}" + ) + file_age = time.time_ns() - signature_before.changed_at_ns + if cache_hit or file_age < _CAGRA_CACHE_MIN_FILE_AGE_NS: + return + + self._verified_data_files = { + signature + for signature in self._verified_data_files + if signature.resolved_path != signature_before.resolved_path + } + self._verified_data_files.add(signature_before) + + def _verify_cagra_data_file( + self, + context: _CagraDataFileContext, + fields: List[_CagraFieldMetadata], + ) -> None: + try: + signature_before = _file_signature(context.data_path) + except OSError: + signature_before = None + + try: + with _CleanupStack() as cleanups: + data_input = context.directory.openInput( + context.data_file, self.IOContext.READONCE + ) + cleanups.add("close CAGRA data input", data_input.close) + self.CodecUtil.checkIndexHeader( + data_input, + _CAGRA_INDEX_CODEC_NAME, + _CAGRA_INDEX_VERSION, + _CAGRA_INDEX_VERSION, + context.segment_info.info.getId(), + context.suffix, + ) + payload_start = int(data_input.getFilePointer()) + payload_end = int(data_input.length()) - int( + self.CodecUtil.footerLength() + ) + if payload_end < payload_start: + raise _CagraVerificationError( + f"CAGRA data file is truncated: {context.data_file!r}" + ) + self._verify_cagra_payload_coverage( + fields, + payload_start, + payload_end, + context.data_file, + ) + self._verify_checksum_with_signature_cache( + context.data_path, + data_input, + signature_before, + context.data_file, + ) + except _CagraVerificationError: + raise + except Exception as exc: + raise _CagraVerificationError( + f"CAGRA-only verification cannot read {context.data_file!r}: {exc}" + ) from exc + + def _verify_field_against_lucene_metadata( + self, + field_infos: Any, + field: _CagraFieldMetadata, + metadata_file: str, + ) -> None: + field_info = field_infos.fieldInfo(field.field_number) + if field_info is None: + raise _CagraVerificationError( + "CAGRA-only verification found an unknown field number " + f"{field.field_number} in {metadata_file!r}" + ) + field_name = str(field_info.getName()) + if field_name != _VECTOR_FIELD: + raise _CagraVerificationError( + "CAGRA-only verification found data for unexpected field " + f"{field_name!r} in {metadata_file!r}" + ) + if int(field_info.getVectorDimension()) != field.dimensions: + raise _CagraVerificationError( + "CAGRA-only verification found dimensions inconsistent " + f"with Lucene field metadata in {metadata_file!r}" + ) + if field_info.getVectorEncoding() != self.VectorEncoding.FLOAT32: + raise _CagraVerificationError( + "CAGRA-only verification found non-FLOAT32 Lucene field " + f"metadata in {metadata_file!r}" + ) + if ( + field_info.getVectorSimilarityFunction() + != self.VectorSimilarityFunction.EUCLIDEAN + ): + raise _CagraVerificationError( + "CAGRA-only verification found non-Euclidean Lucene field " + f"metadata in {metadata_file!r}" + ) + + @staticmethod + def _validate_segment_deletions( + segment_info: Any, segment_name: str + ) -> None: + deletion_count = int(segment_info.getDelCount()) + soft_deletion_count = int(segment_info.getSoftDelCount()) + if ( + segment_info.hasDeletions() + or deletion_count + or soft_deletion_count + ): + raise _CagraVerificationError( + "CAGRA-only verification does not accept committed " + f"deletions in segment {segment_name!r}: " + f"deleted={deletion_count}, " + f"soft_deleted={soft_deletion_count}" + ) + + @staticmethod + def _cagra_metadata_files( + segment_info: Any, segment_name: str + ) -> List[str]: + metadata_files = [] + for file_name in segment_info.files(): + file_name = str(file_name) + if file_name.endswith(_CAGRA_META_EXTENSION): + metadata_files.append(file_name) + metadata_files.sort() + if not metadata_files: + raise _CagraVerificationError( + "CAGRA-only verification found no " + f"{_CAGRA_META_EXTENSION} metadata for segment " + f"{segment_name!r}" + ) + return metadata_files + + def _read_and_verify_cagra_metadata_file( + self, + context: _CagraDataFileContext, + ) -> List[_CagraFieldMetadata]: + metadata_input = context.directory.openChecksumInput( + context.metadata_file + ) + try: + with _CleanupStack() as cleanups: + cleanups.add( + "close CAGRA metadata input", metadata_input.close + ) + self.CodecUtil.checkIndexHeader( + metadata_input, + _CAGRA_META_CODEC_NAME, + _CAGRA_META_VERSION, + _CAGRA_META_VERSION, + context.segment_info.info.getId(), + context.suffix, + ) + fields = self._read_cagra_fields( + metadata_input, context.metadata_file + ) + self.CodecUtil.checkFooter(metadata_input) + return fields + except _CagraVerificationError: + raise + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read " + f"{context.metadata_file!r} as cuVS-Lucene metadata " + f"format v{_CAGRA_META_VERSION}: {exc}" + ) from exc + + def _lucene_vector_field_numbers(self, field_infos: Any) -> set[int]: + field_numbers = set() + for raw_field_info in field_infos: + field_info = self.FieldInfo.cast_(raw_field_info) + if int(field_info.getVectorDimension()) > 0: + field_numbers.add(int(field_info.number)) + return field_numbers + + def _verify_cagra_segment( + self, + index_path: Path, + directory: Any, + segment_info: Any, + ) -> _CagraSegmentVerification: + segment_name = str(segment_info.info.name) + self._validate_segment_deletions(segment_info, segment_name) + metadata_files = self._cagra_metadata_files(segment_info, segment_name) + field_infos = self._read_segment_field_infos(directory, segment_info) + + # Verify every metadata/data pair before interpreting fields across + # the segment. This preserves fail-closed file-validation ordering. + verified_field_sources = [] + for metadata_file in metadata_files: + suffix = self._segment_suffix(segment_name, metadata_file) + context = _CagraDataFileContext( + index_path=index_path, + directory=directory, + segment_info=segment_info, + suffix=suffix, + metadata_file=metadata_file, + ) + fields = self._read_and_verify_cagra_metadata_file(context) + self._verify_cagra_data_file(context, fields) + for field in fields: + verified_field_sources.append( + _CagraFieldSource( + metadata=field, + metadata_file=metadata_file, + ) + ) + + # Compare the verified fields with Lucene's segment-wide view and + # aggregate the values needed for index-wide validation. + field_numbers = set() + vector_count = 0 + dimensions = set() + for field_source in verified_field_sources: + field = field_source.metadata + if field.field_number in field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found duplicate " + f"field {field.field_number} across metadata " + f"files for segment {segment_name!r}" + ) + field_numbers.add(field.field_number) + self._verify_field_against_lucene_metadata( + field_infos, field, field_source.metadata_file + ) + vector_count += field.vector_count + dimensions.add(field.dimensions) + + lucene_field_numbers = self._lucene_vector_field_numbers(field_infos) + if field_numbers != lucene_field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found vector fields without " + "matching CAGRA metadata in segment " + f"{segment_name!r}: metadata={sorted(field_numbers)}, " + f"Lucene={sorted(lucene_field_numbers)}" + ) + + max_documents = int(segment_info.info.maxDoc()) + if vector_count != max_documents: + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{vector_count} vectors for " + f"{max_documents} documents in segment " + f"{segment_name!r}" + ) + return _CagraSegmentVerification( + field_count=len(field_numbers), + vector_count=vector_count, + dimensions=frozenset(dimensions), + ) + + @staticmethod + def _summarize_cagra_segments( + segments: List[_CagraSegmentVerification], + expected_vector_count: Optional[int], + expected_dimensions: Optional[int], + ) -> _CagraIndexVerification: + field_count = 0 + vector_count = 0 + dimensions = set() + for segment in segments: + field_count += segment.field_count + vector_count += segment.vector_count + dimensions.update(segment.dimensions) + if not segments or field_count == 0: + raise _CagraVerificationError( + "CAGRA-only verification found no nonempty vector fields" + ) + if len(dimensions) != 1: + raise _CagraVerificationError( + "CAGRA-only verification found inconsistent vector " + f"dimensions: {sorted(dimensions)}" + ) + + index_dimensions = next(iter(dimensions)) + if ( + expected_vector_count is not None + and vector_count != expected_vector_count + ): + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{vector_count} vectors; expected {expected_vector_count}" + ) + if ( + expected_dimensions is not None + and index_dimensions != expected_dimensions + ): + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{index_dimensions} dimensions; expected " + f"{expected_dimensions}" + ) + return _CagraIndexVerification( + segment_count=len(segments), + field_count=field_count, + vector_count=vector_count, + dimensions=index_dimensions, + ) + + def verify_index( + self, + index_path: Path, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> _CagraIndexVerification: + """Verify every committed vector field is persisted as CAGRA only.""" + self._attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + verified_segments = [] + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + segment_infos = self.SegmentInfos.readLatestCommit(directory) + for raw_segment_info in segment_infos: + segment_info = self.SegmentCommitInfo.cast_(raw_segment_info) + verified_segments.append( + self._verify_cagra_segment( + index_path, directory, segment_info + ) + ) + return self._summarize_cagra_segments( + verified_segments, + expected_vector_count, + expected_dimensions, + ) + + +class _PyLuceneRuntime: + """Own generated PyLucene/Lucene bindings and index operations.""" + + def __init__(self, lucene: Any): + from java.lang import Class, System + from java.nio.file import Paths + from org.apache.lucene.codecs import Codec, CodecUtil + from org.apache.lucene.document import ( + Document, + KnnFloatVectorField, + StoredField, + ) + from org.apache.lucene.index import ( + DirectoryReader, + FieldInfo, + IndexWriter, + IndexWriterConfig, + NoMergePolicy, + SegmentCommitInfo, + SegmentInfos, + VectorEncoding, + VectorSimilarityFunction, + ) + from org.apache.lucene.search import ( + IndexSearcher, + KnnFloatVectorQuery, + ) + from org.apache.lucene.store import FSDirectory, IOContext + + self.lucene = lucene + self.Class = Class + self.System = System + self.Paths = Paths + self.Codec = Codec + self.Document = Document + self.KnnFloatVectorField = KnnFloatVectorField + self.StoredField = StoredField + self.DirectoryReader = DirectoryReader + self.IndexWriter = IndexWriter + self.IndexWriterConfig = IndexWriterConfig + self.NoMergePolicy = NoMergePolicy + self.VectorSimilarityFunction = VectorSimilarityFunction + self.IndexSearcher = IndexSearcher + self.KnnFloatVectorQuery = KnnFloatVectorQuery + self.FSDirectory = FSDirectory + self.IOContext = IOContext + self._codec_cache: Dict[str, Any] = {} + + # Resolve the base cuvs-java provider before codec construction so + # classpath failures have a direct error location. + self.Class.forName("com.nvidia.cuvs.spi.JDKProvider") + self._cagra_verifier = _CagraIndexVerifier( + attach_current_thread=self.attach_current_thread, + paths=Paths, + codec_util=CodecUtil, + field_info=FieldInfo, + segment_commit_info=SegmentCommitInfo, + segment_infos=SegmentInfos, + vector_encoding=VectorEncoding, + vector_similarity_function=VectorSimilarityFunction, + fs_directory=FSDirectory, + io_context=IOContext, + ) + + @classmethod + def create(cls, config: Dict[str, Any]) -> "_PyLuceneRuntime": + return cls(_initialize_pylucene(config)) + + @property + def pylucene_version(self) -> str: + return str(getattr(self.lucene, "VERSION", "unknown")) + + def attach_current_thread(self) -> None: + vm_environment = self.lucene.getVMEnv() + if vm_environment is None: + raise RuntimeError("PyLucene JVM is not initialized") + vm_environment.attachCurrentThread() + + def resolve_codec(self, codec_name: str) -> Any: + self.attach_current_thread() + cached = self._codec_cache.get(codec_name) + if cached is not None: + return cached + + available_codecs = self.Codec.availableCodecs() + if not available_codecs.contains(codec_name): + available = ", ".join(str(name) for name in available_codecs) + raise RuntimeError( + f"{codec_name} was not advertised by Lucene SPI. " + f"Available codecs: {available}" + ) + codec = self.Codec.forName(codec_name) + self._validate_codec(codec, codec_name) + self._codec_cache[codec_name] = codec + return codec + + def resolve_configured_hnsw_codec( + self, m: int, ef_construction: int + ) -> Any: + self.attach_current_thread() + configured_values = { + M_PROPERTY: str(m), + EF_CONSTRUCTION_PROPERTY: str(ef_construction), + } + with _CONFIGURED_CODEC_LOCK: + with _CleanupStack() as cleanups: + for name, value in configured_values.items(): + previous = self.System.getProperty(name) + cleanups.add( + f"restore Java system property {name}", + lambda name=name, previous=previous: ( + _restore_java_property(self.System, name, previous) + ), + ) + self.System.setProperty(name, value) + reflected_codec = self.Class.forName( + CONFIGURED_HNSW_CODEC_CLASS + ).newInstance() + codec = self.Codec.cast_(reflected_codec) + self._validate_codec(codec, _HNSW_CODEC) + expected_diagnostics = f"PyLuceneConfiguredHnswCodec(m={m}, efConstruction={ef_construction})" + if str(codec) != expected_diagnostics: + raise RuntimeError( + "Configured PyLucene codec did not retain the requested " + f"parameters: expected {expected_diagnostics!r}, " + f"got {str(codec)!r}" + ) + return codec + + @staticmethod + def _validate_codec(codec: Any, codec_name: str) -> None: + if str(codec.getName()) != codec_name: + raise RuntimeError( + f"Requested codec {codec_name}, got {codec.getName()}" + ) + if codec.knnVectorsFormat() is None: + raise RuntimeError( + f"{codec_name} did not initialize a Lucene vector format" + ) + + def _java_float_array(self, vector: np.ndarray) -> Any: + return self.lucene.JArray("float")( + tuple(float(value) for value in vector) + ) + + def verify_cagra_index( + self, + index_path: Path, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> _CagraIndexVerification: + return self._cagra_verifier.verify_index( + index_path, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + + def build_index( + self, + index_path: Path, + vectors: np.ndarray, + build_codec: _BuildCodec, + ) -> _IndexTopology: + self.attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + writer_config = self._new_index_writer_config( + build_codec, vector_count=int(vectors.shape[0]) + ) + writer = self.IndexWriter(directory, writer_config) + self._write_and_close_index(writer, vectors) + topology = self._index_topology(directory) + if build_codec.direct_single_segment: + topology.require_direct_single_segment(int(vectors.shape[0])) + else: + topology.validate(int(vectors.shape[0])) + return topology + + def _new_index_writer_config( + self, build_codec: _BuildCodec, *, vector_count: int + ) -> Any: + writer_config = self.IndexWriterConfig() + writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) + writer_config.setCodec(build_codec.java_codec) + if build_codec.codec_name == _CAGRA_CODEC: + # The CAGRA verifier reads the codec's .vemc and .vcag files. + # Lucene controls compound files separately for flushes and merges. + writer_config.setUseCompoundFile(False) + writer_config.getMergePolicy().setNoCFSRatio(0.0) + if build_codec.direct_single_segment: + writer_config.setMaxBufferedDocs(vector_count + 1) + writer_config.setRAMBufferSizeMB( + float(self.IndexWriterConfig.DISABLE_AUTO_FLUSH) + ) + writer_config.setMergePolicy(self.NoMergePolicy.INSTANCE) + return writer_config + + def _index_topology(self, directory: Any) -> _IndexTopology: + reader = self.DirectoryReader.open(directory) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene topology reader", reader.close) + document_counts = [] + vector_counts = [] + for leaf_context in reader.leaves(): + leaf_reader = leaf_context.reader() + values = leaf_reader.getFloatVectorValues(_VECTOR_FIELD) + if values is None: + raise RuntimeError( + "Lucene segment contains no vector values for " + f"{_VECTOR_FIELD!r}" + ) + document_counts.append(int(leaf_reader.numDocs())) + vector_counts.append(int(values.size())) + return _IndexTopology( + segment_document_counts=tuple(document_counts), + segment_vector_counts=tuple(vector_counts), + ) + + def _write_and_close_index(self, writer: Any, vectors: np.ndarray) -> None: + try: + for document_id, vector in enumerate(vectors): + writer.addDocument(self._vector_document(document_id, vector)) + writer.commit() + except BaseException as write_error: + self._rollback_failed_write(writer, write_error) + raise + writer.close() + + @staticmethod + def _rollback_failed_write( + writer: Any, write_error: BaseException + ) -> None: + try: + writer.rollback() + except BaseException as rollback_error: + if isinstance(write_error, Exception): + raise rollback_error from write_error + write_error.add_note( + "IndexWriter rollback also failed: " + f"{type(rollback_error).__name__}: {rollback_error}" + ) + + def _vector_document(self, document_id: int, vector: np.ndarray) -> Any: + document = self.Document() + document.add(self.StoredField(_ID_FIELD, str(document_id))) + document.add( + self.KnnFloatVectorField( + _VECTOR_FIELD, + self._java_float_array(vector), + self.VectorSimilarityFunction.EUCLIDEAN, + ) + ) + return document + + @staticmethod + def _index_vector_dimensions(reader: Any) -> int: + dimensions = set() + for leaf_context in reader.leaves(): + values = leaf_context.reader().getFloatVectorValues(_VECTOR_FIELD) + if values is not None and values.size() > 0: + dimensions.add(int(values.dimension())) + if not dimensions: + raise RuntimeError( + f"Lucene index contains no {_VECTOR_FIELD!r} vectors" + ) + if len(dimensions) != 1: + raise RuntimeError( + f"Lucene index has inconsistent vector dimensions: {dimensions}" + ) + return dimensions.pop() + + def _search_vector( + self, + searcher: Any, + stored_fields: Any, + vector: np.ndarray, + k: int, + num_candidates: int, + ) -> List[_SearchHit]: + query = self.KnnFloatVectorQuery( + _VECTOR_FIELD, + self._java_float_array(vector), + num_candidates, + ) + score_docs = searcher.search(query, k).scoreDocs + hits = [] + for score_doc in score_docs: + stored_id = stored_fields.document(score_doc.doc).get(_ID_FIELD) + if stored_id is None: + raise RuntimeError( + f"Lucene document {score_doc.doc} has no stored ID" + ) + hits.append( + _SearchHit( + document_id=int(stored_id), + score=float(score_doc.score), + ) + ) + return hits + + def search_index( + self, + index_path: Path, + query_vectors: np.ndarray, + k: int, + batch_size: int, + num_candidates: Optional[int] = None, + ) -> _RuntimeSearchResult: + self.attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + reader = self.DirectoryReader.open(directory) + cleanups.add("close Lucene index reader", reader.close) + index_dimensions = self._index_vector_dimensions(reader) + if query_vectors.shape[1] != index_dimensions: + raise ValueError( + "Query vector dimensions do not match the Lucene index: " + f"{query_vectors.shape[1]} != {index_dimensions}" + ) + + document_count = int(reader.numDocs()) + if document_count < 1: + raise RuntimeError("Lucene index contains no documents") + lucene_k = min(k, document_count) + lucene_num_candidates = min( + num_candidates if num_candidates is not None else k, + document_count, + ) + searcher = self.IndexSearcher(reader) + stored_fields = searcher.storedFields() + + all_hits: List[List[_SearchHit]] = [] + batch_latencies_ms: List[float] = [] + for batch_start in range(0, query_vectors.shape[0], batch_size): + start = time.perf_counter() + batch_end = min( + batch_start + batch_size, query_vectors.shape[0] + ) + all_hits.extend( + self._search_vectors( + searcher, + stored_fields, + query_vectors[batch_start:batch_end], + lucene_k, + lucene_num_candidates, + ) + ) + batch_latencies_ms.append( + (time.perf_counter() - start) * 1000.0 + ) + + return _RuntimeSearchResult( + hits=all_hits, + batch_latencies_ms=batch_latencies_ms, + index_dimensions=index_dimensions, + document_count=document_count, + ) + + def _search_vectors( + self, + searcher: Any, + stored_fields: Any, + query_vectors: np.ndarray, + k: int, + num_candidates: int, + ) -> List[List[_SearchHit]]: + hits = [] + for vector in query_vectors: + hits.append( + self._search_vector( + searcher, + stored_fields, + vector, + k, + num_candidates, + ) + ) + return hits + + +class _SelectedAlgorithmGroup(NamedTuple): + algorithm: str + group: str + configuration: dict + metadata: dict + + +@dataclass(frozen=True) +class _GlobalSelectionScope: + algorithms: frozenset[str] + groups: frozenset[str] + + +@dataclass(frozen=True) +class _AlgorithmSelection: + algorithms: Optional[frozenset[str]] + groups: Optional[frozenset[str]] + explicit_groups: frozenset[Tuple[str, str]] + + def _resolve_global_scope( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> Optional[_GlobalSelectionScope]: + available_algorithms = frozenset(algorithm_configs) + if self.algorithms is not None: + unknown_algorithms = self.algorithms - available_algorithms + if unknown_algorithms: + names = ", ".join(sorted(unknown_algorithms)) + raise ValueError( + f"Unknown PyLucene algorithm selector(s): {names}" + ) + + candidate_algorithms = ( + self.algorithms + if self.algorithms is not None + else available_algorithms + ) + if self.groups is not None: + available_groups = set() + for algorithm in candidate_algorithms: + available_groups.update(algorithm_configs[algorithm]) + unknown_groups = self.groups - available_groups + if unknown_groups: + names = ", ".join(sorted(unknown_groups)) + raise ValueError( + f"Unknown PyLucene group selector(s): {names}" + ) + + has_global_selector = ( + self.algorithms is not None or self.groups is not None + ) + if not has_global_selector and self.explicit_groups: + return None + return _GlobalSelectionScope( + algorithms=candidate_algorithms, + groups=( + self.groups if self.groups is not None else frozenset({"base"}) + ), + ) + + def _validate_explicit_groups( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> None: + for algorithm, group in sorted(self.explicit_groups): + if algorithm not in algorithm_configs: + raise ValueError( + f"Unknown PyLucene algorithm in --algo-groups: {algorithm}" + ) + if group not in algorithm_configs[algorithm]: + raise ValueError( + f"Unknown PyLucene group for {algorithm}: {group}" + ) + + def _selected_pairs( + self, + algorithm_configs: Dict[str, Dict[str, Any]], + global_scope: Optional[_GlobalSelectionScope], + ) -> set[Tuple[str, str]]: + selected_pairs = set(self.explicit_groups) + if global_scope is None: + return selected_pairs + + for algorithm, groups in algorithm_configs.items(): + if algorithm not in global_scope.algorithms: + continue + for group in groups: + if group in global_scope.groups: + selected_pairs.add((algorithm, group)) + return selected_pairs + + def resolve( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> List[_SelectedAlgorithmGroup]: + global_scope = self._resolve_global_scope(algorithm_configs) + self._validate_explicit_groups(algorithm_configs) + selected_pairs = self._selected_pairs(algorithm_configs, global_scope) + + selected_groups = [] + for algorithm, groups in algorithm_configs.items(): + for group, group_config in groups.items(): + if (algorithm, group) not in selected_pairs: + continue + selected_groups.append( + _SelectedAlgorithmGroup( + algorithm=algorithm, + group=group, + configuration=group_config, + metadata={}, + ) + ) + return selected_groups + + +@dataclass(frozen=True) +class _BenchmarkConfigContext: + dataset: str + dataset_path: str + subset_scope: Optional[str] + runtime_config: Dict[str, Any] + + +class PyLuceneConfigLoader(ConfigLoader): + """Load PyLucene algorithm configurations.""" + + def __init__(self, config_path: Optional[Union[str, os.PathLike]] = None): + self.config_path = ( + os.fspath(config_path) + if config_path is not None + else os.path.join( + os.path.dirname(os.path.realpath(__file__)), "../config" + ) + ) + + @property + def backend_type(self) -> str: + return "pylucene" + + @staticmethod + def _parse_name_filter( + value: Optional[str], + ) -> Optional[frozenset[str]]: + if not value: + return None + + names = set() + for raw_name in value.split(","): + name = raw_name.strip() + if name: + names.add(name) + return frozenset(names) or None + + @staticmethod + def _parse_algorithm_group_filter( + value: Optional[str], + ) -> frozenset[Tuple[str, str]]: + selections = set() + if not value: + return frozenset() + + for selection in value.split(","): + algorithm, separator, group = selection.strip().partition(".") + if not separator or not algorithm or not group: + raise ValueError( + "algo_groups entries must use ., " + f"got {selection!r}" + ) + selections.add((algorithm, group)) + return frozenset(selections) + + @classmethod + def _selection_from_options( + cls, options: Dict[str, Any] + ) -> _AlgorithmSelection: + return _AlgorithmSelection( + algorithms=cls._parse_name_filter(options.get("algorithms")), + groups=cls._parse_name_filter(options.get("groups")), + explicit_groups=cls._parse_algorithm_group_filter( + options.get("algo_groups") + ), + ) + + def _match_backend_algorithm_config( + self, config: Any + ) -> Optional[Tuple[str, Dict[str, Any]]]: + if not isinstance(config, dict): + return None + + algorithm = config.get("name") + if not isinstance(algorithm, str) or not algorithm: + return None + + declared_backend = config.get("backend") + if declared_backend is None: + belongs_to_backend = algorithm.startswith("pylucene_") + else: + belongs_to_backend = declared_backend == self.backend_type + if not belongs_to_backend: + return None + return algorithm, config.get("groups", {}) + + def _load_algorithm_configs( + self, algorithm_files: List[str] + ) -> Dict[str, Dict[str, Any]]: + algorithm_configs = {} + for algorithm_file in algorithm_files: + matched_config = self._match_backend_algorithm_config( + self.load_yaml_file(algorithm_file) + ) + if matched_config is None: + continue + algorithm, groups = matched_config + # Later files are explicit overrides and determine output order. + algorithm_configs.pop(algorithm, None) + algorithm_configs[algorithm] = groups + return algorithm_configs + + def _discover_algo_groups( + self, + dataset_conf: dict, + dataset: str, + dataset_path: str, + **kwargs, + ) -> List[Tuple[str, str, dict, dict]]: + algorithm_files = self.gather_algorithm_configs( + self.config_path, kwargs.get("algorithm_configuration") + ) + algorithm_configs = self._load_algorithm_configs(algorithm_files) + selection = self._selection_from_options(kwargs) + return selection.resolve(algorithm_configs) + + @staticmethod + def _runtime_config(options: Dict[str, Any]) -> Dict[str, Any]: + runtime_keys = ( + "cuvs_java_jar", + "cuvs_lucene_jar", + "java_library_path", + "jvm_args", + ) + runtime_config = {} + for key in runtime_keys: + value = options.get(key) + if value is not None: + runtime_config[key] = value + return runtime_config + + @staticmethod + def _subset_scope(subset_size: Any) -> Optional[str]: + subset_size = _validate_subset_size(subset_size) + if subset_size is None: + return None + return f"subset{subset_size}" + + @staticmethod + def _effective_parameters( + build_combinations: List[Dict[str, Any]], + search_combinations: List[Dict[str, Any]], + *, + tune_mode: bool, + tune_build_params: Optional[Dict[str, Any]], + tune_search_params: Optional[Dict[str, Any]], + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + use_tuned_parameters = tune_mode and tune_build_params is not None + if not use_tuned_parameters: + return build_combinations, search_combinations + if not build_combinations: + raise ValueError("PyLucene tune mode requires build parameters") + build_params = { + **build_combinations[0], + **tune_build_params, + } + search_params = [tune_search_params] if tune_search_params else [{}] + return [build_params], search_params + + @staticmethod + def _index_label( + algorithm: str, + group: str, + subset_scope: Optional[str], + build_params: Dict[str, Any], + ) -> str: + identity = format_artifact_identity( + algorithm, group, subset_scope, description="PyLucene" + ) + parameters = _normalize_build_params(build_params) + label = f"{identity}[codec={parameters['codec']}]" + if parameters["codec"] == _HNSW_CODEC: + direct_single_segment = str( + parameters["direct_single_segment"] + ).lower() + label += ( + f"[m={parameters['m']}]" + f"[ef_construction={parameters['ef_construction']}]" + f"[direct_single_segment={direct_single_segment}]" + ) + return label + + @classmethod + def _benchmark_config( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: Dict[str, Any], + search_params: List[Dict[str, Any]], + ) -> BenchmarkConfig: + build_params = _normalize_build_params(build_params) + search_params = _normalize_search_params( + search_params, codec_name=build_params["codec"] + ) + dataset = validate_path_component( + context.dataset, "PyLucene dataset name" + ) + index_label = cls._index_label( + algorithm, group, context.subset_scope, build_params + ) + index_root = Path(context.dataset_path, dataset, "index") + index_path = index_root / index_label + index_config = IndexConfig( + name=index_label, + algo=algorithm, + build_param=build_params, + search_params=search_params, + file=str(index_path), + ) + backend_config = { + "name": index_label, + "algo": algorithm, + "group": group, + "index_name": index_label, + "index_root": str(index_root), + "result_scope": context.subset_scope, + "codec": build_params["codec"], + "requires_gpu": build_params["codec"] == _CAGRA_CODEC, + **context.runtime_config, + } + return BenchmarkConfig( + indexes=[index_config], + backend_config=backend_config, + ) + + @classmethod + def _group_benchmark_configs( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: List[Dict[str, Any]], + search_params: List[Dict[str, Any]], + ) -> List[BenchmarkConfig]: + benchmark_configs = [] + for params in build_params: + effective_params = _normalize_build_params(params) + benchmark_configs.append( + cls._benchmark_config( + context, + algorithm, + group, + effective_params, + search_params, + ) + ) + return benchmark_configs + + def _build_benchmark_configs( + self, + dataset_config: DatasetConfig, + dataset_conf: dict, + dataset: str, + dataset_path: str, + expanded_groups: List[Tuple[str, str, dict, List, List, dict]], + **kwargs, + ) -> List[BenchmarkConfig]: + context = _BenchmarkConfigContext( + dataset=dataset, + dataset_path=dataset_path, + subset_scope=self._subset_scope(dataset_config.subset_size), + runtime_config=self._runtime_config(kwargs), + ) + tune_mode = kwargs.get("_tune_mode", False) + tune_build_params = kwargs.get("_tune_build_params") + tune_search_params = kwargs.get("_tune_search_params") + + benchmark_configs = [] + for ( + algorithm, + group, + _group_config, + build_combinations, + search_combinations, + _group_metadata, + ) in expanded_groups: + build_params, search_params = self._effective_parameters( + build_combinations, + search_combinations, + tune_mode=tune_mode, + tune_build_params=tune_build_params, + tune_search_params=tune_search_params, + ) + benchmark_configs.extend( + self._group_benchmark_configs( + context, + algorithm, + group, + build_params, + search_params, + ) + ) + + return benchmark_configs + + +class PyLuceneBackend(BenchmarkBackend): + """Build and search cuVS-Lucene indexes through PyLucene.""" + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + self._runtime: Optional[_PyLuceneRuntime] = None + + @property + def algo(self) -> str: + return self.config.get("algo", "pylucene") + + def _result_identity(self) -> Dict[str, Any]: + identity = { + "group": self.config.get("group", "base"), + "index_name": self.config.get( + "index_name", self.config.get("name", self.algo) + ), + } + result_scope = self.config.get("result_scope") + if result_scope is not None: + identity["result_scope"] = result_scope + return identity + + def _index_metadata( + self, build_parameters: Dict[str, Any] + ) -> Dict[str, Any]: + codec_name = build_parameters["codec"] + return { + **self._result_identity(), + "compound_file_policy": _COMPOUND_FILE_POLICY[codec_name], + **build_parameters, + } + + def cleanup(self) -> None: + """Release Python references; the process-wide JVM remains active.""" + self._runtime = None + + def _get_runtime(self) -> _PyLuceneRuntime: + if self._runtime is None: + self._runtime = _PyLuceneRuntime.create(self.config) + return self._runtime + + def _trusted_index_root(self, index_path: Path) -> Path: + configured_root = self.config.get("index_root") + if configured_root is None: + # Direct backend users explicitly supply IndexConfig.file. The + # built-in loader supplies an independent root for CLI runs. + return index_path.parent + return Path(os.fspath(configured_root)) + + def _validate_index_location(self, index_path: Path) -> None: + trusted_root = self._trusted_index_root(index_path).resolve() + if index_path.resolve().parent != trusted_root: + raise ValueError( + "PyLucene index path must be an immediate child of its " + f"configured root {trusted_root}: {index_path.resolve()}" + ) + + def _remove_index(self, index_path: Path) -> None: + _safe_remove_index( + index_path, trusted_index_root=self._trusted_index_root(index_path) + ) + + def _failed_build_result( + self, + error_message: str, + index_path: str = "", + build_params: Optional[Dict[str, Any]] = None, + ) -> BuildResult: + return BuildResult( + index_path=index_path, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params or {}, + metadata=self._result_identity(), + success=False, + error_message=error_message, + ) + + def _failed_search_result( + self, + k: int, + error_message: str, + search_params: Optional[List[Dict[str, Any]]] = None, + ) -> SearchResult: + result_k = max(0, k) + return SearchResult( + neighbors=np.empty((0, result_k), dtype=np.int64), + distances=np.empty((0, result_k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=search_params or [], + metadata=self._result_identity(), + success=False, + error_message=error_message, + ) + + def _verify_existing_index( + self, + index_path: Path, + build_parameters: Dict[str, Any], + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> Tuple[_IndexProvenanceVerification, Dict[str, Any]]: + """Verify backend ownership and codec-specific persisted data.""" + codec_name = build_parameters["codec"] + if codec_name == _CAGRA_CODEC: + provenance = _verify_cagra_provenance( + index_path, + expected_build_parameters=build_parameters, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + cagra_verification = self._get_runtime().verify_cagra_index( + index_path, + expected_vector_count=provenance.vector_count, + expected_dimensions=provenance.dimensions, + ) + metadata = { + "cagra_provenance": provenance.to_metadata(), + "cagra_verification": cagra_verification.to_metadata(), + } + else: + provenance = _verify_hnsw_provenance( + index_path, + codec_name, + expected_build_parameters=build_parameters, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + metadata = {"hnsw_verification": provenance.to_metadata()} + return provenance, metadata + + def _dry_run_build_result( + self, + index_path: Path, + build_parameters: Dict[str, Any], + ) -> BuildResult: + print( + f"[dry_run] Would build PyLucene index '{index_path}' " + f"with {build_parameters}" + ) + return BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_parameters, + metadata=self._index_metadata(build_parameters), + success=True, + ) + + def _decide_existing_index( + self, + dataset: Dataset, + index_path: Path, + build_parameters: Dict[str, Any], + force: bool, + ) -> _ExistingIndexDecision: + if not index_path.exists(): + return _ExistingIndexDecision.build() + if not index_path.is_dir(): + return _ExistingIndexDecision.reject( + self._failed_build_result( + f"PyLucene index path is not a directory: {index_path}", + str(index_path), + build_parameters, + ) + ) + if not _has_lucene_segments(index_path): + return _ExistingIndexDecision.reject( + self._failed_build_result( + "Existing PyLucene index directory does not contain " + f"a Lucene segments file: {index_path}", + str(index_path), + build_parameters, + ) + ) + if force: + return _ExistingIndexDecision.build() + + metadata: Dict[str, Any] = { + **self._index_metadata(build_parameters), + "skipped": True, + } + try: + _validate_metric(dataset) + expected_shape = _expected_training_shape(dataset) + provenance, verification_metadata = self._verify_existing_index( + index_path, + build_parameters, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=( + expected_shape[1] if expected_shape else None + ), + ) + metadata["segment_count"] = provenance.segment_count + metadata.update(verification_metadata) + result = BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=_index_size(index_path), + algorithm=self.algo, + build_params=build_parameters, + metadata=metadata, + success=True, + ) + except Exception as exc: + return _ExistingIndexDecision.reject( + self._failed_build_result( + _exception_summary(exc), + str(index_path), + build_parameters, + ) + ) + + return _ExistingIndexDecision.reuse(result) + + @staticmethod + def _training_vectors_for_build( + dataset: Dataset, codec_name: str + ) -> np.ndarray: + _validate_metric(dataset) + vectors = _validate_float32_matrix( + dataset.training_vectors, "training_vectors" + ) + if vectors.shape[0] < 2: + raise ValueError( + f"{codec_name} requires at least two training vectors; " + "cuVS-Lucene does not invoke cuVS for a single-vector index" + ) + return vectors + + @staticmethod + def _resolve_build_codec( + runtime: _PyLuceneRuntime, build_parameters: Dict[str, Any] + ) -> _BuildCodec: + codec_name = build_parameters["codec"] + if codec_name == _HNSW_CODEC: + java_codec = runtime.resolve_configured_hnsw_codec( + build_parameters["m"], + build_parameters["ef_construction"], + ) + else: + java_codec = runtime.resolve_codec(codec_name) + return _BuildCodec( + codec_name=codec_name, + java_codec=java_codec, + writer_policy=_EXPECTED_WRITER_POLICY[codec_name], + build_parameters=build_parameters, + ) + + @staticmethod + def _verify_and_persist_built_index_provenance( + runtime: _PyLuceneRuntime, + index_path: Path, + build_parameters: Dict[str, Any], + vectors: np.ndarray, + topology: _IndexTopology, + ) -> Dict[str, Any]: + codec_name = build_parameters["codec"] + vector_count = int(vectors.shape[0]) + dimensions = int(vectors.shape[1]) + if codec_name == _CAGRA_CODEC: + cagra_verification = runtime.verify_cagra_index( + index_path, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + _write_cagra_provenance( + index_path, + vector_count=vector_count, + dimensions=dimensions, + build_parameters=build_parameters, + segment_count=topology.segment_count, + ) + provenance = _verify_cagra_provenance( + index_path, + expected_build_parameters=build_parameters, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return { + "cagra_provenance": provenance.to_metadata(), + "cagra_verification": cagra_verification.to_metadata(), + } + + _write_hnsw_provenance( + index_path, + codec_name, + vector_count=vector_count, + dimensions=dimensions, + build_parameters=build_parameters, + segment_count=topology.segment_count, + ) + verification = _verify_hnsw_provenance( + index_path, + codec_name, + expected_build_parameters=build_parameters, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return {"hnsw_verification": verification.to_metadata()} + + def _cleanup_partial_index( + self, index_path: Path, created_for_build: bool + ) -> Optional[Exception]: + if not created_for_build or not index_path.exists(): + return None + try: + self._remove_index(index_path) + except Exception as exc: + return exc + return None + + def build( + self, + dataset: Dataset, + indexes: List[IndexConfig], + force: bool = False, + dry_run: bool = False, + ) -> BuildResult: + """Build one local Lucene index with the selected cuVS codec.""" + if len(indexes) != 1: + return self._failed_build_result( + "PyLucene backend requires exactly one index configuration" + ) + + index_config = indexes[0] + requested_build_params = index_config.build_param + index_path = Path(index_config.file) + try: + build_parameters = _normalize_build_params( + requested_build_params, self.config + ) + self._validate_index_location(index_path) + except Exception as exc: + return self._failed_build_result( + str(exc), str(index_path), requested_build_params + ) + + if dry_run: + return self._dry_run_build_result(index_path, build_parameters) + + existing_index_decision = self._decide_existing_index( + dataset, index_path, build_parameters, force + ) + if existing_index_decision.action is _ExistingIndexAction.BUILD: + return self._build_new_index(dataset, index_path, build_parameters) + + return existing_index_decision.completed_result() + + def _build_new_index( + self, + dataset: Dataset, + index_path: Path, + build_parameters: Dict[str, Any], + ) -> BuildResult: + """Build a validated index that is not eligible for reuse.""" + + created_for_build = False + try: + codec_name = build_parameters["codec"] + vectors = self._training_vectors_for_build(dataset, codec_name) + runtime = self._get_runtime() + # Preflight must succeed before an existing index is replaced. + build_codec = self._resolve_build_codec(runtime, build_parameters) + + if index_path.exists(): + self._remove_index(index_path) + index_path.mkdir(parents=True) + created_for_build = True + + start = time.perf_counter() + topology = runtime.build_index(index_path, vectors, build_codec) + build_time = time.perf_counter() - start + + metadata = { + **self._index_metadata(build_parameters), + "pylucene_version": runtime.pylucene_version, + "writer_policy": build_codec.writer_policy, + **topology.to_metadata(), + } + metadata.update( + self._verify_and_persist_built_index_provenance( + runtime, + index_path, + build_parameters, + vectors, + topology, + ) + ) + + return BuildResult( + index_path=str(index_path), + build_time_seconds=build_time, + index_size_bytes=_index_size(index_path), + algorithm=self.algo, + build_params=build_parameters, + metadata=metadata, + success=True, + ) + except Exception as exc: + cleanup_error = self._cleanup_partial_index( + index_path, created_for_build + ) + error_message = _exception_summary(exc) + if cleanup_error is not None: + error_message += ( + "; failed to remove partial index: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + return self._failed_build_result( + error_message, + str(index_path), + build_parameters, + ) + except BaseException as exc: + cleanup_error = self._cleanup_partial_index( + index_path, created_for_build + ) + if cleanup_error is not None: + exc.add_note( + "Failed to remove partial index: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + + def _resolve_search_plans( + self, + index_config: IndexConfig, + *, + k: int, + batch_size: int, + mode: str, + search_threads: Optional[Union[int, str]], + ) -> List[_SearchPlan]: + build_parameters = _normalize_build_params( + index_config.build_param, self.config + ) + codec_name = build_parameters["codec"] + if k < 1: + raise ValueError("k must be positive") + if batch_size < 1: + raise ValueError("batch_size must be positive") + if mode != "latency": + raise ValueError( + "PyLucene backend currently supports only latency mode" + ) + if search_threads not in (None, 1, "1"): + raise ValueError( + "PyLucene backend currently supports only one search thread" + ) + search_parameters = _normalize_search_params( + index_config.search_params or [{}], codec_name=codec_name, k=k + ) + if codec_name == _CAGRA_CODEC and k > 1024: + raise ValueError( + "CuVS2510GPUSearchCodec benchmarks support k <= 1024 " + "to avoid cuVS-Lucene search paths that can use GPU " + "brute-force search above that limit" + ) + return [ + _SearchPlan( + index_path=Path(index_config.file), + codec_name=codec_name, + build_parameters=build_parameters, + search_parameters=parameters, + num_candidates=parameters.get("num_candidates", k), + k=k, + batch_size=batch_size, + mode=mode, + ) + for parameters in search_parameters + ] + + def _dry_run_search_result( + self, + plan: _SearchPlan, + ) -> SearchResult: + print( + f"[dry_run] Would search PyLucene index '{plan.index_path}' " + f"with codec={plan.codec_name}, k={plan.k}, " + f"num_candidates={plan.num_candidates}, " + f"batch_size={plan.batch_size}" + ) + return SearchResult( + neighbors=np.empty((0, plan.k), dtype=np.int64), + distances=np.empty((0, plan.k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=[plan.search_parameters], + metadata={ + **self._index_metadata(plan.build_parameters), + "top_k": plan.k, + "num_candidates": plan.num_candidates, + "batch_size": plan.batch_size, + "mode": plan.mode, + }, + success=True, + ) + + @staticmethod + def _load_search_inputs(dataset: Dataset) -> _SearchInputs: + _validate_metric(dataset) + query_vectors = _validate_float32_matrix( + dataset.query_vectors, "query_vectors" + ) + expected_shape = _expected_training_shape(dataset) + if ( + expected_shape is not None + and expected_shape[1] != query_vectors.shape[1] + ): + raise ValueError( + "Query vector dimensions do not match the dataset: " + f"{query_vectors.shape[1]} != {expected_shape[1]}" + ) + return _SearchInputs( + query_vectors=query_vectors, + expected_training_shape=expected_shape, + ) + + def _execute_search( + self, dataset: Dataset, plan: _SearchPlan + ) -> SearchResult: + end_to_end_start = time.perf_counter() + inputs = self._load_search_inputs(dataset) + query_vectors = inputs.query_vectors + expected_shape = inputs.expected_training_shape + provenance, verification_metadata = self._verify_existing_index( + plan.index_path, + plan.build_parameters, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=int(query_vectors.shape[1]), + ) + runtime = self._get_runtime() + + runtime_result = runtime.search_index( + plan.index_path, + query_vectors, + plan.k, + plan.batch_size, + plan.num_candidates, + ) + end_to_end_time_ms = (time.perf_counter() - end_to_end_start) * 1000.0 + processed = _process_search_result( + runtime_result, + provenance, + query_count=int(query_vectors.shape[0]), + k=plan.k, + batch_size=plan.batch_size, + ) + metadata = { + **self._index_metadata(plan.build_parameters), + "pylucene_version": runtime.pylucene_version, + "index_dimensions": runtime_result.index_dimensions, + "document_count": runtime_result.document_count, + "segment_count": provenance.segment_count, + "top_k": plan.k, + "num_candidates": plan.num_candidates, + "batch_size": plan.batch_size, + "num_batches": processed.num_batches, + "mode": plan.mode, + "latency_seconds": processed.latency_seconds, + "end_to_end_time_ms": end_to_end_time_ms, + "non_query_overhead_time_ms": max( + 0.0, end_to_end_time_ms - processed.search_time_ms + ), + } + metadata.update(verification_metadata) + + return SearchResult( + neighbors=processed.neighbors, + distances=processed.distances, + search_time_ms=processed.search_time_ms, + queries_per_second=processed.queries_per_second, + recall=0.0, + algorithm=self.algo, + search_params=[plan.search_parameters], + latency_percentiles=processed.latency_percentiles, + metadata=metadata, + success=True, + ) + + def search( + self, + dataset: Dataset, + indexes: List[IndexConfig], + k: int, + batch_size: int = 10000, + mode: str = "latency", + force: bool = False, + search_threads: Optional[int] = None, + dry_run: bool = False, + ) -> List[SearchResult]: + """Search one local Lucene index and report per-batch latency.""" + if len(indexes) != 1: + return [ + self._failed_search_result( + k, + "PyLucene backend requires exactly one index configuration", + ) + ] + + index_config = indexes[0] + try: + plans = self._resolve_search_plans( + index_config, + k=k, + batch_size=batch_size, + mode=mode, + search_threads=search_threads, + ) + except Exception as exc: + return [ + self._failed_search_result( + k, + str(exc), + index_config.search_params or [{}], + ) + ] + + if dry_run: + return [self._dry_run_search_result(plan) for plan in plans] + + index_path = plans[0].index_path + if not index_path.is_dir(): + return [ + self._failed_search_result( + k, + f"PyLucene index directory does not exist: {index_path}", + [plan.search_parameters for plan in plans], + ) + ] + + results = [] + for plan in plans: + try: + results.append(self._execute_search(dataset, plan)) + except Exception as exc: + results.append( + self._failed_search_result( + k, + _exception_summary(exc), + [plan.search_parameters], + ) + ) + return results + + +__all__ = ["PyLuceneBackend", "PyLuceneConfigLoader"] diff --git a/python/cuvs_bench/cuvs_bench/backends/search_spaces.py b/python/cuvs_bench/cuvs_bench/backends/search_spaces.py index edb488f77e..e601abbb39 100644 --- a/python/cuvs_bench/cuvs_bench/backends/search_spaces.py +++ b/python/cuvs_bench/cuvs_bench/backends/search_spaces.py @@ -159,6 +159,22 @@ }, }, # ========================================================================= + # PyLucene/cuVS HNSW + # ========================================================================= + "pylucene_cuvs_hnsw": { + "build": { + "m": {"type": "int", "min": 1, "max": 512}, + "ef_construction": {"type": "int", "min": 1, "max": 512}, + }, + "search": { + "num_candidates": { + "type": "int", + "min": "top_k", + "max": 500, + }, + }, + }, + # ========================================================================= # Elasticsearch GPU HNSW (hnsw, int8_hnsw, int4_hnsw, bbq_hnsw) # Per ES-GPU-API-REFERENCE.md: index_options (m, ef_construction), # knn (num_candidates) diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml new file mode 100644 index 0000000000..2e272895b7 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml @@ -0,0 +1,10 @@ +name: pylucene_cuvs_cagra +groups: + base: + build: + codec: ["CuVS2510GPUSearchCodec"] + search: {} + test: + build: + codec: ["CuVS2510GPUSearchCodec"] + search: {} diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml new file mode 100644 index 0000000000..16c5bab34b --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml @@ -0,0 +1,20 @@ +name: pylucene_cuvs_hnsw +groups: + base: + build: + codec: + - "Lucene101AcceleratedHNSWCodec" + m: + - 32 + ef_construction: + - 32 + direct_single_segment: + - false + search: {} + test: + build: + codec: ["Lucene101AcceleratedHNSWCodec"] + m: [32] + ef_construction: [32] + direct_single_segment: [false] + search: {} diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py b/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py index 7600101439..36bc5539a7 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py @@ -1,10 +1,15 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from .orchestrator import BenchmarkOrchestrator -from .config_loaders import ConfigLoader, BenchmarkConfig, DatasetConfig, CppGBenchConfigLoader +from .config_loaders import ( + ConfigLoader, + BenchmarkConfig, + DatasetConfig, + CppGBenchConfigLoader, +) from ..backends.registry import ( get_backend_class, list_backends, @@ -12,6 +17,7 @@ get_config_loader, ) from ..backends.opensearch import OpenSearchConfigLoader +from ..backends.pylucene import PyLuceneConfigLoader __all__ = [ # Main orchestrator @@ -22,6 +28,7 @@ "DatasetConfig", "CppGBenchConfigLoader", "OpenSearchConfigLoader", + "PyLuceneConfigLoader", # Registry functions "get_backend_class", "list_backends", @@ -34,10 +41,12 @@ # Register built-in config loaders # ============================================================================ + def _register_builtin_loaders(): """Register built-in config loaders.""" register_config_loader("cpp_gbench", CppGBenchConfigLoader) register_config_loader("opensearch", OpenSearchConfigLoader) + register_config_loader("pylucene", PyLuceneConfigLoader) # Auto-register when module is imported diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py index e3af2fe58e..7d6375ce96 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -362,7 +362,7 @@ def gather_algorithm_configs( list A list of paths to the algorithm configuration files. """ - algos_conf_fs = os.listdir(os.path.join(config_path, "algos")) + algos_conf_fs = sorted(os.listdir(os.path.join(config_path, "algos"))) algos_conf_fs = [ os.path.join(config_path, "algos", f) for f in algos_conf_fs @@ -373,7 +373,7 @@ def gather_algorithm_configs( if os.path.isdir(algorithm_configuration): algos_conf_fs += [ os.path.join(algorithm_configuration, f) - for f in os.listdir(algorithm_configuration) + for f in sorted(os.listdir(algorithm_configuration)) if f.endswith((".yaml", ".yml")) ] elif os.path.isfile(algorithm_configuration): diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index abffad0217..826742bef3 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -10,7 +10,7 @@ benchmark runs across different backends using the registry pattern. """ -from typing import List, Optional, Union +from typing import Any, List, Mapping, Optional, Union import numpy as np @@ -29,6 +29,17 @@ def _should_compute_recall(result: SearchResult) -> bool: return result.success and result.neighbors.size > 0 +def _resolve_tune_bound( + bound: Any, bound_values: Optional[Mapping[str, Any]] +) -> Any: + """Resolve a symbolic tune bound from the current trial context.""" + if not isinstance(bound, str): + return bound + if bound_values is None or bound not in bound_values: + raise ValueError(f"Unable to resolve symbolic tune bound {bound!r}") + return bound_values[bound] + + class BenchmarkOrchestrator: """ Orchestrator for running benchmarks using the pluggable backend system. @@ -389,18 +400,23 @@ def _run_tune( all_results: List[Union[BuildResult, SearchResult]] = [] def suggest_params( - trial, param_space: dict, build_params: dict = None + trial, + param_space: dict, + bound_values: Optional[Mapping[str, Any]] = None, ) -> dict: """Suggest parameters from search space using Optuna trial.""" params = {} for param, spec in param_space.items(): if spec["type"] == "int": - max_val = spec["max"] - # Handle dynamic constraints (e.g., nprobe <= nlist) - if isinstance(max_val, str) and build_params: - max_val = build_params.get(max_val, 1000) + min_val = _resolve_tune_bound(spec["min"], bound_values) + max_val = _resolve_tune_bound(spec["max"], bound_values) + if min_val > max_val: + raise ValueError( + f"Invalid tune range for {param!r}: minimum " + f"{min_val} exceeds maximum {max_val}" + ) params[param] = trial.suggest_int( - param, spec["min"], max_val, log=spec.get("log", False) + param, min_val, max_val, log=spec.get("log", False) ) elif spec["type"] == "float": params[param] = trial.suggest_float( @@ -421,8 +437,9 @@ def objective(trial) -> float: build_params = suggest_params(trial, search_space.get("build", {})) # Suggest search parameters (may depend on build params) + search_bound_values = {**build_params, "top_k": count} search_params_dict = suggest_params( - trial, search_space.get("search", {}), build_params + trial, search_space.get("search", {}), search_bound_values ) # Run single trial with these specific parameters diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 225a3d3f9c..4344550e76 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -14,11 +14,41 @@ from .data_export import ( convert_json_to_csv_build, convert_json_to_csv_search, + validate_dataset_name, write_results_to_csv, ) +from ..backends.base import SearchResult from ..orchestrator import BenchmarkOrchestrator +def _run_failed(results, mode, *, allow_empty=False): + if not results: + return not allow_empty + if mode == "sweep": + return any(not result.success for result in results) + return not any( + isinstance(result, SearchResult) and result.success + for result in results + ) + + +def _raise_for_failed_run(results, mode, *, allow_empty=False): + if not _run_failed(results, mode, allow_empty=allow_empty): + return + + failures = [result for result in results if not result.success] + if failures: + details = "; ".join( + f"{result.algorithm}: {result.error_message or 'benchmark failed'}" + for result in failures + ) + elif results: + details = f"{mode} mode produced no successful search result" + else: + details = f"{mode} mode produced no benchmark results" + raise click.ClickException(details) + + @click.command() @click.option( "--subset-size", @@ -257,6 +287,11 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ + try: + validate_dataset_name(dataset) + except ValueError as exc: + raise click.BadParameter(str(exc), param_hint="--dataset") from exc + if data_export: click.echo( "Warning: --data-export is deprecated because benchmark runs now " @@ -312,6 +347,10 @@ def main( ) if dry_run: + # Native dry runs print planned commands without result objects. + _raise_for_failed_run( + results, mode, allow_empty=backend_type == "cpp_gbench" + ) return if backend_type == "cpp_gbench": @@ -320,7 +359,16 @@ def main( if search: convert_json_to_csv_search(dataset, dataset_path) else: - write_results_to_csv(results, dataset, dataset_path, count, batch_size) + write_results_to_csv( + results, + dataset, + dataset_path, + count, + batch_size, + search_requested=search, + ) + + _raise_for_failed_run(results, mode) if __name__ == "__main__": diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 50afa7a5ea..393103ed5a 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -10,6 +10,11 @@ import pandas as pd +from .._validation import ( + format_artifact_identity, + validate_path_component, + validate_result_component, +) from ..backends.base import BuildResult, SearchResult skip_build_cols = set( @@ -53,8 +58,37 @@ } -def write_results_to_csv(results, dataset, dataset_path, count, batch_size): +def validate_dataset_name(dataset): + """Validate a dataset name before using it as a path component.""" + return validate_path_component(dataset, "benchmark dataset name") + + +def _result_stem(algorithm, group, scope=None): + parts = [ + validate_result_component(algorithm, "benchmark algorithm name"), + validate_result_component(group, "benchmark group name"), + ] + if scope is not None: + parts.append( + validate_result_component(scope, "benchmark result scope") + ) + return ",".join(parts) + + +def _series_name(algorithm, group, scope): + return format_artifact_identity(algorithm, group, scope) + + +def write_results_to_csv( + results, + dataset, + dataset_path, + count, + batch_size, + search_requested=False, +): """Write Python-backend results using the existing plotting CSV schema.""" + validate_dataset_name(dataset) grouped = defaultdict(list) for result in results: group = result.metadata.get("group") @@ -64,29 +98,43 @@ def write_results_to_csv(results, dataset, dataset_path, count, batch_size): ): continue method = "build" if isinstance(result, BuildResult) else "search" - grouped[(method, result.algorithm, group)].append(result) + scope = result.metadata.get("result_scope") + grouped[(method, result.algorithm, group, scope)].append(result) - for (method, algorithm, group), group_results in grouped.items(): + for (method, algorithm, group, scope), group_results in grouped.items(): if method == "build": _write_build_results( - group_results, algorithm, group, dataset, dataset_path + group_results, + algorithm, + group, + scope, + dataset, + dataset_path, ) else: _write_search_results( group_results, algorithm, group, + scope, dataset, dataset_path, count, batch_size, ) + if search_requested: + _remove_search_artifacts_after_failed_builds( + grouped, dataset, dataset_path, count, batch_size + ) + -def _write_build_results(results, algorithm, group, dataset, dataset_path): +def _write_build_results( + results, algorithm, group, scope, dataset, dataset_path +): output_dir = os.path.join(dataset_path, dataset, "result", "build") os.makedirs(output_dir, exist_ok=True) - algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + algo_name = _series_name(algorithm, group, scope) rows = [] for result in results: @@ -98,14 +146,18 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): **result.build_params, **metadata, "algo_name": algo_name, - "index_name": result.index_path, + "index_name": result.metadata.get( + "index_name", result.index_path + ), "time": result.build_time_seconds, } ) columns = ["algo_name", "index_name", "time"] dataframe = pd.DataFrame(rows) - build_file = os.path.join(output_dir, f"{algorithm},{group}.csv") + build_file = os.path.join( + output_dir, f"{_result_stem(algorithm, group, scope)}.csv" + ) complete_run = all( result.success and not result.metadata.get("skipped") @@ -131,11 +183,18 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): def _write_search_results( - results, algorithm, group, dataset, dataset_path, count, batch_size + results, + algorithm, + group, + scope, + dataset, + dataset_path, + count, + batch_size, ): output_dir = os.path.join(dataset_path, dataset, "result", "search") os.makedirs(output_dir, exist_ok=True) - algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + algo_name = _series_name(algorithm, group, scope) rows = [] for result in results: @@ -148,6 +207,7 @@ def _write_search_results( rows.append( { **search_params, + **(result.latency_percentiles or {}), **metadata, "algo_name": algo_name, "index_name": result.metadata["index_name"], @@ -179,7 +239,7 @@ def _write_search_results( dataset, "result", "build", - f"{algorithm},{group}.csv", + f"{_result_stem(algorithm, group, scope)}.csv", ) if os.path.exists(build_file): build = pd.read_csv(build_file).drop_duplicates( @@ -194,7 +254,7 @@ def _write_search_results( how="left", ) - stem = f"{algorithm},{group},k{count},bs{batch_size}" + stem = f"{_result_stem(algorithm, group, scope)},k{count},bs{batch_size}" raw_file = os.path.join(output_dir, f"{stem},raw.csv") dataframe.to_csv(raw_file, index=False) frontier_file = os.path.join(output_dir, f"{stem}.json") @@ -203,7 +263,7 @@ def _write_search_results( def _scalar_metadata(metadata): - reserved = {"group", "index_name", "latency_seconds"} + reserved = {"group", "index_name", "latency_seconds", "result_scope"} return { key: value for key, value in metadata.items() @@ -212,6 +272,26 @@ def _scalar_metadata(metadata): } +def _remove_search_artifacts_after_failed_builds( + grouped, dataset, dataset_path, count, batch_size +): + search_dir = os.path.join(dataset_path, dataset, "result", "search") + for (method, algorithm, group, scope), results in grouped.items(): + if method != "build" or any(result.success for result in results): + continue + search_key = ("search", algorithm, group, scope) + if search_key in grouped: + continue + stem = ( + f"{_result_stem(algorithm, group, scope)},k{count},bs{batch_size}" + ) + for suffix in (",raw.csv", ",throughput.csv", ",latency.csv"): + try: + os.remove(os.path.join(search_dir, f"{stem}{suffix}")) + except FileNotFoundError: + pass + + def read_json_files(dataset, dataset_path, method): """ Yield file paths, algo names, and loaded JSON data as pandas DataFrames. @@ -231,6 +311,7 @@ def read_json_files(dataset, dataset_path, method): A tuple containing the file path, algorithm name, and the DataFrame of JSON content. """ + validate_dataset_name(dataset) dir_path = os.path.join(dataset_path, dataset, "result", method) if not os.path.isdir(dir_path): return diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_execution_path_utils.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_execution_path_utils.py new file mode 100644 index 0000000000..3126a4af97 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_execution_path_utils.py @@ -0,0 +1,803 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import struct +import tempfile +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache +from pathlib import Path +from typing import Any, Iterator + +from cuvs_bench.backends.pylucene import _PyLuceneRuntime + +ID_FIELD = "id" +VECTOR_FIELD = "vector" +FILTER_FIELD = "filter-status" +FILTER_INCLUDED_VALUE = "included" +FILTER_EXCLUDED_VALUE = "excluded" + +CAGRA_TEST_CODEC_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$CagraSearchCodec" +) +CPU_HNSW_TEST_CODEC_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$CpuHnswCodec" +) +CAGRA_BUILT_HNSW_BASE_LAYER_TEST_CODEC_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$CagraBuiltHnswBaseLayerCodec" +) +CAGRA_BUILT_HNSW_THREE_LAYER_TEST_CODEC_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$CagraBuiltHnswThreeLayerCodec" +) +CAGRA_TEST_QUERY_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$CagraSearchQuery" +) +HNSW_GRAPH_VERIFYING_QUERY_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$HnswGraphVerifyingQuery" +) +HNSW_GRAPH_VERIFYING_EXACT_QUERY_CLASS = ( + "com.nvidia.cuvs.lucene.PyLuceneTestSupport$HnswGraphVerifyingExactQuery" +) +CAGRA_VECTOR_READER_CLASS = "com.nvidia.cuvs.lucene.CuVS2510GPUVectorsReader" + +QUERY_PROPERTIES = { + "field": "cuvs.lucene.pylucene.query.field", + "target": "cuvs.lucene.pylucene.query.target", + "k": "cuvs.lucene.pylucene.query.k", + "i_top_k": "cuvs.lucene.pylucene.query.iTopK", + "search_width": "cuvs.lucene.pylucene.query.searchWidth", + "expected_hnsw_m": "cuvs.lucene.pylucene.query.expectedHnswM", + "filter_field": "cuvs.lucene.pylucene.query.filterField", + "filter_value": "cuvs.lucene.pylucene.query.filterValue", +} + +_FLOAT32 = struct.Struct("!f") + + +@dataclass(frozen=True) +class IndexScenario: + name: str + codec_name: str + document_count: int + dimensions: int + top_k: int + segment_count: int = 1 + force_merge_segment_count: int = 0 + disable_automatic_merges: bool = True + document_ids_without_vectors: frozenset[int] = frozenset() + document_ids_to_delete: frozenset[int] = frozenset() + additional_query_document_ids: tuple[int, ...] = () + document_ids_accepted_by_filter: frozenset[int] = frozenset() + filter_query_document_id: int | None = None + codec_factory_class: str = "" + use_cagra_search_query: bool = False + use_exact_hnsw_filter_query: bool = False + expected_hnsw_m: int = 0 + search_width: int = 1 + i_top_k: int = 64 + + +class QueryDocumentState(Enum): + SEARCHABLE = "searchable" + WITHOUT_VECTOR = "without-vector" + DELETED = "deleted" + + +@dataclass(frozen=True) +class QueryObservation: + query_document_id: int + hit_ids: tuple[str, ...] + query_class: str + query_document_state: QueryDocumentState + + +@dataclass(frozen=True) +class FilteredQueryObservation: + query_document_id: int + hit_ids: tuple[str, ...] + query_class: str + + +@dataclass(frozen=True) +class VectorReaderMetadata: + vector_count: int + dimensions: tuple[int, ...] + reader_classes: tuple[str, ...] + hnsw_layer_counts: tuple[int, ...] + + +@dataclass(frozen=True) +class IndexRun: + index_files: tuple[str, ...] + pre_merge_segment_count: int + segment_count: int + live_document_count: int + max_document_count: int + vector_count: int + vector_dimensions: tuple[int, ...] + vector_reader_classes: tuple[str, ...] + hnsw_layer_counts: tuple[int, ...] + searchable_vector_document_ids: tuple[int, ...] + document_ids_without_vectors: frozenset[int] + deleted_document_ids: frozenset[int] + query_observations: tuple[QueryObservation, ...] + filtered_query_observation: FilteredQueryObservation | None + writer_telemetry: dict[str, str] + + +@dataclass +class PyLuceneContext: + runtime: _PyLuceneRuntime + codec_class: Any + query_class: Any + class_class: Any + system_class: Any + jarray: Any + codec_cache: dict[str, Any] + + +@lru_cache(maxsize=None) +def deterministic_float32_vector( + document_id: int, dimensions: int +) -> tuple[float, ...]: + value = ((document_id + 1) * 2654435761) & 0xFFFFFFFF + vector = [] + for dimension in range(dimensions): + value = (1664525 * value + 1013904223 + dimension * 17) & 0xFFFFFFFF + component = (value / 4294967295.0) * 2.0 - 1.0 + vector.append(_FLOAT32.unpack(_FLOAT32.pack(component))[0]) + return tuple(vector) + + +def to_java_float_array(jarray: Any, values: tuple[float, ...]) -> Any: + return jarray("float")(values) + + +def initialize_pylucene_context( + backend_config: dict[str, str], + expected_spi_codecs: tuple[str, ...], +) -> PyLuceneContext: + from org.apache.lucene.search import Query + + runtime = _PyLuceneRuntime.create(backend_config) + runtime.attach_current_thread() + class_class = runtime.Class + codec_class = runtime.Codec + + # Resolve classpath failures before constructing codecs or indexes. + class_class.forName(CPU_HNSW_TEST_CODEC_CLASS) + class_class.forName(CAGRA_TEST_CODEC_CLASS) + class_class.forName(CAGRA_BUILT_HNSW_BASE_LAYER_TEST_CODEC_CLASS) + class_class.forName(CAGRA_BUILT_HNSW_THREE_LAYER_TEST_CODEC_CLASS) + class_class.forName(CAGRA_TEST_QUERY_CLASS) + class_class.forName(HNSW_GRAPH_VERIFYING_QUERY_CLASS) + class_class.forName(HNSW_GRAPH_VERIFYING_EXACT_QUERY_CLASS) + + available_codecs = codec_class.availableCodecs() + for codec_name in expected_spi_codecs: + if not available_codecs.contains(codec_name): + raise RuntimeError( + f"{codec_name} was not advertised by Lucene SPI; " + f"available codecs: {available_codecs}" + ) + + return PyLuceneContext( + runtime=runtime, + codec_class=codec_class, + query_class=Query, + class_class=class_class, + system_class=runtime.System, + jarray=runtime.lucene.JArray, + codec_cache={}, + ) + + +def _codec_for_scenario( + scenario: IndexScenario, context: PyLuceneContext +) -> Any: + cache_key = scenario.codec_factory_class or f"spi:{scenario.codec_name}" + cached = context.codec_cache.get(cache_key) + if cached is not None: + return cached + + if scenario.codec_factory_class: + reflected = context.class_class.forName( + scenario.codec_factory_class + ).newInstance() + codec = context.codec_class.cast_(reflected) + else: + available_codecs = context.codec_class.availableCodecs() + if not available_codecs.contains(scenario.codec_name): + raise RuntimeError( + f"{scenario.codec_name} was not advertised by Lucene SPI; " + f"available codecs: {available_codecs}" + ) + codec = context.codec_class.forName(scenario.codec_name) + + if codec.getName() != scenario.codec_name: + raise RuntimeError( + f"Expected codec {scenario.codec_name}, got {codec.getName()}" + ) + context.codec_cache[cache_key] = codec + return codec + + +def _writer_telemetry(codec: Any) -> dict[str, str]: + description = str(codec.knnVectorsFormat()) + _, separator, payload = description.partition("(") + if not separator or not payload.endswith(")"): + raise RuntimeError( + f"Malformed vector-format diagnostics: {description!r}" + ) + + telemetry: dict[str, str] = {} + for item in payload[:-1].split(";"): + key, item_separator, value = item.partition("=") + if not item_separator or not key: + raise RuntimeError( + f"Malformed vector-format diagnostic item {item!r} in {description!r}" + ) + telemetry[key] = value + return telemetry + + +@contextmanager +def _temporary_system_properties( + system_class: Any, properties: dict[str, str] +) -> Iterator[None]: + previous = {name: system_class.getProperty(name) for name in properties} + try: + for name, value in properties.items(): + system_class.setProperty(name, value) + yield + finally: + for name, value in previous.items(): + if value is None: + system_class.clearProperty(name) + else: + system_class.setProperty(name, value) + + +def _new_cagra_search_query( + context: PyLuceneContext, + target: tuple[float, ...], + top_k: int, + i_top_k: int, + search_width: int, + filter_field: str = "", + filter_value: str = "", +) -> Any: + properties = { + QUERY_PROPERTIES["field"]: VECTOR_FIELD, + QUERY_PROPERTIES["target"]: ",".join( + f"{value:.9g}" for value in target + ), + QUERY_PROPERTIES["k"]: str(top_k), + QUERY_PROPERTIES["i_top_k"]: str(max(top_k, i_top_k)), + QUERY_PROPERTIES["search_width"]: str(search_width), + } + if filter_field: + properties[QUERY_PROPERTIES["filter_field"]] = filter_field + properties[QUERY_PROPERTIES["filter_value"]] = filter_value + with _temporary_system_properties(context.system_class, properties): + reflected = context.class_class.forName( + CAGRA_TEST_QUERY_CLASS + ).newInstance() + return context.query_class.cast_(reflected) + + +def _new_hnsw_graph_verifying_query( + context: PyLuceneContext, + target: tuple[float, ...], + top_k: int, + expected_m: int, + filter_field: str = "", + filter_value: str = "", + *, + exact_search: bool = False, +) -> Any: + properties = { + QUERY_PROPERTIES["field"]: VECTOR_FIELD, + QUERY_PROPERTIES["target"]: ",".join( + f"{value:.9g}" for value in target + ), + QUERY_PROPERTIES["k"]: str(top_k), + QUERY_PROPERTIES["expected_hnsw_m"]: str(expected_m), + } + if filter_field: + properties[QUERY_PROPERTIES["filter_field"]] = filter_field + properties[QUERY_PROPERTIES["filter_value"]] = filter_value + query_class = ( + HNSW_GRAPH_VERIFYING_EXACT_QUERY_CLASS + if exact_search + else HNSW_GRAPH_VERIFYING_QUERY_CLASS + ) + with _temporary_system_properties(context.system_class, properties): + reflected = context.class_class.forName(query_class).newInstance() + return context.query_class.cast_(reflected) + + +def _validate_document_ids(scenario: IndexScenario) -> None: + valid_document_ids = set(range(scenario.document_count)) + referenced_document_ids = ( + set(scenario.document_ids_without_vectors) + | set(scenario.document_ids_to_delete) + | set(scenario.additional_query_document_ids) + | set(scenario.document_ids_accepted_by_filter) + ) + if scenario.filter_query_document_id is not None: + referenced_document_ids.add(scenario.filter_query_document_id) + invalid_document_ids = referenced_document_ids - valid_document_ids + if invalid_document_ids: + raise ValueError( + f"{scenario.name}: document IDs are outside the index: " + f"{sorted(invalid_document_ids)}" + ) + + conflicting_document_ids = ( + scenario.document_ids_without_vectors & scenario.document_ids_to_delete + ) + if conflicting_document_ids: + raise ValueError( + f"{scenario.name}: documents cannot be both vectorless and deleted: " + f"{sorted(conflicting_document_ids)}" + ) + + if scenario.filter_query_document_id is None: + if scenario.document_ids_accepted_by_filter: + raise ValueError( + f"{scenario.name}: accepted filter documents require a " + "filtered query" + ) + return + + if not scenario.document_ids_accepted_by_filter: + raise ValueError( + f"{scenario.name}: filtered query has no accepted documents" + ) + if scenario.filter_query_document_id in ( + scenario.document_ids_without_vectors | scenario.document_ids_to_delete + ): + raise ValueError( + f"{scenario.name}: filtered query document must have a live " + f"vector: {scenario.filter_query_document_id}" + ) + if ( + scenario.filter_query_document_id + in scenario.document_ids_accepted_by_filter + ): + raise ValueError( + f"{scenario.name}: filtered query document must be rejected by " + "the filter" + ) + + +def _representative_query_document_ids( + scenario: IndexScenario, searchable_vector_document_ids: tuple[int, ...] +) -> tuple[int, ...]: + targets = (0, scenario.document_count // 2, scenario.document_count - 1) + query_document_ids: list[int] = [] + for target in targets: + nearest = min( + searchable_vector_document_ids, + key=lambda document_id: abs(document_id - target), + ) + if nearest not in query_document_ids: + query_document_ids.append(nearest) + return tuple(query_document_ids) + + +def _query_document_state( + query_document_id: int, + document_ids_without_vectors: frozenset[int], + deleted_document_ids: frozenset[int], +) -> QueryDocumentState: + if query_document_id in document_ids_without_vectors: + return QueryDocumentState.WITHOUT_VECTOR + if query_document_id in deleted_document_ids: + return QueryDocumentState.DELETED + return QueryDocumentState.SEARCHABLE + + +def segment_document_id_ranges( + document_count: int, segment_count: int +) -> tuple[range, ...]: + effective_segment_count = max(1, min(segment_count, document_count)) + documents_per_segment, remainder = divmod( + document_count, effective_segment_count + ) + ranges = [] + start_document_id = 0 + for segment_id in range(effective_segment_count): + documents_in_segment = documents_per_segment + ( + 1 if segment_id < remainder else 0 + ) + end_document_id = start_document_id + documents_in_segment + ranges.append(range(start_document_id, end_document_id)) + start_document_id = end_document_id + return tuple(ranges) + + +def _segment_end_document_ids(scenario: IndexScenario) -> set[int]: + return { + document_ids.stop - 1 + for document_ids in segment_document_id_ranges( + scenario.document_count, scenario.segment_count + ) + } + + +def _segment_count(directory: Any) -> int: + from org.apache.lucene.index import DirectoryReader + + reader = DirectoryReader.open(directory) + try: + return sum(1 for _ in reader.leaves()) + finally: + reader.close() + + +def _new_writer_config(codec: Any, suppress_merges: bool) -> Any: + from org.apache.lucene.index import ( + IndexWriterConfig, + LogDocMergePolicy, + NoMergePolicy, + SerialMergeScheduler, + ) + + config = IndexWriterConfig() + config.setCodec(codec) + config.setUseCompoundFile(False) + config.setMergeScheduler(SerialMergeScheduler()) + if suppress_merges: + config.setMergePolicy(NoMergePolicy.INSTANCE) + else: + merge_policy = LogDocMergePolicy() + merge_policy.setMergeFactor(1000) + config.setMergePolicy(merge_policy) + return config + + +def _hit_ids(stored_fields: Any, hits: Any) -> tuple[str, ...]: + return tuple(stored_fields.document(hit.doc).get(ID_FIELD) for hit in hits) + + +def _vector_reader_observations( + reader: Any, +) -> VectorReaderMetadata: + from org.apache.lucene.codecs.hnsw import HnswGraphProvider + from org.apache.lucene.index import SegmentReader + + vector_count = 0 + dimensions: list[int] = [] + reader_classes: list[str] = [] + hnsw_layer_counts: list[int] = [] + + for leaf_reader_context in reader.leaves(): + leaf_reader = leaf_reader_context.reader() + values = leaf_reader.getFloatVectorValues(VECTOR_FIELD) + if values is not None: + vector_count += values.size() + dimensions.append(values.dimension()) + + if not SegmentReader.instance_(leaf_reader): + reader_classes.append(str(leaf_reader.getClass().getName())) + continue + + segment_reader = SegmentReader.cast_(leaf_reader) + vector_reader = segment_reader.getVectorReader() + reader_classes.append(str(vector_reader.getClass().getName())) + if HnswGraphProvider.instance_(vector_reader): + graph_provider = HnswGraphProvider.cast_(vector_reader) + hnsw_layer_counts.append( + graph_provider.getGraph(VECTOR_FIELD).numLevels() + ) + + return VectorReaderMetadata( + vector_count=vector_count, + dimensions=tuple(dimensions), + reader_classes=tuple(reader_classes), + hnsw_layer_counts=tuple(hnsw_layer_counts), + ) + + +def _write_index_and_apply_deletions( + directory: Any, + codec: Any, + scenario: IndexScenario, + context: PyLuceneContext, + document_ids_without_vectors: frozenset[int], + document_ids_to_delete: frozenset[int], +) -> None: + from org.apache.lucene.document import ( + Document, + Field, + KnnFloatVectorField, + StringField, + ) + from org.apache.lucene.index import ( + IndexWriter, + Term, + VectorSimilarityFunction, + ) + + writer_config = _new_writer_config( + codec, suppress_merges=scenario.disable_automatic_merges + ) + writer = IndexWriter(directory, writer_config) + try: + segment_end_document_ids = _segment_end_document_ids(scenario) + for document_id in range(scenario.document_count): + document = Document() + document.add( + StringField(ID_FIELD, f"doc-{document_id}", Field.Store.YES) + ) + if document_id not in document_ids_without_vectors: + vector = deterministic_float32_vector( + document_id, scenario.dimensions + ) + document.add( + KnnFloatVectorField( + VECTOR_FIELD, + to_java_float_array(context.jarray, vector), + VectorSimilarityFunction.EUCLIDEAN, + ) + ) + if scenario.filter_query_document_id is not None: + filter_value = ( + FILTER_INCLUDED_VALUE + if document_id in scenario.document_ids_accepted_by_filter + else FILTER_EXCLUDED_VALUE + ) + document.add( + StringField( + FILTER_FIELD, + filter_value, + Field.Store.NO, + ) + ) + writer.addDocument(document) + if document_id in segment_end_document_ids: + writer.commit() + + for document_id in document_ids_to_delete: + writer.deleteDocuments(Term(ID_FIELD, f"doc-{document_id}")) + if document_ids_to_delete: + writer.commit() + finally: + writer.close() + + +def _force_merge_if_requested( + directory: Any, codec: Any, force_merge_segment_count: int +) -> None: + if not force_merge_segment_count: + return + + from org.apache.lucene.index import IndexWriter + + writer = IndexWriter( + directory, _new_writer_config(codec, suppress_merges=False) + ) + try: + writer.forceMerge(force_merge_segment_count) + writer.commit() + finally: + writer.close() + + +def _new_vector_query( + scenario: IndexScenario, + context: PyLuceneContext, + query_document_id: int, + top_k: int, + filter_field: str = "", + filter_value: str = "", +) -> Any: + from org.apache.lucene.index import Term + from org.apache.lucene.search import KnnFloatVectorQuery, TermQuery + + target = deterministic_float32_vector( + query_document_id, scenario.dimensions + ) + if scenario.use_cagra_search_query: + return _new_cagra_search_query( + context, + target, + top_k, + scenario.i_top_k, + scenario.search_width, + filter_field, + filter_value, + ) + if scenario.expected_hnsw_m: + return _new_hnsw_graph_verifying_query( + context, + target, + top_k, + scenario.expected_hnsw_m, + filter_field, + filter_value, + exact_search=( + scenario.use_exact_hnsw_filter_query and bool(filter_field) + ), + ) + filter_query = ( + TermQuery(Term(filter_field, filter_value)) if filter_field else None + ) + if filter_query is None: + return KnnFloatVectorQuery( + VECTOR_FIELD, + to_java_float_array(context.jarray, target), + top_k, + ) + return KnnFloatVectorQuery( + VECTOR_FIELD, + to_java_float_array(context.jarray, target), + top_k, + filter_query, + ) + + +def _run_filtered_query( + searcher: Any, + stored_fields: Any, + context: PyLuceneContext, + scenario: IndexScenario, + query_document_id: int, + accepted_document_count: int, +) -> FilteredQueryObservation: + top_k = min(scenario.top_k, accepted_document_count) + query = _new_vector_query( + scenario, + context, + query_document_id, + top_k, + FILTER_FIELD, + FILTER_INCLUDED_VALUE, + ) + hits = searcher.search(query, top_k).scoreDocs + return FilteredQueryObservation( + query_document_id=query_document_id, + hit_ids=_hit_ids(stored_fields, hits), + query_class=str(query.getClass().getName()), + ) + + +def run_index_scenario( + scenario: IndexScenario, context: PyLuceneContext +) -> IndexRun: + context.runtime.attach_current_thread() + + from java.nio.file import Paths + from org.apache.lucene.index import DirectoryReader + from org.apache.lucene.search import IndexSearcher + from org.apache.lucene.store import FSDirectory + + codec = _codec_for_scenario(scenario, context) + _validate_document_ids(scenario) + document_ids_without_vectors = scenario.document_ids_without_vectors + document_ids_to_delete = scenario.document_ids_to_delete + searchable_vector_document_ids = tuple( + document_id + for document_id in range(scenario.document_count) + if document_id not in document_ids_without_vectors + and document_id not in document_ids_to_delete + ) + if not searchable_vector_document_ids: + raise RuntimeError( + f"{scenario.name}: no searchable vectors are available" + ) + filter_accepted_searchable_document_ids = tuple( + document_id + for document_id in searchable_vector_document_ids + if document_id in scenario.document_ids_accepted_by_filter + ) + if ( + scenario.filter_query_document_id is not None + and not filter_accepted_searchable_document_ids + ): + raise RuntimeError( + f"{scenario.name}: filter accepts no searchable vectors" + ) + representative_query_document_ids = _representative_query_document_ids( + scenario, searchable_vector_document_ids + ) + + with ExitStack() as stack: + index_path = stack.enter_context( + tempfile.TemporaryDirectory( + prefix=f"cuvs-lucene-pylucene-{scenario.name}-" + ) + ) + directory = FSDirectory.open(Paths.get(index_path)) + try: + _write_index_and_apply_deletions( + directory, + codec, + scenario, + context, + document_ids_without_vectors, + document_ids_to_delete, + ) + + pre_merge_segment_count = _segment_count(directory) + _force_merge_if_requested( + directory, codec, scenario.force_merge_segment_count + ) + + index_files = tuple( + sorted(path.name for path in Path(index_path).iterdir()) + ) + reader = DirectoryReader.open(directory) + try: + vector_reader_metadata = _vector_reader_observations(reader) + searcher = IndexSearcher(reader) + stored_fields = searcher.storedFields() + top_k = min( + scenario.top_k, len(searchable_vector_document_ids) + ) + observations = [] + query_document_ids = tuple( + dict.fromkeys( + representative_query_document_ids + + scenario.additional_query_document_ids + ) + ) + for query_document_id in query_document_ids: + query = _new_vector_query( + scenario, context, query_document_id, top_k + ) + hits = searcher.search(query, top_k).scoreDocs + observations.append( + QueryObservation( + query_document_id=query_document_id, + hit_ids=_hit_ids(stored_fields, hits), + query_class=str(query.getClass().getName()), + query_document_state=_query_document_state( + query_document_id, + document_ids_without_vectors, + document_ids_to_delete, + ), + ) + ) + + filtered_query_observation = ( + _run_filtered_query( + searcher, + stored_fields, + context, + scenario, + scenario.filter_query_document_id, + len(filter_accepted_searchable_document_ids), + ) + if scenario.filter_query_document_id is not None + else None + ) + + result = IndexRun( + index_files=index_files, + pre_merge_segment_count=pre_merge_segment_count, + segment_count=sum(1 for _ in reader.leaves()), + live_document_count=reader.numDocs(), + max_document_count=reader.maxDoc(), + vector_count=vector_reader_metadata.vector_count, + vector_dimensions=vector_reader_metadata.dimensions, + vector_reader_classes=vector_reader_metadata.reader_classes, + hnsw_layer_counts=vector_reader_metadata.hnsw_layer_counts, + searchable_vector_document_ids=searchable_vector_document_ids, + document_ids_without_vectors=document_ids_without_vectors, + deleted_document_ids=document_ids_to_delete, + query_observations=tuple(observations), + filtered_query_observation=filtered_query_observation, + writer_telemetry=_writer_telemetry(codec), + ) + finally: + reader.close() + finally: + directory.close() + + return result diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_live_test_config.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_live_test_config.py new file mode 100644 index 0000000000..488501ecd4 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_live_test_config.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared configuration for opt-in PyLucene integration tests.""" + +from __future__ import annotations + +import importlib +import os +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType + +import pytest + +OPT_IN_ENV = "CUVS_BENCH_PYLUCENE_INTEGRATION" +CUVS_JAVA_JAR_ENV = "CUVS_LUCENE_CUVS_JAVA_JAR" +CUVS_LUCENE_JAR_ENV = "CUVS_LUCENE_JAR" +PYLUCENE_TEST_CLASSES_ENV = "PYLUCENE_TEST_CLASSES" + +PYTHON_PACKAGE_ROOT = Path(__file__).resolve().parents[3] +JAVA_TEST_SOURCE_ROOT = PYTHON_PACKAGE_ROOT / "tests" / "java" + +_EXPECTED_NESTED_CLASSES = { + "PyLuceneTestSupport.java": ( + "CagraBuiltHnswBaseLayerCodec", + "CagraBuiltHnswThreeLayerCodec", + "CagraSearchCodec", + "CagraSearchQuery", + "CpuHnswCodec", + "DiagnosticKnnVectorsFormat", + "HnswGraphVerifyingExactQuery", + "HnswGraphVerifyingQuery", + "QueryProperties", + ), + "PyLuceneWriterSelectionCodec.java": ("WriterSelectionFormat",), +} + + +@dataclass(frozen=True) +class PyLuceneRuntimeFixture: + """Resolved runtime inputs shared by every live PyLucene test.""" + + backend_config: dict[str, str] + test_classes: Path + + +def _required_jar(env_name: str) -> Path: + configured = os.environ.get(env_name) + if not configured: + pytest.fail(f"{env_name} must point to the required runtime jar") + + jar = Path(configured).expanduser() + if not jar.is_file(): + pytest.fail(f"{env_name} does not point to an existing file: {jar}") + return jar.resolve() + + +def _find_javac() -> str: + javac = shutil.which("javac") + environment_javac = Path(sys.prefix) / "lib" / "jvm" / "bin" / "javac" + if javac is None and environment_javac.is_file(): + javac = str(environment_javac) + if javac is None: + pytest.fail("javac is required for the PyLucene integration tests") + return javac + + +def _expected_class_files(java_sources: tuple[Path, ...]) -> tuple[Path, ...]: + expected = [] + for source in java_sources: + relative_class = source.relative_to(JAVA_TEST_SOURCE_ROOT).with_suffix( + ".class" + ) + expected.append(relative_class) + for nested_class in _EXPECTED_NESTED_CLASSES.get(source.name, ()): + expected.append( + relative_class.with_name( + f"{relative_class.stem}${nested_class}.class" + ) + ) + return tuple(expected) + + +def compile_test_java_sources( + output_dir: Path, + cuvs_java_jar: Path, + cuvs_lucene_jar: Path, + pylucene_classpath: str, +) -> None: + """Compile every PyLucene test adapter together before JVM startup.""" + java_sources = tuple(sorted(JAVA_TEST_SOURCE_ROOT.rglob("*.java"))) + if not java_sources: + pytest.fail( + "No PyLucene Java test sources were found under " + f"{JAVA_TEST_SOURCE_ROOT}" + ) + + output_dir.mkdir(parents=True, exist_ok=True) + compile_classpath = os.pathsep.join( + (str(cuvs_java_jar), str(cuvs_lucene_jar), pylucene_classpath) + ) + completed = subprocess.run( + [ + _find_javac(), + "--release", + "22", + "-classpath", + compile_classpath, + "-d", + str(output_dir), + *(str(source) for source in java_sources), + ], + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + pytest.fail( + "Could not compile the PyLucene Java test adapters:\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + + missing_classes = tuple( + relative + for relative in _expected_class_files(java_sources) + if not (output_dir / relative).is_file() + ) + if missing_classes: + missing = ", ".join(str(path) for path in missing_classes) + pytest.fail( + "Compiled PyLucene test adapters are incomplete under " + f"{output_dir}: missing {missing}" + ) + + +def _load_uninitialized_pylucene() -> ModuleType: + try: + lucene = importlib.import_module("lucene") + except ImportError as exc: + pytest.fail(f"PyLucene is not importable: {exc}") + + from cuvs_bench.backends.pylucene import _validate_pylucene_version + + _validate_pylucene_version(lucene) + if lucene.getVMEnv() is not None: + pytest.fail( + "PyLucene's process-wide JVM was initialized before the shared " + "integration-test bootstrap; run the suite in a fresh process" + ) + return lucene + + +def configure_pylucene_runtime( + tmp_path_factory: pytest.TempPathFactory, +) -> PyLuceneRuntimeFixture: + """Resolve inputs and prepare the sole live-test JVM classpath.""" + if os.environ.get(OPT_IN_ENV) != "1": + pytest.skip(f"set {OPT_IN_ENV}=1 to run PyLucene integration tests") + + cuvs_java_jar = _required_jar(CUVS_JAVA_JAR_ENV) + cuvs_lucene_jar = _required_jar(CUVS_LUCENE_JAR_ENV) + java_library_path = os.environ.get("JAVA_LIBRARY_PATH") or os.environ.get( + "LD_LIBRARY_PATH" + ) + if not java_library_path: + pytest.fail( + "JAVA_LIBRARY_PATH or LD_LIBRARY_PATH must provide the native " + "cuVS runtime libraries" + ) + + lucene = _load_uninitialized_pylucene() + test_classes = tmp_path_factory.mktemp("pylucene-java-test-classes") + compile_test_java_sources( + test_classes, + cuvs_java_jar, + cuvs_lucene_jar, + str(lucene.CLASSPATH), + ) + classpath_entries = str(lucene.CLASSPATH).split(os.pathsep) + if str(test_classes) in classpath_entries: + pytest.fail( + "The temporary PyLucene test classes directory was already on " + "the process classpath before bootstrap" + ) + lucene.CLASSPATH = os.pathsep.join( + (str(test_classes), str(lucene.CLASSPATH)) + ) + + return PyLuceneRuntimeFixture( + backend_config={ + "cuvs_java_jar": str(cuvs_java_jar), + "cuvs_lucene_jar": str(cuvs_lucene_jar), + "java_library_path": java_library_path, + }, + test_classes=test_classes, + ) diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_test_utils.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_test_utils.py new file mode 100644 index 0000000000..0b9cf6dc90 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/_pylucene_test_utils.py @@ -0,0 +1,244 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Shared fakes and data builders for PyLucene unit tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.pylucene import ( + PyLuceneBackend, + _CagraIndexVerification, + _IndexTopology, + _RuntimeSearchResult, + _SearchHit, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig + +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" + + +class _FakeRuntime: + pylucene_version = "test" + + def __init__( + self, + *, + index_dimensions: int = 4, + document_count: int = 10, + hits: list[list[_SearchHit]] | None = None, + batch_latencies_ms: list[float] | None = None, + validate_query_dimensions: bool = True, + ): + self.index_dimensions = index_dimensions + self.document_count = document_count + self.hits = hits + self.batch_latencies_ms = batch_latencies_ms + self.validate_query_dimensions = validate_query_dimensions + self.resolve_calls = [] + self.configured_resolve_calls = [] + self.build_calls = [] + self.search_calls = [] + self.verification_calls = [] + self.resolve_error = None + self.build_error = None + self.topology = None + self.search_error = None + self.verification_error = None + + def resolve_codec(self, codec_name): + self.resolve_calls.append(codec_name) + if self.resolve_error is not None: + raise self.resolve_error + return codec_name + + def resolve_configured_hnsw_codec(self, m, ef_construction): + self.configured_resolve_calls.append((m, ef_construction)) + if self.resolve_error is not None: + raise self.resolve_error + return (_HNSW_CODEC, m, ef_construction) + + def build_index(self, index_path, vectors, build_codec): + self.build_calls.append((index_path, vectors.copy(), build_codec)) + if self.build_error is not None: + (index_path / "partial").write_bytes(b"partial") + raise self.build_error + self.document_count = int(vectors.shape[0]) + (index_path / "segments_1").write_bytes(b"index") + topology = self.topology or _IndexTopology( + segment_document_counts=(self.document_count,), + segment_vector_counts=(self.document_count,), + ) + if build_codec.direct_single_segment: + topology.require_direct_single_segment(self.document_count) + else: + topology.validate(self.document_count) + return topology + + def verify_cagra_index( + self, + index_path, + expected_vector_count=None, + expected_dimensions=None, + ): + self.verification_calls.append( + (index_path, expected_vector_count, expected_dimensions) + ) + if self.verification_error is not None: + raise self.verification_error + return _CagraIndexVerification( + segment_count=1, + field_count=1, + vector_count=( + expected_vector_count + if expected_vector_count is not None + else 10 + ), + dimensions=( + expected_dimensions if expected_dimensions is not None else 4 + ), + ) + + def search_index( + self, index_path, query_vectors, k, batch_size, num_candidates=None + ): + self.search_calls.append( + ( + index_path, + query_vectors.copy(), + k, + batch_size, + num_candidates, + ) + ) + if self.search_error is not None: + raise self.search_error + if ( + self.validate_query_dimensions + and query_vectors.shape[1] != self.index_dimensions + ): + raise ValueError( + "Query vector dimensions do not match the Lucene index" + ) + hits = self.hits + if hits is None: + hits = [ + [_SearchHit(document_id=0, score=1.0)] + for _ in range(query_vectors.shape[0]) + ] + batch_latencies_ms = self.batch_latencies_ms + if batch_latencies_ms is None: + batch_latencies_ms = [ + 1.0 for _ in range(0, query_vectors.shape[0], batch_size) + ] + return _RuntimeSearchResult( + hits=hits, + batch_latencies_ms=batch_latencies_ms, + index_dimensions=self.index_dimensions, + document_count=self.document_count, + ) + + +def _dataset( + *, + n_base: int = 10, + n_queries: int = 2, + dimensions: int = 4, + dtype=np.float32, + distance_metric: str = "euclidean", +) -> Dataset: + rng = np.random.default_rng(7) + base = rng.random((n_base, dimensions)).astype(dtype) + queries = rng.random((n_queries, dimensions)).astype(dtype) + return Dataset( + name="test", + training_vectors=base, + query_vectors=queries, + distance_metric=distance_metric, + ) + + +def _index( + index_path: Path, + *, + codec: str = _HNSW_CODEC, + build_params: dict | None = None, + search_params: list[dict] | None = None, +) -> IndexConfig: + parameters = {"codec": codec} + if build_params: + parameters.update(build_params) + return IndexConfig( + name="pylucene-test", + algo="pylucene_cuvs_hnsw", + build_param=parameters, + search_params=[{}] if search_params is None else search_params, + file=str(index_path), + ) + + +def _backend( + runtime: _FakeRuntime | None = None, + *, + codec: str = _HNSW_CODEC, +) -> PyLuceneBackend: + backend = PyLuceneBackend( + { + "name": "pylucene-test", + "algo": "pylucene_cuvs_hnsw", + "codec": codec, + } + ) + if runtime is not None: + backend._runtime = runtime + return backend + + +def _prepare_hnsw_index( + index_path: Path, + *, + codec: str = _HNSW_CODEC, + vector_count: int = 10, + dimensions: int = 4, +) -> Path: + index_path.mkdir(parents=True, exist_ok=True) + (index_path / "segments_1").write_bytes(b"index") + pylucene_backend._write_hnsw_provenance( + index_path, + codec, + vector_count=vector_count, + dimensions=dimensions, + ) + return index_path / pylucene_backend._HNSW_PROVENANCE_FILE + + +def _prepare_cagra_index( + index_path: Path, + *, + vector_count: int = 10, + dimensions: int = 4, +) -> Path: + index_path.mkdir(parents=True, exist_ok=True) + (index_path / "segments_1").write_bytes(b"index") + pylucene_backend._write_cagra_provenance( + index_path, + vector_count=vector_count, + dimensions=dimensions, + ) + return index_path / pylucene_backend._CAGRA_PROVENANCE_FILE + + +def _write_test_bin(path: Path, data: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as file: + write_bin_header(file, data.shape[0], data.shape[1]) + np.ascontiguousarray(data).tofile(file) diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/conftest.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/conftest.py new file mode 100644 index 0000000000..49e1848d2f --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/conftest.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Session-scoped bootstrap for live PyLucene tests.""" + +from __future__ import annotations + +import pytest + +from cuvs_bench.tests.pylucene._pylucene_live_test_config import ( + PyLuceneRuntimeFixture, + configure_pylucene_runtime, +) + + +@pytest.fixture(scope="session") +def pylucene_runtime_config( + tmp_path_factory: pytest.TempPathFactory, +) -> PyLuceneRuntimeFixture: + """Prepare the process-global JVM classpath before any live test runs.""" + return configure_pylucene_runtime(tmp_path_factory) diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/pylucene_cpu_fallback_probe.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/pylucene_cpu_fallback_probe.py new file mode 100644 index 0000000000..f9f71a9cac --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/pylucene_cpu_fallback_probe.py @@ -0,0 +1,114 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Fresh-process probe for the cuVS-Lucene HNSW CPU fallback.""" + +import os +import re +import tempfile +from pathlib import Path + +import lucene +import numpy as np + +from cuvs_bench.backends._pylucene_java import ( + EF_CONSTRUCTION_PROPERTY, + M_PROPERTY, +) +from cuvs_bench.backends.pylucene import ( + _BuildCodec, + _PyLuceneRuntime, + _validate_pylucene_version, +) +from cuvs_bench.tests.pylucene._pylucene_live_test_config import ( + PYLUCENE_TEST_CLASSES_ENV, +) + +_CPU_HNSW_WRITER = ( + "org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter" +) +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_WRITER_SELECTION_CODEC = "com.nvidia.cuvs.bench.PyLuceneWriterSelectionCodec" +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} + + +def _writer_diagnostics(java_codec): + diagnostics = str(java_codec.knnVectorsFormat()) + match = re.search( + r"writerClass=([^,)]+), fieldsWriterCalls=(\d+)", diagnostics + ) + if match is None: + raise AssertionError(f"Unexpected writer diagnostics: {diagnostics}") + return match.group(1), int(match.group(2)) + + +def main(): + _validate_pylucene_version(lucene) + lucene.CLASSPATH = os.pathsep.join( + (os.environ[PYLUCENE_TEST_CLASSES_ENV], lucene.CLASSPATH) + ) + runtime = _PyLuceneRuntime.create( + { + "cuvs_java_jar": os.environ["CUVS_LUCENE_CUVS_JAVA_JAR"], + "cuvs_lucene_jar": os.environ["CUVS_LUCENE_JAR"], + "java_library_path": os.environ["JAVA_LIBRARY_PATH"], + } + ) + runtime.System.setProperty(M_PROPERTY, "16") + runtime.System.setProperty(EF_CONSTRUCTION_PROPERTY, "48") + reflected_codec = runtime.Class.forName( + _WRITER_SELECTION_CODEC + ).newInstance() + java_codec = runtime.Codec.cast_(reflected_codec) + vectors = ( + np.random.default_rng(174).standard_normal((4, 32)).astype(np.float32) + ) + + with tempfile.TemporaryDirectory(prefix="pylucene-cpu-fallback-") as temp: + index_path = Path(temp) + runtime.build_index( + index_path, + vectors, + _BuildCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec, + writer_policy="gpu-with-cpu-fallback", + build_parameters={ + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 16, + "ef_construction": 48, + }, + ), + ) + writer_class, writer_calls = _writer_diagnostics(java_codec) + if writer_class != _CPU_HNSW_WRITER: + raise AssertionError( + f"Expected {_CPU_HNSW_WRITER}, found {writer_class}" + ) + + search = runtime.search_index( + index_path, + vectors[[2]], + k=1, + batch_size=1, + num_candidates=1, + ) + first_hit = search.hits[0][0].document_id + + if first_hit != 2: + raise AssertionError(f"Expected document 2, found {first_hit}") + print( + f"writerClass={writer_class} fieldsWriterCalls={writer_calls} " + f"firstHit={first_hit}" + ) + + +if __name__ == "__main__": + main() diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_backend.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_backend.py new file mode 100644 index 0000000000..3d41ff0f23 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_backend.py @@ -0,0 +1,1390 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Build and search behavior tests for the PyLucene backend.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _SearchHit +from cuvs_bench.tests.pylucene._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, +) + + +def _only_search_result(results): + assert len(results) == 1 + return results[0] + + +@pytest.mark.parametrize( + ("score", "distance"), + [(1.0, 0.0), (0.5, 1.0), (0.2, 4.0), (0.0, np.inf)], +) +def test_score_to_squared_euclidean(score, distance): + assert pylucene_backend._score_to_squared_euclidean( + score + ) == pytest.approx(distance) + + +def test_build_dry_run_does_not_initialize_pylucene(tmp_path): + backend = _backend() + result = backend.build( + _dataset(), [_index(tmp_path / "index")], dry_run=True + ) + + assert result.success + assert backend._runtime is None + assert result.metadata["codec"] == _HNSW_CODEC + assert result.metadata["compound_file_policy"] == "lucene-default" + + +def test_build_creates_index_and_records_hnsw_writer_policy(tmp_path): + runtime = _FakeRuntime() + backend = _backend(runtime) + index_path = tmp_path / "index" + + result = backend.build(_dataset(), [_index(index_path)], force=True) + + assert result.success + assert result.index_size_bytes > len(b"index") + assert result.metadata["writer_policy"] == "gpu-with-cpu-fallback" + assert result.metadata["compound_file_policy"] == "lucene-default" + assert result.metadata["hnsw_verification"] == { + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 4, + "codec": _HNSW_CODEC, + "build_parameters": { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + }, + "writer_policy": "gpu-with-cpu-fallback", + "compound_file_policy": "lucene-default", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + assert runtime.configured_resolve_calls == [(32, 32)] + assert runtime.resolve_calls == [] + assert len(runtime.build_calls) == 1 + assert runtime.build_calls[0][2].writer_policy == ("gpu-with-cpu-fallback") + assert runtime.build_calls[0][2].build_parameters == { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + } + assert result.metadata["segment_count"] == 1 + assert result.metadata["segment_document_counts"] == [10] + assert result.metadata["segment_vector_counts"] == [10] + assert runtime.verification_calls == [] + + +def test_build_verifies_persisted_cagra_index(tmp_path): + runtime = _FakeRuntime() + index_path = tmp_path / "index" + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + force=True, + ) + + assert result.success + assert result.metadata["compound_file_policy"] == "disabled" + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 4, + "codec": _CAGRA_CODEC, + "build_parameters": {"codec": _CAGRA_CODEC}, + "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + assert (index_path / pylucene_backend._CAGRA_PROVENANCE_FILE).is_file() + + +def test_build_rejects_persisted_cagra_fallback_and_removes_index(tmp_path): + runtime = _FakeRuntime() + runtime.verification_error = RuntimeError("persisted brute-force index") + index_path = tmp_path / "index" + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + force=True, + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert len(runtime.build_calls) == 1 + assert not index_path.exists() + + +def test_build_reuses_existing_index_without_runtime(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + backend = _backend(runtime) + + result = backend.build(_dataset(), [_index(index_path)]) + + assert result.success + assert result.metadata["skipped"] is True + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-with-cpu-fallback-provenance" + ) + assert runtime.resolve_calls == [] + assert runtime.build_calls == [] + assert runtime.verification_calls == [] + + +def test_build_reports_reused_index_sizing_failure(tmp_path, monkeypatch): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + + def fail_index_size(_index_path): + raise PermissionError("cannot read index size") + + monkeypatch.setattr(pylucene_backend, "_index_size", fail_index_size) + + result = _backend(runtime).build(_dataset(), [_index(index_path)]) + + assert not result.success + assert "cannot read index size" in result.error_message + assert runtime.build_calls == [] + + +def test_build_rejects_reused_hnsw_index_without_provenance(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + runtime = _FakeRuntime() + + result = _backend(runtime).build(_dataset(), [_index(index_path)]) + + assert not result.success + assert "provenance manifest is missing" in result.error_message + assert runtime.build_calls == [] + + +def test_build_verifies_reused_cagra_index(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime() + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert result.success + assert result.metadata["skipped"] is True + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"]["status"] == ( + "gpu-cagra-provenance" + ) + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + assert runtime.build_calls == [] + + +def test_build_rejects_reused_cagra_index_that_cannot_be_verified(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + segments_file = index_path / "segments_1" + runtime = _FakeRuntime() + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert segments_file.read_bytes() == b"index" + assert runtime.build_calls == [] + + +def test_build_rejects_reused_cagra_index_without_provenance_before_runtime( + tmp_path, +): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + runtime = _FakeRuntime() + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert not result.success + assert "CAGRA provenance manifest is missing" in result.error_message + assert runtime.verification_calls == [] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize("force", [False, True]) +def test_build_rejects_unrelated_existing_directory(tmp_path, force): + index_path = tmp_path / "index" + index_path.mkdir() + unrelated_file = index_path / "unrelated" + unrelated_file.write_bytes(b"not a Lucene index") + runtime = _FakeRuntime() + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=force + ) + + assert not result.success + assert "does not contain a Lucene segments file" in result.error_message + assert unrelated_file.exists() + + +def test_build_rejects_existing_index_path_that_is_a_file(tmp_path): + index_path = tmp_path / "index" + index_path.write_bytes(b"not a directory") + + result = _backend(_FakeRuntime()).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "not a directory" in result.error_message + assert index_path.read_bytes() == b"not a directory" + + +def test_build_force_replaces_existing_index(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + runtime = _FakeRuntime() + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert result.success + assert not old_file.exists() + assert (index_path / "segments_1").exists() + + +def test_build_refuses_index_outside_configured_root(tmp_path): + allowed_root = tmp_path / "dataset" / "index" + external_index = tmp_path / "external-index" + external_index.mkdir() + (external_index / "segments_1").write_bytes(b"existing") + old_file = external_index / "old" + old_file.write_bytes(b"must remain") + runtime = _FakeRuntime() + backend = _backend(runtime) + backend.config["index_root"] = str(allowed_root) + + result = backend.build(_dataset(), [_index(external_index)], force=True) + + assert not result.success + assert "immediate child" in result.error_message + assert old_file.read_bytes() == b"must remain" + assert runtime.resolve_calls == [] + assert runtime.build_calls == [] + + +def test_safe_remove_rechecks_configured_root(tmp_path): + allowed_root = tmp_path / "allowed" + external_index = tmp_path / "external-index" + external_index.mkdir() + old_file = external_index / "old" + old_file.write_bytes(b"must remain") + + with pytest.raises(ValueError, match="outside its configured root"): + pylucene_backend._safe_remove_index(external_index, allowed_root) + + assert old_file.read_bytes() == b"must remain" + + +def test_build_preserves_existing_index_if_codec_preflight_fails(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + runtime = _FakeRuntime() + runtime.resolve_error = RuntimeError("codec unavailable") + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "codec unavailable" in result.error_message + assert old_file.exists() + + +def test_build_accepts_production_hnsw_fallback_policy(tmp_path): + runtime = _FakeRuntime() + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert result.success + assert result.metadata["writer_policy"] == "gpu-with-cpu-fallback" + assert runtime.build_calls[0][2].writer_policy == ("gpu-with-cpu-fallback") + + +@pytest.mark.parametrize( + ("dataset", "error"), + [ + (_dataset(dtype=np.float64), "float32"), + (_dataset(distance_metric="cosine"), "Euclidean"), + ], +) +def test_build_rejects_unsupported_dataset(dataset, error, tmp_path): + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert error in result.error_message + + +def test_build_rejects_nonfinite_vectors(tmp_path): + dataset = _dataset() + dataset.training_vectors[0, 0] = np.nan + + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert "finite" in result.error_message + + +@pytest.mark.parametrize( + "codec", + [ + _HNSW_CODEC, + _CAGRA_CODEC, + ], + ids=["hnsw", "cagra"], +) +def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec): + runtime = _FakeRuntime() + + result = _backend(runtime, codec=codec).build( + _dataset(n_base=1), + [_index(tmp_path / "index", codec=codec)], + force=True, + ) + + assert not result.success + assert "at least two training vectors" in result.error_message + assert "does not invoke cuVS" in result.error_message + assert not (tmp_path / "index").exists() + + +def test_build_rejects_unknown_codec(tmp_path): + result = _backend(_FakeRuntime()).build( + _dataset(), + [_index(tmp_path / "index", codec="UnknownCodec")], + force=True, + ) + + assert not result.success + assert "Unsupported PyLucene codec" in result.error_message + + +@pytest.mark.parametrize( + "codec", + [None, "", 0, False], + ids=["none", "empty", "zero", "false"], +) +def test_explicit_invalid_codec_does_not_fall_back_to_backend_config( + tmp_path, codec +): + index = _index(tmp_path / "index") + index.build_param["codec"] = codec + backend = _backend(codec=_HNSW_CODEC) + + build_result = backend.build(_dataset(), [index], dry_run=True) + search_result = _only_search_result( + backend.search(_dataset(), [index], k=3, dry_run=True) + ) + + for result in (build_result, search_result): + assert not result.success + assert f"Unsupported PyLucene codec {codec!r}" in result.error_message + + +def test_build_rejects_unsupported_parameter(tmp_path): + index = _index(tmp_path / "index") + index.build_param["ignored"] = 1 + + result = _backend(_FakeRuntime()).build(_dataset(), [index]) + + assert not result.success + assert "Unsupported PyLucene build parameter" in result.error_message + + +def test_hnsw_build_parameters_have_canonical_defaults(): + assert pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC} + ) == { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + } + + +@pytest.mark.parametrize("name", ["m", "ef_construction"]) +@pytest.mark.parametrize("value", [1, 512]) +def test_hnsw_build_parameter_boundaries_are_inclusive(name, value): + parameters = pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC, name: value} + ) + + assert parameters[name] == value + + +@pytest.mark.parametrize("name", ["m", "ef_construction"]) +@pytest.mark.parametrize( + "value", + [0, 513, True, 1.5, "32", None], + ids=["below-min", "above-max", "bool", "float", "string", "none"], +) +def test_hnsw_build_parameters_reject_invalid_values(name, value): + with pytest.raises(ValueError, match=rf"{name} must be an integer"): + pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC, name: value} + ) + + +@pytest.mark.parametrize("value", [0, 1, None, "true"]) +def test_direct_single_segment_rejects_non_boolean_values(value): + with pytest.raises(TypeError, match="must be a boolean"): + pylucene_backend._normalize_build_params( + { + "codec": _HNSW_CODEC, + "direct_single_segment": value, + } + ) + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("m", 32), + ("ef_construction", 64), + ("direct_single_segment", False), + ], +) +def test_cagra_rejects_hnsw_only_build_parameters(name, value): + with pytest.raises(ValueError, match="apply only to"): + pylucene_backend._normalize_build_params( + {"codec": _CAGRA_CODEC, name: value} + ) + + +@pytest.mark.parametrize("parameters", [None, [], "codec"]) +def test_build_parameters_must_be_a_mapping(parameters): + with pytest.raises(TypeError, match="must be a mapping"): + pylucene_backend._normalize_build_params(parameters) + + +def test_build_propagates_configured_hnsw_parameters_to_runtime(tmp_path): + runtime = _FakeRuntime() + index_path = tmp_path / "index" + index = _index( + index_path, + build_params={"m": 24, "ef_construction": 96}, + ) + + result = _backend(runtime).build(_dataset(), [index], force=True) + + assert result.success + assert runtime.configured_resolve_calls == [(24, 96)] + assert runtime.resolve_calls == [] + build_codec = runtime.build_calls[0][2] + assert build_codec.java_codec == (_HNSW_CODEC, 24, 96) + assert build_codec.build_parameters == { + "codec": _HNSW_CODEC, + "m": 24, + "ef_construction": 96, + "direct_single_segment": False, + } + assert result.build_params == build_codec.build_parameters + + +@pytest.mark.parametrize("index_count", [0, 2]) +def test_build_requires_exactly_one_index(index_count, tmp_path): + indexes = [ + _index(tmp_path / f"index-{index_id}") + for index_id in range(index_count) + ] + + result = _backend(_FakeRuntime()).build(_dataset(), indexes, force=True) + + assert not result.success + assert "exactly one" in result.error_message + + +@pytest.mark.parametrize( + ("vectors", "error"), + [ + (np.empty((0, 4), dtype=np.float32), "at least one vector"), + (np.ones(4, dtype=np.float32), "two-dimensional"), + (np.ones((1, 4097), dtype=np.float32), "4096"), + ], +) +def test_build_rejects_invalid_vector_shape(vectors, error, tmp_path): + dataset = _dataset() + dataset.training_vectors = vectors + + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert error in result.error_message + + +def test_build_removes_partial_index_after_runtime_failure(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("Java build failed") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "Java build failed" in result.error_message + assert not index_path.exists() + + +def test_build_removes_index_when_direct_segment_topology_is_not_one( + tmp_path, +): + runtime = _FakeRuntime() + runtime.topology = pylucene_backend._IndexTopology( + segment_document_counts=(5, 5), + segment_vector_counts=(5, 5), + ) + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), + [ + _index( + index_path, + build_params={"direct_single_segment": True}, + ) + ], + force=True, + ) + + assert not result.success + assert "requested one committed Lucene segment" in result.error_message + assert not index_path.exists() + + +def test_build_reports_chained_runtime_failure(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("rollback failed") + runtime.build_error.__cause__ = RuntimeError("add failed") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "RuntimeError: rollback failed" in result.error_message + assert "caused by RuntimeError: add failed" in result.error_message + assert not index_path.exists() + + +def test_build_removes_partial_index_and_reraises_interrupt(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = KeyboardInterrupt() + index_path = tmp_path / "index" + + with pytest.raises(KeyboardInterrupt): + _backend(runtime).build(_dataset(), [_index(index_path)], force=True) + + assert not index_path.exists() + + +def test_build_preserves_interrupt_when_partial_index_cleanup_fails( + tmp_path, monkeypatch +): + runtime = _FakeRuntime() + runtime.build_error = KeyboardInterrupt() + index_path = tmp_path / "index" + + def fail_cleanup(_index_path, trusted_index_root=None): + raise PermissionError("cleanup denied") + + monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) + with pytest.raises(KeyboardInterrupt) as exc_info: + _backend(runtime).build(_dataset(), [_index(index_path)], force=True) + + assert exc_info.value.__notes__ == [ + "Failed to remove partial index: PermissionError: cleanup denied" + ] + assert index_path.exists() + + +def test_build_reports_partial_index_cleanup_failure(tmp_path, monkeypatch): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("Java build failed") + index_path = tmp_path / "index" + + def fail_cleanup(_index_path, trusted_index_root=None): + raise PermissionError("cleanup denied") + + monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "Java build failed" in result.error_message + assert "failed to remove partial index" in result.error_message + assert "cleanup denied" in result.error_message + assert index_path.exists() + + +def test_search_dry_run_does_not_initialize_pylucene(tmp_path): + backend = _backend() + + result = _only_search_result( + backend.search( + _dataset(), [_index(tmp_path / "index")], k=3, dry_run=True + ) + ) + + assert result.success + assert backend._runtime is None + assert result.neighbors.shape == (0, 3) + assert result.search_params == [{"num_candidates": 3}] + + +def test_search_dry_run_returns_one_result_per_candidate_setting(tmp_path): + backend = _backend() + candidates = [150, 200, 300] + index = _index( + tmp_path / "index", + search_params=[ + {"num_candidates": num_candidates} for num_candidates in candidates + ], + ) + + results = backend.search(_dataset(), [index], k=150, dry_run=True) + + assert len(results) == 3 + assert all(result.success for result in results) + assert backend._runtime is None + assert [result.search_params for result in results] == [ + [{"num_candidates": num_candidates}] for num_candidates in candidates + ] + assert [result.metadata["num_candidates"] for result in results] == ( + candidates + ) + + +def test_search_associates_each_candidate_setting_with_its_result(tmp_path): + class _SelectiveFailureRuntime(_FakeRuntime): + def search_index( + self, + index_path, + query_vectors, + k, + batch_size, + num_candidates=None, + ): + result = super().search_index( + index_path, + query_vectors, + k, + batch_size, + num_candidates, + ) + if num_candidates == 4: + raise RuntimeError("candidate-specific failure") + return result + + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _SelectiveFailureRuntime() + candidates = [3, 4, 5] + index = _index( + index_path, + search_params=[ + {"num_candidates": num_candidates} for num_candidates in candidates + ], + ) + + results = _backend(runtime).search(_dataset(), [index], k=3) + + assert len(results) == 3 + assert [result.success for result in results] == [True, False, True] + assert [result.search_params for result in results] == [ + [{"num_candidates": num_candidates}] for num_candidates in candidates + ] + assert "candidate-specific failure" in results[1].error_message + assert [call[4] for call in runtime.search_calls] == candidates + + +def test_search_converts_hits_scores_and_padding(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=3, score=1.0), + _SearchHit(document_id=7, score=0.5), + ], + [_SearchHit(document_id=2, score=0.2)], + ] + ) + + result = _only_search_result( + _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=1 + ) + ) + + assert result.success + np.testing.assert_array_equal( + result.neighbors, + np.array([[3, 7, -1], [2, -1, -1]], dtype=np.int64), + ) + np.testing.assert_allclose(result.distances[0, :2], [0.0, 1.0]) + assert result.distances[1, 0] == pytest.approx(4.0) + assert np.isinf(result.distances[:, -1]).all() + assert result.search_time_ms == pytest.approx(2.0) + assert result.metadata["latency_seconds"] == pytest.approx(0.001) + assert result.queries_per_second == pytest.approx(1000.0) + assert result.latency_percentiles == { + "p50": 1.0, + "p95": 1.0, + "p99": 1.0, + } + assert result.metadata["num_batches"] == 2 + assert result.metadata["top_k"] == 3 + assert result.metadata["num_candidates"] == 3 + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-with-cpu-fallback-provenance" + ) + assert "per_search_param_results" not in result.metadata + assert runtime.search_calls[0][2:] == (3, 1, 3) + assert runtime.verification_calls == [] + + +def test_search_end_to_end_timing_starts_before_input_verification( + tmp_path, monkeypatch +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + backend = _backend(_FakeRuntime()) + events = [] + times = iter((10.0, 10.012)) + load_search_inputs = backend._load_search_inputs + + def record_time(): + events.append("time") + return next(times) + + def record_input_loading(dataset): + events.append("inputs") + return load_search_inputs(dataset) + + monkeypatch.setattr(pylucene_backend.time, "perf_counter", record_time) + monkeypatch.setattr(backend, "_load_search_inputs", record_input_loading) + + result = _only_search_result( + backend.search(_dataset(), [_index(index_path)], k=3, batch_size=2) + ) + + assert result.success + assert events == ["time", "inputs", "time"] + assert result.metadata["end_to_end_time_ms"] == pytest.approx(12.0) + assert result.metadata["non_query_overhead_time_ms"] == pytest.approx(11.0) + + +@pytest.mark.parametrize( + ("runtime", "error"), + [ + ( + _FakeRuntime(document_count=9), + "document count does not match index provenance", + ), + ( + _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=3, score=1.0), + _SearchHit(document_id=3, score=0.5), + ], + [_SearchHit(document_id=2, score=1.0)], + ] + ), + "duplicate stored ID", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=10, score=1.0)], + [_SearchHit(document_id=2, score=1.0)], + ] + ), + "out-of-range stored ID", + ), + ( + _FakeRuntime(hits=[[_SearchHit(document_id=0, score=1.0)]]), + "unexpected number of query results", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=float("nan"))], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "non-finite score", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=1.1)], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "outside the Euclidean range", + ), + ( + _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=0, score=1.0), + _SearchHit(document_id=1, score=0.9), + _SearchHit(document_id=2, score=0.8), + _SearchHit(document_id=3, score=0.7), + ], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "too many hits", + ), + ( + _FakeRuntime(batch_latencies_ms=[1.0, 2.0]), + "unexpected number of batch latencies", + ), + ( + _FakeRuntime(batch_latencies_ms=[float("nan")]), + "invalid batch latency", + ), + ], + ids=[ + "document-count", + "duplicate-id", + "out-of-range-id", + "query-count", + "score", + "score-range", + "hit-count", + "latency-count", + "batch-latency", + ], +) +def test_search_rejects_invalid_runtime_results(tmp_path, runtime, error): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + result = _only_search_result( + _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=2 + ) + ) + + assert not result.success + assert error in result.error_message + + +def test_search_verifies_cagra_index_before_querying(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime() + + result = _only_search_result( + _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + ) + + assert result.success + assert runtime.verification_calls == [(index_path, 10, 4)] + assert len(runtime.search_calls) == 1 + assert result.metadata["cagra_provenance"]["status"] == ( + "gpu-cagra-provenance" + ) + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + + +def test_search_rejects_unverified_cagra_index_before_querying(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime() + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _only_search_result( + _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize("manifest_state", ["missing", "stale"]) +def test_search_rejects_invalid_cagra_provenance_before_runtime( + tmp_path, manifest_state +): + index_path = tmp_path / "index" + manifest_path = _prepare_cagra_index(index_path) + if manifest_state == "missing": + manifest_path.unlink() + else: + (index_path / "segments_1").write_bytes(b"changed commit") + runtime = _FakeRuntime() + + result = _only_search_result( + _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + ) + + assert not result.success + assert "CAGRA provenance" in result.error_message + assert runtime.verification_calls == [] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize( + ("index", "k", "batch_size", "search_threads", "mode", "error"), + [ + ( + _index(Path("/tmp/index")), + 0, + 10000, + None, + "latency", + "k must be positive", + ), + ( + _index(Path("/tmp/index")), + 3, + 10000, + 2, + "latency", + "one search thread", + ), + ( + _index(Path("/tmp/index")), + 3, + 10000, + None, + "throughput", + "only latency mode", + ), + ( + _index(Path("/tmp/index")), + 3, + 0, + None, + "latency", + "batch_size must be positive", + ), + ( + _index(Path("/tmp/index"), search_params=[{"search_width": 16}]), + 3, + 10000, + None, + "latency", + "Unsupported PyLucene search parameter", + ), + ( + _index(Path("/tmp/index"), codec=_CAGRA_CODEC), + 1025, + 10000, + None, + "latency", + "GPU brute-force", + ), + ], +) +def test_search_rejects_unsupported_options( + index, k, batch_size, search_threads, mode, error +): + result = _only_search_result( + _backend(_FakeRuntime()).search( + _dataset(), + [index], + k=k, + batch_size=batch_size, + mode=mode, + search_threads=search_threads, + dry_run=True, + ) + ) + + assert not result.success + assert error in result.error_message + + +@pytest.mark.parametrize( + "num_candidates", + [0, -1, True, 1.5, "3", None], + ids=["zero", "negative", "bool", "float", "string", "none"], +) +def test_search_rejects_invalid_num_candidates(num_candidates, tmp_path): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[{"num_candidates": num_candidates}], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "num_candidates must be a positive integer" in result.error_message + + +def test_search_rejects_num_candidates_below_top_k(tmp_path): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[{"num_candidates": 2}], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "greater than or equal to k (3)" in result.error_message + + +def test_search_rejects_num_candidates_for_cagra(tmp_path): + result = _only_search_result( + _backend(codec=_CAGRA_CODEC).search( + _dataset(), + [ + _index( + tmp_path / "index", + codec=_CAGRA_CODEC, + search_params=[{"num_candidates": 3}], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "num_candidates applies only to" in result.error_message + + +@pytest.mark.parametrize("search_parameters", [1, None, "candidates"]) +def test_search_parameter_entries_must_be_mappings( + search_parameters, tmp_path +): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[search_parameters], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "Every PyLucene search parameter must be a mapping" in ( + result.error_message + ) + + +def test_resolve_search_plans_normalizes_the_complete_request(tmp_path): + index = _index(tmp_path / "index", search_params=[]) + + plans = _backend()._resolve_search_plans( + index, + k=3, + batch_size=7, + mode="latency", + search_threads="1", + ) + + assert plans == [ + pylucene_backend._SearchPlan( + index_path=tmp_path / "index", + codec_name=_HNSW_CODEC, + build_parameters={ + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + }, + search_parameters={"num_candidates": 3}, + num_candidates=3, + k=3, + batch_size=7, + mode="latency", + ) + ] + + +def test_search_validation_preserves_first_error_precedence(tmp_path): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[{"unsupported": True}], + ) + ], + k=0, + batch_size=0, + mode="throughput", + search_threads=2, + dry_run=True, + ) + ) + + assert not result.success + assert result.error_message == "k must be positive" + + +@pytest.mark.parametrize("index_count", [0, 2]) +def test_search_requires_exactly_one_index(index_count, tmp_path): + indexes = [ + _index(tmp_path / f"index-{index_id}") + for index_id in range(index_count) + ] + + result = _only_search_result( + _backend(_FakeRuntime()).search(_dataset(), indexes, k=3) + ) + + assert not result.success + assert "exactly one" in result.error_message + + +def test_search_rejects_missing_index(tmp_path): + result = _only_search_result( + _backend(_FakeRuntime()).search( + _dataset(), [_index(tmp_path / "missing")], k=3 + ) + ) + + assert not result.success + assert "does not exist" in result.error_message + + +def test_search_rejects_empty_query_vectors(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + dataset = _dataset() + dataset.query_vectors = np.empty((0, 4), dtype=np.float32) + + result = _only_search_result( + _backend(_FakeRuntime()).search(dataset, [_index(index_path)], k=3) + ) + + assert not result.success + assert "at least one vector" in result.error_message + + +def test_search_reports_runtime_query_dimension_mismatch(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime(index_dimensions=8) + + result = _only_search_result( + _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) + ) + + assert not result.success + assert "dimensions do not match the Lucene index" in result.error_message + + +def test_search_rejects_runtime_dimensions_that_disagree_with_provenance( + tmp_path, +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path, dimensions=4) + runtime = _FakeRuntime( + index_dimensions=8, + validate_query_dimensions=False, + ) + + result = _only_search_result( + _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) + ) + + assert not result.success + assert ( + "Lucene index dimensions do not match index provenance" + in result.error_message + ) + assert len(runtime.search_calls) == 1 + + +def test_search_rejects_query_dimensions_that_disagree_with_dataset(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + dataset = _dataset(dimensions=4) + dataset.query_vectors = np.zeros((2, 8), dtype=np.float32) + runtime = _FakeRuntime() + + result = _only_search_result( + _backend(runtime).search(dataset, [_index(index_path)], k=3) + ) + + assert not result.success + assert "Query vector dimensions do not match the dataset" in ( + result.error_message + ) + assert runtime.search_calls == [] + + +def test_search_reports_runtime_failure(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + runtime.search_error = RuntimeError("Java search failed") + + result = _only_search_result( + _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + ) + + assert not result.success + assert "Java search failed" in result.error_message + + +@pytest.mark.parametrize( + "manifest_state", + ["missing", "corrupt", "stale", "wrong-codec"], +) +def test_search_rejects_invalid_hnsw_provenance_before_runtime( + tmp_path, manifest_state +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + if manifest_state == "missing": + manifest_path.unlink() + elif manifest_state == "corrupt": + manifest_path.write_text("{") + elif manifest_state == "stale": + (index_path / "segments_1").write_bytes(b"changed commit") + else: + payload = json.loads(manifest_path.read_text()) + payload["codec"] = "DifferentCodec" + manifest_path.write_text(json.dumps(payload)) + + runtime = _FakeRuntime() + result = _only_search_result( + _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + ) + + assert not result.success + assert "provenance" in result.error_message.lower() + assert runtime.search_calls == [] + + +def test_cleanup_releases_runtime_reference(): + backend = _backend(_FakeRuntime()) + + backend.cleanup() + + assert backend._runtime is None diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_cagra_verifier.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_cagra_verifier.py new file mode 100644 index 0000000000..fa320fdde4 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_cagra_verifier.py @@ -0,0 +1,815 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Persisted CAGRA verifier unit tests.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend + + +class _FakeCagraMetadataInput: + def __init__(self, integers, variable_longs): + self._integers = iter(integers) + self._variable_longs = iter(variable_longs) + + def readInt(self): + return next(self._integers) + + def readVLong(self): + return next(self._variable_longs) + + +class _FakeChecksumInput(_FakeCagraMetadataInput): + def __init__(self, integers, variable_longs): + super().__init__(integers, variable_longs) + self.closed = False + + def close(self): + self.closed = True + + +class _FakeIndexInput: + def __init__(self, payload_start=57, payload_length=1000): + self.payload_start = payload_start + self.file_length = payload_start + payload_length + 16 + self.closed = False + + def getFilePointer(self): + return self.payload_start + + def length(self): + return self.file_length + + def close(self): + self.closed = True + + +class _FakeCodecUtil: + def __init__(self, error_at=None): + self.error_at = error_at + self.header_calls = [] + self.footer_calls = [] + self.checksum_calls = [] + self.retrieved_checksum_calls = [] + + def checkIndexHeader(self, *args): + self.header_calls.append(args) + if self.error_at == "header": + raise RuntimeError("unsupported metadata version") + return 0 + + def checkFooter(self, metadata_input): + self.footer_calls.append(metadata_input) + if self.error_at == "footer": + raise RuntimeError("invalid checksum footer") + + def footerLength(self): + return 16 + + def checksumEntireFile(self, data_input): + self.checksum_calls.append(data_input) + if self.error_at == "data-checksum": + raise RuntimeError("invalid data checksum") + return 0 + + def retrieveChecksum(self, data_input): + self.retrieved_checksum_calls.append(data_input) + return 0 + + +class _FakeFieldInfos: + def __init__(self, field_infos): + self._field_infos = { + field_info.number: field_info for field_info in field_infos + } + + def fieldInfo(self, field_number): + return self._field_infos.get(field_number) + + def __iter__(self): + return iter(self._field_infos.values()) + + +def _fake_cagra_index_verifier( + metadata_input, + *, + metadata_files=("_0.vemc",), + codec_util=None, + data_payload_length=1000, + field_name="vector", + field_dimensions=32, + field_updates=False, + max_documents=4, + deletion_count=0, + soft_deletion_count=0, + extra_vector_field=False, + data_error=None, +): + data_input = _FakeIndexInput(payload_length=data_payload_length) + vector_encoding = object() + vector_similarity = object() + field_info = SimpleNamespace( + number=1, + getName=lambda: field_name, + getVectorDimension=lambda: field_dimensions, + getVectorEncoding=lambda: vector_encoding, + getVectorSimilarityFunction=lambda: vector_similarity, + ) + all_field_infos = [field_info] + if extra_vector_field: + all_field_infos.append( + SimpleNamespace( + number=2, + getName=lambda: "extra-vector", + getVectorDimension=lambda: field_dimensions, + getVectorEncoding=lambda: vector_encoding, + getVectorSimilarityFunction=lambda: vector_similarity, + ) + ) + field_infos = _FakeFieldInfos(all_field_infos) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: field_infos + ) + + def open_data(_file_name, _context): + if data_error is not None: + raise data_error + return data_input + + directory = SimpleNamespace( + openChecksumInput=lambda _file_name: metadata_input, + openInput=open_data, + close=lambda: None, + ) + segment_info = SimpleNamespace( + info=SimpleNamespace( + name="_0", + getId=lambda: b"segment-id", + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: max_documents, + ), + files=lambda: metadata_files, + hasFieldUpdates=lambda: field_updates, + hasDeletions=lambda: deletion_count > 0, + getDelCount=lambda: deletion_count, + getSoftDelCount=lambda: soft_deletion_count, + ) + verifier = pylucene_backend._CagraIndexVerifier( + attach_current_thread=lambda: None, + paths=SimpleNamespace(get=lambda path: path), + codec_util=codec_util or _FakeCodecUtil(), + field_info=SimpleNamespace( + cast_=lambda raw_field_info: raw_field_info + ), + segment_commit_info=SimpleNamespace( + cast_=lambda raw_segment_info: raw_segment_info + ), + segment_infos=SimpleNamespace( + readLatestCommit=lambda _directory: [segment_info] + ), + vector_encoding=SimpleNamespace(FLOAT32=vector_encoding), + vector_similarity_function=SimpleNamespace( + EUCLIDEAN=vector_similarity + ), + fs_directory=SimpleNamespace(open=lambda _path: directory), + io_context=SimpleNamespace(READONCE=object()), + ) + verifier._test_data_input = data_input + verifier._test_directory = directory + verifier._test_segment_info = segment_info + return verifier + + +def _cagra_data_context(verifier, index_path): + return pylucene_backend._CagraDataFileContext( + index_path=index_path, + directory=verifier._test_directory, + segment_info=verifier._test_segment_info, + suffix="", + metadata_file="_0.vemc", + ) + + +@pytest.mark.parametrize( + ("metadata_file", "expected_suffix"), + [("_0.vemc", ""), ("_0_CuVS_0.vemc", "CuVS_0")], +) +def test_cagra_segment_suffix(metadata_file, expected_suffix): + assert ( + pylucene_backend._CagraIndexVerifier._segment_suffix( + "_0", metadata_file + ) + == expected_suffix + ) + + +def test_cagra_segment_suffix_rejects_unrelated_metadata_file(): + with pytest.raises(RuntimeError, match="does not match segment"): + pylucene_backend._CagraIndexVerifier._segment_suffix("_0", "_1.vemc") + + +def test_read_cagra_fields_accepts_only_persisted_cagra_data(): + metadata_input = _FakeCagraMetadataInput( + integers=[ + 1, + 1, + 0, + 32, + 4, + 2, + 1, + 0, + 32, + 0, + -1, + ], + variable_longs=[57, 1000, 1057, 0, 0, 0, 0, 0], + ) + + assert pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) == [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + + +@pytest.mark.parametrize( + ("cagra_length", "brute_force_length", "error"), + [ + (0, 512, "persisted brute-force"), + (1000, 512, "persisted brute-force"), + ], +) +def test_read_cagra_fields_rejects_non_cagra_only_data( + cagra_length, brute_force_length, error +): + metadata_input = _FakeCagraMetadataInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[ + 57, + cagra_length, + 57 + cagra_length, + brute_force_length, + ], + ) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +def test_read_cagra_fields_rejects_data_for_empty_field(): + metadata_input = _FakeCagraMetadataInput( + integers=[1, 1, 0, 32, 0, -1], + variable_longs=[57, 1, 58, 0], + ) + + with pytest.raises(RuntimeError, match="empty field"): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +@pytest.mark.parametrize( + ("encoding", "similarity", "dimensions", "error"), + [ + (0, 0, 32, "encoding ordinal 0"), + (1, 1, 32, "similarity ordinal 1"), + (1, 0, 0, "invalid field metadata"), + (1, 0, 4097, "invalid field metadata"), + ], +) +def test_read_cagra_fields_rejects_unsupported_vector_semantics( + encoding, similarity, dimensions, error +): + metadata_input = _FakeCagraMetadataInput( + integers=[1, encoding, similarity, dimensions, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +def test_verify_cagra_index_validates_header_footer_and_expected_count( + tmp_path, +): + metadata_input = _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + metadata_input, codec_util=codec_util + ) + + verification = verifier.verify_index( + tmp_path, expected_vector_count=4, expected_dimensions=32 + ) + + assert verification.to_metadata() == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 4, + "dimensions": 32, + } + assert codec_util.header_calls == [ + ( + metadata_input, + "Lucene102CuVSVectorsFormatMeta", + 0, + 0, + b"segment-id", + "", + ), + ( + verifier._test_data_input, + "Lucene102CuVSVectorsFormatIndex", + 0, + 0, + b"segment-id", + "", + ), + ] + assert codec_util.footer_calls == [metadata_input] + assert codec_util.checksum_calls == [verifier._test_data_input] + assert metadata_input.closed is True + assert verifier._test_data_input.closed is True + + +def test_verify_cagra_index_traverses_segments_and_closes_each_input( + tmp_path, +): + events = [] + metadata_inputs = { + name: _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + for name in ("_0.vemc", "_1.vemc") + } + data_inputs = {name: _FakeIndexInput() for name in ("_0.vcag", "_1.vcag")} + verifier = _fake_cagra_index_verifier(metadata_inputs["_0.vemc"]) + + def segment(name): + field_info = SimpleNamespace( + number=1, + getName=lambda: "vector", + getVectorDimension=lambda: 32, + getVectorEncoding=lambda: verifier.VectorEncoding.FLOAT32, + getVectorSimilarityFunction=( + lambda: verifier.VectorSimilarityFunction.EUCLIDEAN + ), + ) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: ( + _FakeFieldInfos([field_info]) + ) + ) + return SimpleNamespace( + info=SimpleNamespace( + name=name, + getId=lambda: name.encode(), + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: 4, + ), + files=lambda: (f"{name}.vemc",), + hasFieldUpdates=lambda: False, + hasDeletions=lambda: False, + getDelCount=lambda: 0, + getSoftDelCount=lambda: 0, + ) + + def open_metadata(file_name): + events.append(f"open {file_name}") + metadata_input = metadata_inputs[file_name] + + def close_metadata(): + metadata_input.closed = True + events.append(f"close {file_name}") + + metadata_input.close = close_metadata + return metadata_input + + def open_data(file_name, _context): + events.append(f"open {file_name}") + data_input = data_inputs[file_name] + + def close_data(): + data_input.closed = True + events.append(f"close {file_name}") + + data_input.close = close_data + return data_input + + verifier.SegmentInfos.readLatestCommit = lambda _directory: [ + segment("_0"), + segment("_1"), + ] + verifier._test_directory.openChecksumInput = open_metadata + verifier._test_directory.openInput = open_data + verifier._test_directory.close = lambda: events.append("close directory") + + verification = verifier.verify_index(tmp_path) + + assert verification == pylucene_backend._CagraIndexVerification( + segment_count=2, + field_count=2, + vector_count=8, + dimensions=32, + ) + assert events == [ + "open _0.vemc", + "close _0.vemc", + "open _0.vcag", + "close _0.vcag", + "open _1.vemc", + "close _1.vemc", + "open _1.vcag", + "close _1.vcag", + "close directory", + ] + assert all(item.closed for item in metadata_inputs.values()) + assert all(item.closed for item in data_inputs.values()) + + +@pytest.mark.parametrize("error_at", ["header", "footer"]) +def test_verify_cagra_index_fails_closed_on_invalid_format(error_at, tmp_path): + metadata_input = _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + verifier = _fake_cagra_index_verifier( + metadata_input, codec_util=_FakeCodecUtil(error_at=error_at) + ) + + with pytest.raises(RuntimeError, match="metadata format v0"): + verifier.verify_index(tmp_path) + + assert metadata_input.closed is True + + +def test_verify_cagra_index_rejects_missing_segment_metadata(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), metadata_files=() + ) + + with pytest.raises(RuntimeError, match="no .vemc metadata"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_only_empty_fields(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 0, -1], + variable_longs=[0, 0, 0, 0], + ), + data_payload_length=0, + ) + + with pytest.raises(RuntimeError, match="without matching CAGRA metadata"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_vector_count_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + ) + + with pytest.raises(RuntimeError, match="4 vectors; expected 5"): + verifier.verify_index(tmp_path, expected_vector_count=5) + + +def test_verify_cagra_index_rejects_dimension_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + ) + + with pytest.raises(RuntimeError, match="32 dimensions; expected 16"): + verifier.verify_index(tmp_path, expected_dimensions=16) + + +@pytest.mark.parametrize( + ("runtime_kwargs", "expected_counts"), + [ + ({"deletion_count": 1}, "deleted=1, soft_deleted=0"), + ({"soft_deletion_count": 1}, "deleted=0, soft_deleted=1"), + ], + ids=["hard-deletion", "soft-deletion"], +) +def test_verify_cagra_index_rejects_committed_deletions( + tmp_path, runtime_kwargs, expected_counts +): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + **runtime_kwargs, + ) + + with pytest.raises(RuntimeError, match=expected_counts): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_segment_document_count_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + max_documents=5, + ) + + with pytest.raises(RuntimeError, match="4 vectors for 5 documents"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_unaccounted_vector_field(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + extra_vector_field=True, + ) + + with pytest.raises(RuntimeError, match=r"metadata=\[1\], Lucene=\[1, 2\]"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_duplicate_field_across_metadata_files( + tmp_path, +): + metadata_inputs = { + name: _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + for name in ("_0.vemc", "_0_CuVS_0.vemc") + } + verifier = _fake_cagra_index_verifier( + metadata_inputs["_0.vemc"], + metadata_files=tuple(metadata_inputs), + ) + verifier._test_directory.openChecksumInput = metadata_inputs.__getitem__ + verifier._test_directory.openInput = ( + lambda _file_name, _context: _FakeIndexInput() + ) + + with pytest.raises(RuntimeError, match="duplicate field 1"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_inconsistent_dimensions_across_segments( + tmp_path, +): + metadata_inputs = { + "_0.vemc": _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + "_1.vemc": _FakeChecksumInput( + integers=[1, 1, 0, 16, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + } + verifier = _fake_cagra_index_verifier(metadata_inputs["_0.vemc"]) + + def segment(name, dimensions): + field_info = SimpleNamespace( + number=1, + getName=lambda: "vector", + getVectorDimension=lambda: dimensions, + getVectorEncoding=lambda: verifier.VectorEncoding.FLOAT32, + getVectorSimilarityFunction=( + lambda: verifier.VectorSimilarityFunction.EUCLIDEAN + ), + ) + field_infos = _FakeFieldInfos([field_info]) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: field_infos + ) + return SimpleNamespace( + info=SimpleNamespace( + name=name, + getId=lambda: name.encode(), + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: 4, + ), + files=lambda: (f"{name}.vemc",), + hasFieldUpdates=lambda: False, + hasDeletions=lambda: False, + getDelCount=lambda: 0, + getSoftDelCount=lambda: 0, + ) + + segments = [segment("_0", 32), segment("_1", 16)] + verifier.SegmentInfos.readLatestCommit = lambda _directory: segments + verifier._test_directory.openChecksumInput = metadata_inputs.__getitem__ + verifier._test_directory.openInput = ( + lambda _file_name, _context: _FakeIndexInput() + ) + + with pytest.raises(RuntimeError, match=r"dimensions: \[16, 32\]"): + verifier.verify_index(tmp_path) + + +@pytest.mark.parametrize( + ("runtime_kwargs", "error"), + [ + ({"field_name": "other"}, "unexpected field"), + ({"field_dimensions": 16}, "field metadata"), + ({"field_updates": True}, "field-info updates"), + ({"data_payload_length": 999}, "do not exactly cover"), + ], +) +def test_verify_cagra_index_rejects_foreign_or_inconsistent_index( + tmp_path, runtime_kwargs, error +): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + **runtime_kwargs, + ) + + with pytest.raises(RuntimeError, match=error): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_corrupt_data_file(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + codec_util=_FakeCodecUtil(error_at="data-checksum"), + ) + + with pytest.raises(RuntimeError, match="invalid data checksum"): + verifier.verify_index(tmp_path) + + assert verifier._test_data_input.closed is True + + +def test_verify_cagra_index_rejects_missing_data_file(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + data_error=FileNotFoundError("_0.vcag"), + ) + + with pytest.raises(RuntimeError, match="cannot read '_0.vcag'"): + verifier.verify_index(tmp_path) + + +def test_file_signature_exposes_named_stat_fields(tmp_path): + data_path = tmp_path / "data.vcag" + data_path.write_bytes(b"data") + file_stat = data_path.stat() + + signature = pylucene_backend._file_signature(data_path) + + assert signature == pylucene_backend._FileSignature( + resolved_path=str(data_path.resolve()), + device=file_stat.st_dev, + inode=file_stat.st_ino, + size=file_stat.st_size, + modified_at_ns=file_stat.st_mtime_ns, + changed_at_ns=file_stat.st_ctime_ns, + ) + + +def test_cagra_data_checksum_is_cached_for_unchanged_file( + tmp_path, monkeypatch +): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + (tmp_path / "_0.vcag").write_bytes(b"data") + data_ctime_ns = (tmp_path / "_0.vcag").stat().st_ctime_ns + monkeypatch.setattr( + pylucene_backend.time, + "time_ns", + lambda: ( + data_ctime_ns + pylucene_backend._CAGRA_CACHE_MIN_FILE_AGE_NS + 1 + ), + ) + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + for _ in range(2): + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [verifier._test_data_input] + assert codec_util.retrieved_checksum_calls == [verifier._test_data_input] + + +def test_cagra_data_checksum_cache_invalidates_on_size_change(tmp_path): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + data_path = tmp_path / "_0.vcag" + data_path.write_bytes(b"data") + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + verifier._verify_cagra_data_file(context, fields) + data_path.write_bytes(b"changed-size") + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [ + verifier._test_data_input, + verifier._test_data_input, + ] + assert codec_util.retrieved_checksum_calls == [] + + +def test_cagra_data_checksum_cache_invalidates_same_size_restored_mtime( + tmp_path, +): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + data_path = tmp_path / "_0.vcag" + data_path.write_bytes(b"data") + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + verifier._verify_cagra_data_file(context, fields) + original_stat = data_path.stat() + data_path.write_bytes(b"evil") + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + assert data_path.stat().st_mtime_ns == original_stat.st_mtime_ns + + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [ + verifier._test_data_input, + verifier._test_data_input, + ] + assert codec_util.retrieved_checksum_calls == [] diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_cli_config.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_cli_config.py new file mode 100644 index 0000000000..35edcc7f69 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_cli_config.py @@ -0,0 +1,1204 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""CLI and configuration tests for the PyLucene backend.""" + +from __future__ import annotations + +import csv +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from click.testing import CliRunner + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends import get_registry +from cuvs_bench.backends.base import SearchResult +from cuvs_bench.backends.pylucene import ( + PyLuceneConfigLoader, + _SearchHit, +) +from cuvs_bench.backends.registry import list_config_loaders +from cuvs_bench.backends.search_spaces import get_search_space +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator +from cuvs_bench.tests.pylucene._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _write_test_bin, +) + + +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} + + +def _install_single_trial_optuna(monkeypatch): + suggested_ranges = {} + selected_parameters = {} + + def suggest_int(name, minimum, maximum, *, log=False): + suggested_ranges[name] = (minimum, maximum, log) + selected_parameters[name] = minimum + return minimum + + trial = SimpleNamespace( + number=0, + suggest_int=suggest_int, + suggest_float=lambda name, minimum, maximum, **kwargs: minimum, + suggest_categorical=lambda name, choices: choices[0], + ) + study = SimpleNamespace( + best_trial=trial, + best_params=selected_parameters, + best_value=0.0, + ) + + def optimize(objective, **kwargs): + study.best_value = objective(trial) + + study.optimize = optimize + fake_optuna = SimpleNamespace( + TrialPruned=RuntimeError, + create_study=lambda **kwargs: study, + logging=SimpleNamespace(WARNING=0, set_verbosity=lambda level: None), + ) + monkeypatch.setitem(sys.modules, "optuna", fake_optuna) + return suggested_ranges + + +def _hnsw_index_name( + algorithm: str, + group: str, + *, + subset_scope: str | None = None, + m: int = 32, + ef_construction: int = 32, + direct_single_segment: bool = False, +) -> str: + identity = f"{algorithm}[group={group}]" + if subset_scope is not None: + identity = f"{identity}[scope={subset_scope}]" + return ( + f"{identity}[codec={_HNSW_CODEC}]" + f"[m={m}]" + f"[ef_construction={ef_construction}]" + f"[direct_single_segment={str(direct_single_segment).lower()}]" + ) + + +def _prepare_cli_dataset(dataset_path: Path) -> None: + training_vectors = np.asarray( + [ + [0.0, 0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + query_vectors = training_vectors[:2].copy() + groundtruth_neighbors = np.asarray([[0], [1]], dtype=np.int32) + dataset_dir = dataset_path / "test-dataset" + _write_test_bin(dataset_dir / "base.fbin", training_vectors) + _write_test_bin(dataset_dir / "query.fbin", query_vectors) + _write_test_bin( + dataset_dir / "groundtruth.neighbors.ibin", groundtruth_neighbors + ) + + +def _pylucene_cli_args( + backend_config: Path, + config_dir: Path, + dataset_path: Path, + *, + dry_run: bool = False, +) -> list[str]: + args = [ + "--backend-config", + str(backend_config), + "--dataset-configuration", + str(config_dir / "datasets" / "datasets.yaml"), + "--configuration", + str(config_dir / "algos" / "pylucene_test.yaml"), + "--dataset", + "test-dataset", + "--dataset-path", + str(dataset_path), + "--algorithms", + "pylucene_test", + "--groups", + "test", + "--batch-size", + "2", + "-k", + "1", + "-m", + "latency", + "--build", + "--search", + "--force", + ] + if dry_run: + args.append("--dry-run") + return args + + +@pytest.fixture +def config_dir(tmp_path): + (tmp_path / "datasets").mkdir() + (tmp_path / "datasets" / "datasets.yaml").write_text( + """\ +- name: test-dataset + base_file: test-dataset/base.fbin + query_file: test-dataset/query.fbin + groundtruth_neighbors_file: test-dataset/groundtruth.neighbors.ibin + distance: euclidean + dims: 4 +""" + ) + (tmp_path / "algos").mkdir() + (tmp_path / "algos" / "pylucene_test.yaml").write_text( + f"""\ +name: pylucene_test +groups: + base: + build: + codec: [{_HNSW_CODEC}, {_CAGRA_CODEC}] + search: {{}} + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + (tmp_path / "algos" / "unrelated.yaml").write_text( + """\ +name: unrelated +groups: + base: + build: {} + search: {} +""" + ) + return tmp_path + + +def test_backend_and_loader_are_registered(): + assert get_registry().is_registered("pylucene") + assert list_config_loaders()["pylucene"] is PyLuceneConfigLoader + + +def test_import_does_not_load_pylucene(): + subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; import cuvs_bench.backends; " + "assert 'lucene' not in sys.modules" + ), + ], + check=True, + ) + + +def test_cli_backend_config_supports_pylucene_dry_run(config_dir, tmp_path): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + result_path = dataset_path / "test-dataset" / "result" + assert not result_path.exists() + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args( + backend_config, + config_dir, + dataset_path, + dry_run=True, + ), + ) + + assert result.exit_code == 0, result.output + assert "Would build PyLucene index" in result.output + assert "Would search PyLucene index" in result.output + assert not result_path.exists() + + +def test_cli_build_search_persists_metrics_and_build_join( + config_dir, tmp_path, monkeypatch +): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + _prepare_cli_dataset(dataset_path) + result_path = dataset_path / "test-dataset" / "result" + assert not result_path.exists() + + runtime = _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=1.0)], + [_SearchHit(document_id=1, score=1.0)], + ] + ) + monkeypatch.setattr( + pylucene_backend._PyLuceneRuntime, + "create", + staticmethod(lambda _config: runtime), + ) + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args(backend_config, config_dir, dataset_path), + ) + + assert result.exit_code == 0, result.output + index_name = _hnsw_index_name("pylucene_test", "test") + index_path = dataset_path / "test-dataset" / "index" / index_name + assert (index_path / "segments_1").is_file() + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + build_csv = result_path / "build" / "pylucene_test,test.csv" + with build_csv.open(newline="") as file: + build_rows = list(csv.DictReader(file)) + assert len(build_rows) == 1 + build_row = build_rows[0] + assert build_row["index_name"] == index_name + assert float(build_row["time"]) > 0.0 + assert build_row["codec"] == _HNSW_CODEC + assert build_row["writer_policy"] == "gpu-with-cpu-fallback" + + raw_csv = result_path / "search" / "pylucene_test,test,k1,bs2,raw.csv" + with raw_csv.open(newline="") as file: + csv_rows = list(csv.DictReader(file)) + assert len(csv_rows) == 1 + csv_row = csv_rows[0] + assert csv_row["index_name"] == index_name + assert float(csv_row["recall"]) == pytest.approx(1.0) + assert float(csv_row["throughput"]) == pytest.approx(2000.0) + assert float(csv_row["latency"]) == pytest.approx(0.001) + assert float(csv_row["build time"]) > 0.0 + assert float(csv_row["p50"]) == pytest.approx(1.0) + assert (raw_csv.parent / "pylucene_test,test,k1,bs2,latency.csv").is_file() + assert ( + raw_csv.parent / "pylucene_test,test,k1,bs2,throughput.csv" + ).is_file() + + +def test_cli_returns_nonzero_without_persisting_failed_measurement( + config_dir, tmp_path, monkeypatch +): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + _prepare_cli_dataset(dataset_path) + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("intentional build failure") + monkeypatch.setattr( + pylucene_backend._PyLuceneRuntime, + "create", + staticmethod(lambda _config: runtime), + ) + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args(backend_config, config_dir, dataset_path), + ) + + assert result.exit_code != 0 + assert "intentional build failure" in result.output + build_csv = ( + dataset_path + / "test-dataset" + / "result" + / "build" + / "pylucene_test,test.csv" + ) + assert not build_csv.exists() + assert not (dataset_path / "test-dataset" / "result" / "search").exists() + + +def test_config_loader_expands_codecs_and_forwards_runtime_config(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + dataset_config, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="base", + cuvs_java_jar="/artifacts/cuvs-java.jar", + cuvs_lucene_jar="/artifacts/cuvs-lucene.jar", + java_library_path="/native", + jvm_args=["-Xms1g"], + ) + + assert dataset_config.distance == "euclidean" + assert len(configs) == 2 + assert {config.indexes[0].build_param["codec"] for config in configs} == { + _HNSW_CODEC, + _CAGRA_CODEC, + } + for config in configs: + assert config.indexes[0].search_params == [{}] + assert config.backend_config["cuvs_java_jar"].endswith("cuvs-java.jar") + assert config.backend_config["cuvs_lucene_jar"].endswith( + "cuvs-lucene.jar" + ) + assert config.backend_config["java_library_path"] == "/native" + assert config.backend_config["jvm_args"] == ["-Xms1g"] + codec = config.indexes[0].build_param["codec"] + assert config.backend_config["requires_gpu"] is (codec == _CAGRA_CODEC) + assert config.backend_config["group"] == "base" + assert config.backend_config["index_name"] == config.index_name + assert config.backend_config["index_root"] == str( + Path("/datasets/test-dataset/index") + ) + assert config.backend_config["result_scope"] is None + + +def test_hnsw_build_defaults_are_canonical_and_share_one_identity(config_dir): + _, configs = PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + ) + + assert len(configs) == 1 + config = configs[0] + assert config.indexes[0].build_param == _DEFAULT_HNSW_BUILD_PARAMETERS + assert config.index_name == _hnsw_index_name("pylucene_test", "test") + assert PyLuceneConfigLoader._index_label( + "pylucene_test", + "test", + None, + {"codec": _HNSW_CODEC}, + ) == PyLuceneConfigLoader._index_label( + "pylucene_test", + "test", + None, + _DEFAULT_HNSW_BUILD_PARAMETERS, + ) + + +def test_config_loader_expands_hnsw_build_and_search_sweeps( + config_dir, tmp_path +): + algorithm_config = tmp_path / "pylucene_hnsw_sweep.yaml" + algorithm_config.write_text( + f"""\ +backend: pylucene +name: pylucene_hnsw_sweep +groups: + requested: + build: + codec: [{_HNSW_CODEC}] + m: [16, 24, 32] + search: + num_candidates: [150, 200, 300, 600] + build_grid: + build: + codec: [{_HNSW_CODEC}] + m: [16, 32] + ef_construction: [48, 64] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, requested_configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_hnsw_sweep", + groups="requested", + algorithm_configuration=str(algorithm_config), + ) + + assert len(requested_configs) == 3 + assert len({config.index_path for config in requested_configs}) == 3 + assert ( + sum( + len(config.indexes[0].search_params) + for config in requested_configs + ) + == 12 + ) + for config, m in zip(requested_configs, (16, 24, 32), strict=True): + assert config.indexes[0].build_param == { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": m, + } + assert config.index_name == _hnsw_index_name( + "pylucene_hnsw_sweep", "requested", m=m + ) + assert config.indexes[0].search_params == [ + {"num_candidates": 150}, + {"num_candidates": 200}, + {"num_candidates": 300}, + {"num_candidates": 600}, + ] + + _, build_grid_configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_hnsw_sweep", + groups="build_grid", + algorithm_configuration=str(algorithm_config), + ) + + assert len(build_grid_configs) == 4 + assert { + ( + config.indexes[0].build_param["m"], + config.indexes[0].build_param["ef_construction"], + ) + for config in build_grid_configs + } == {(16, 48), (16, 64), (32, 48), (32, 64)} + + +def test_pylucene_tune_space_uses_runtime_top_k_for_candidates(): + assert get_search_space("pylucene_cuvs_hnsw")["search"] == { + "num_candidates": { + "type": "int", + "min": "top_k", + "max": 500, + } + } + + +def test_pylucene_tune_resolves_candidate_minimum_from_top_k(monkeypatch): + suggested_ranges = _install_single_trial_optuna(monkeypatch) + orchestrator = BenchmarkOrchestrator.__new__(BenchmarkOrchestrator) + trial_arguments = {} + + def run_trial(**kwargs): + trial_arguments.update(kwargs) + return [ + SearchResult( + neighbors=np.empty((0, 0), dtype=np.int64), + distances=np.empty((0, 0), dtype=np.float32), + search_time_ms=1.0, + queries_per_second=1.0, + recall=1.0, + algorithm="pylucene_cuvs_hnsw", + search_params=[kwargs["search_params"]], + ) + ] + + orchestrator._run_trial = run_trial + results = orchestrator._run_tune( + constraints={"recall": "maximize"}, + n_trials=1, + build=True, + search=True, + force=False, + dry_run=False, + count=150, + batch_size=1, + search_mode="latency", + search_threads=1, + algorithms="pylucene_cuvs_hnsw", + ) + + assert suggested_ranges["num_candidates"] == (150, 500, False) + assert trial_arguments["search_params"] == {"num_candidates": 150} + assert results[0].search_params == [{"num_candidates": 150}] + + +def test_pylucene_tune_rejects_top_k_above_candidate_ceiling(monkeypatch): + _install_single_trial_optuna(monkeypatch) + orchestrator = BenchmarkOrchestrator.__new__(BenchmarkOrchestrator) + + with pytest.raises(ValueError, match="minimum 501 exceeds maximum 500"): + orchestrator._run_tune( + constraints={"recall": "maximize"}, + n_trials=1, + build=True, + search=True, + force=False, + dry_run=False, + count=501, + batch_size=1, + search_mode="latency", + search_threads=1, + algorithms="pylucene_cuvs_hnsw", + ) + + +@pytest.mark.parametrize( + ("field_name", "field_value"), + [ + ("m", True), + ("m", 0), + ("m", 513), + ("ef_construction", False), + ("ef_construction", 0), + ("ef_construction", 513), + ], +) +def test_hnsw_build_parameters_reject_booleans_and_values_outside_range( + field_name, field_value +): + with pytest.raises(ValueError, match=field_name): + pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC, field_name: field_value} + ) + + +def test_hnsw_build_parameter_boundaries_are_supported(): + assert pylucene_backend._normalize_build_params( + { + "codec": _HNSW_CODEC, + "m": 1, + "ef_construction": 512, + } + ) == { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 1, + "ef_construction": 512, + } + + +@pytest.mark.parametrize("field_name", ["m", "ef_construction"]) +def test_hnsw_build_parameters_are_rejected_for_cagra(field_name): + with pytest.raises(ValueError, match=field_name): + pylucene_backend._normalize_build_params( + {"codec": _CAGRA_CODEC, field_name: 32} + ) + + +@pytest.mark.parametrize( + ("build_yaml", "search_yaml", "error"), + [ + ( + f"codec: [{_HNSW_CODEC}]\n ignored: [1]", + "{}", + "Unsupported PyLucene build parameter.*ignored", + ), + ( + f"codec: [{_HNSW_CODEC}]", + "ignored: [1]", + "Unsupported PyLucene search parameter.*ignored", + ), + ], +) +def test_config_loader_rejects_unsupported_parameters( + config_dir, tmp_path, build_yaml, search_yaml, error +): + custom_config = tmp_path / "unsupported-parameters.yaml" + custom_config.write_text( + f"""\ +backend: pylucene +name: pylucene_unsupported +groups: + test: + build: + {build_yaml} + search: + {search_yaml} +""" + ) + + with pytest.raises(ValueError, match=error): + PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_unsupported", + groups="test", + algorithm_configuration=str(custom_config), + ) + + +@pytest.mark.parametrize( + ("algorithm_name", "group_name", "match"), + [ + ("../../../escaped", "test", "PyLucene algorithm name"), + ("pylucene_safe", "../escaped", "PyLucene group name"), + ("pylucene,ambiguous", "test", "PyLucene algorithm name"), + ], +) +def test_config_loader_rejects_unsafe_artifact_identity( + config_dir, tmp_path, algorithm_name, group_name, match +): + custom_config = tmp_path / "unsafe-identity.yaml" + custom_config.write_text( + f"""\ +backend: pylucene +name: {algorithm_name} +groups: + {group_name}: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + + with pytest.raises(ValueError, match=match): + PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms=(None if "," in algorithm_name else algorithm_name), + groups=group_name, + algorithm_configuration=str(custom_config), + ) + + +def test_config_loader_scopes_artifact_identity_by_dataset_subset(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + identities = [] + + for subset_size, subset_scope in ( + (None, None), + (2, "subset2"), + (3, "subset3"), + ): + dataset_config, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + subset_size=subset_size, + count=1, + batch_size=2, + ) + + assert dataset_config.subset_size == subset_size + assert len(configs) == 1 + config = configs[0] + index_name = _hnsw_index_name( + "pylucene_test", "test", subset_scope=subset_scope + ) + assert config.index_name == index_name + assert config.index_path == ( + Path("/datasets/test-dataset/index") / index_name + ) + assert config.backend_config["group"] == "test" + assert config.backend_config["index_name"] == index_name + assert config.backend_config["result_scope"] == subset_scope + identities.append( + ( + config.index_name, + config.index_path, + config.backend_config["result_scope"], + ) + ) + + assert len(set(identities)) == 3 + + +@pytest.mark.parametrize( + ("first", "second"), + [ + (("foo.subset2", "base", None), ("foo", "base", "subset2")), + (("a_b", "c", None), ("a", "b_c", None)), + ], +) +def test_index_labels_are_injective(first, second): + build_params = _DEFAULT_HNSW_BUILD_PARAMETERS + + first_label = PyLuceneConfigLoader._index_label( + first[0], first[1], first[2], build_params + ) + second_label = PyLuceneConfigLoader._index_label( + second[0], second[1], second[2], build_params + ) + + assert first_label != second_label + + +@pytest.mark.parametrize("subset_size", [0, -1, True, "../other"]) +def test_config_loader_rejects_unsafe_subset_identity(config_dir, subset_size): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match="positive integer"): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + subset_size=subset_size, + ) + + +@pytest.mark.parametrize( + ("backend_declaration", "algorithm_name"), + [ + ("", "pylucene_custom"), + ("backend: pylucene\n", "custom_lucene"), + ], + ids=["parsed-name", "declared-backend"], +) +def test_config_loader_discovers_custom_filename_by_parsed_identity( + config_dir, + tmp_path, + backend_declaration, + algorithm_name, +): + custom_config = tmp_path / "arbitrary-custom-name.yaml" + custom_config.write_text( + f"""\ +{backend_declaration}name: {algorithm_name} +groups: + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms=algorithm_name, + groups="test", + algorithm_configuration=str(custom_config), + ) + + assert len(configs) == 1 + assert configs[0].indexes[0].algo == algorithm_name + assert configs[0].index_name == _hnsw_index_name(algorithm_name, "test") + + +def test_config_loader_excludes_explicit_other_backend(config_dir, tmp_path): + custom_config = tmp_path / "misleading-pylucene-name.yaml" + custom_config.write_text( + f"""\ +backend: cpp_gbench +name: pylucene_other_backend +groups: + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises( + ValueError, + match="Unknown PyLucene algorithm selector.*pylucene_other_backend", + ): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_other_backend", + groups="test", + algorithm_configuration=str(custom_config), + ) + + +def test_config_loader_honors_algorithm_and_group_filters(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + ) + assert len(configs) == 1 + assert configs[0].index_name == _hnsw_index_name("pylucene_test", "test") + + +def test_algorithm_config_discovery_is_deterministic(tmp_path): + config_path = tmp_path / "config" + bundled_algorithms = config_path / "algos" + bundled_algorithms.mkdir(parents=True) + bundled_zeta = bundled_algorithms / "zeta.yaml" + bundled_alpha = bundled_algorithms / "alpha.yml" + bundled_zeta.touch() + bundled_alpha.touch() + (bundled_algorithms / "ignored.txt").touch() + + custom_algorithms = tmp_path / "custom" + custom_algorithms.mkdir() + custom_zeta = custom_algorithms / "zeta.yml" + custom_alpha = custom_algorithms / "alpha.yaml" + custom_zeta.touch() + custom_alpha.touch() + + files = PyLuceneConfigLoader( + config_path=config_path + ).gather_algorithm_configs(config_path, str(custom_algorithms)) + + assert files == [ + str(bundled_alpha), + str(bundled_zeta), + str(custom_alpha), + str(custom_zeta), + ] + + +def test_duplicate_algorithm_config_uses_last_definition_and_position( + config_dir, monkeypatch +): + loader = PyLuceneConfigLoader(config_path=config_dir) + config_files = ["alpha-original", "beta", "alpha-override"] + configs_by_file = { + "alpha-original": { + "name": "pylucene_alpha", + "groups": {"original": {}}, + }, + "beta": { + "name": "pylucene_beta", + "groups": {"base": {}}, + }, + "alpha-override": { + "name": "pylucene_alpha", + "groups": {"override": {}}, + }, + } + monkeypatch.setattr(loader, "load_yaml_file", configs_by_file.__getitem__) + + algorithm_configs = loader._load_algorithm_configs(config_files) + + assert list(algorithm_configs) == ["pylucene_beta", "pylucene_alpha"] + assert algorithm_configs["pylucene_alpha"] == {"override": {}} + + +def test_config_loader_unions_global_and_algorithm_specific_groups(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="base", + algo_groups="pylucene_test.test", + ) + + assert len(configs) == 3 + assert ( + sum( + config.index_name.startswith("pylucene_test[group=test]") + for config in configs + ) + == 1 + ) + + +def test_config_loader_selects_only_explicit_algorithm_group(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algo_groups="pylucene_test.test", + ) + + assert len(configs) == 1 + assert configs[0].index_name.startswith("pylucene_test[group=test]") + + +@pytest.mark.parametrize( + ("selectors", "expected_groups"), + [ + ( + {}, + ["pylucene_alpha,base", "pylucene_beta,base"], + ), + ( + {"algorithms": "pylucene_beta,pylucene_alpha"}, + ["pylucene_alpha,base", "pylucene_beta,base"], + ), + ( + {"groups": "shared"}, + ["pylucene_alpha,shared", "pylucene_beta,shared"], + ), + ( + { + "algorithms": "pylucene_alpha,pylucene_beta", + "groups": "alpha_only,shared", + }, + [ + "pylucene_alpha,shared", + "pylucene_alpha,alpha_only", + "pylucene_beta,shared", + ], + ), + ( + {"algo_groups": "pylucene_beta.beta_only"}, + ["pylucene_beta,beta_only"], + ), + ( + { + "algorithms": "pylucene_alpha", + "algo_groups": "pylucene_beta.beta_only", + }, + ["pylucene_alpha,base", "pylucene_beta,beta_only"], + ), + ( + { + "groups": "alpha_only", + "algo_groups": "pylucene_beta.beta_only", + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,beta_only"], + ), + ( + { + "algorithms": "pylucene_beta", + "groups": "shared", + "algo_groups": "pylucene_alpha.alpha_only", + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,shared"], + ), + ( + { + "algorithms": "pylucene_alpha", + "groups": "shared", + "algo_groups": "pylucene_alpha.shared", + }, + ["pylucene_alpha,shared"], + ), + ( + { + "algo_groups": ( + "pylucene_beta.beta_only,pylucene_alpha.alpha_only" + ) + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,beta_only"], + ), + ( + { + "algorithms": " pylucene_beta, pylucene_beta ", + "groups": " base,base ", + }, + ["pylucene_beta,base"], + ), + ], + ids=[ + "defaults", + "algorithm", + "group", + "algorithm-and-group", + "explicit-pair", + "algorithm-plus-explicit-pair", + "group-plus-explicit-pair", + "all-selectors", + "overlapping-selectors", + "explicit-pair-order-follows-config", + "duplicate-selectors", + ], +) +def test_config_loader_combines_global_and_explicit_selectors( + config_dir, monkeypatch, selectors, expected_groups +): + alpha_config = config_dir / "algos" / "pylucene_alpha.yaml" + alpha_config.write_text( + f"""\ +name: pylucene_alpha +groups: + base: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + shared: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + alpha_only: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + beta_config = config_dir / "algos" / "pylucene_beta.yaml" + beta_config.write_text( + f"""\ +name: pylucene_beta +groups: + base: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + shared: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + beta_only: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + monkeypatch.setattr( + loader, + "gather_algorithm_configs", + lambda *_args: [str(alpha_config), str(beta_config)], + ) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + **selectors, + ) + + selected_groups = [ + f"{config.algo},{config.backend_config['group']}" for config in configs + ] + assert selected_groups == expected_groups + + +def test_algorithm_selection_resolution_preserves_inputs_and_source_order(): + alpha_base = {"build": {"codec": [_HNSW_CODEC]}} + beta_base = {"build": {"codec": [_HNSW_CODEC]}} + algorithm_configs = { + "pylucene_beta": {"base": beta_base}, + "pylucene_alpha": {"base": alpha_base}, + } + original_items = [ + (algorithm, list(groups.items())) + for algorithm, groups in algorithm_configs.items() + ] + selection = pylucene_backend._AlgorithmSelection( + algorithms=frozenset({"pylucene_alpha", "pylucene_beta"}), + groups=None, + explicit_groups=frozenset(), + ) + + selected = selection.resolve(algorithm_configs) + + assert [(group.algorithm, group.group) for group in selected] == [ + ("pylucene_beta", "base"), + ("pylucene_alpha", "base"), + ] + assert selected[0].configuration is beta_base + assert selected[1].configuration is alpha_base + assert [ + (algorithm, list(groups.items())) + for algorithm, groups in algorithm_configs.items() + ] == original_items + + +def test_algorithm_selection_rejects_group_outside_selected_algorithms(): + selection = pylucene_backend._AlgorithmSelection( + algorithms=frozenset({"pylucene_alpha"}), + groups=frozenset({"beta_only"}), + explicit_groups=frozenset(), + ) + algorithm_configs = { + "pylucene_alpha": {"base": {}}, + "pylucene_beta": {"beta_only": {}}, + } + + with pytest.raises( + ValueError, + match="Unknown PyLucene group selector\\(s\\): beta_only", + ): + selection.resolve(algorithm_configs) + + +def test_algorithm_selection_reports_explicit_errors_deterministically(): + selection = pylucene_backend._AlgorithmSelection( + algorithms=None, + groups=None, + explicit_groups=frozenset( + { + ("pylucene_zeta", "base"), + ("pylucene_beta", "base"), + } + ), + ) + + with pytest.raises(ValueError) as error: + selection.resolve({"pylucene_alpha": {"base": {}}}) + + assert str(error.value) == ( + "Unknown PyLucene algorithm in --algo-groups: pylucene_beta" + ) + + +def test_config_loader_rejects_malformed_algorithm_group(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match="."): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algo_groups="pylucene_test", + ) + + +@pytest.mark.parametrize( + ("selectors", "error"), + [ + ( + {"algorithms": "not-pylucene", "groups": "base"}, + "Unknown PyLucene algorithm selector.*not-pylucene", + ), + ( + {"algorithms": "pylucene_test", "groups": "missing"}, + "Unknown PyLucene group selector.*missing", + ), + ( + { + "algorithms": "pylucene_test", + "groups": "base", + "algo_groups": "not-pylucene.test", + }, + "Unknown PyLucene algorithm in --algo-groups.*not-pylucene", + ), + ( + { + "algorithms": "pylucene_test", + "groups": "base", + "algo_groups": "pylucene_test.missing", + }, + "Unknown PyLucene group for pylucene_test.*missing", + ), + ], + ids=[ + "algorithm", + "global-group", + "algo-group-algorithm", + "algo-group-group", + ], +) +def test_config_loader_rejects_unknown_selectors(config_dir, selectors, error): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match=error): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + **selectors, + ) + + +@pytest.mark.parametrize( + ("algorithm", "expected_codecs"), + [ + ( + "pylucene_cuvs_hnsw", + {_HNSW_CODEC}, + ), + ("pylucene_cuvs_cagra", {_CAGRA_CODEC}), + ], +) +def test_shipped_algorithm_configs_load(algorithm, expected_codecs): + _, configs = PyLuceneConfigLoader().load( + dataset="test-data", + dataset_path="/datasets", + algorithms=algorithm, + groups="base", + ) + + assert { + config.indexes[0].build_param["codec"] for config in configs + } == expected_codecs diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_execution_paths.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_execution_paths.py new file mode 100644 index 0000000000..8c07040a07 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_execution_paths.py @@ -0,0 +1,1056 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""PyLucene end-to-end coverage for CPU HNSW and GPU cuVS search paths. + +The parametrized cases cover segment and force-merge topologies, CAGRA search +widths, persisted HNSW layer counts, deletions, filters, and brute-force recall. +GPU-required cases assert that cuVS ran and did not silently fall back to the CPU path. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from enum import Enum + +import pytest + +from cuvs_bench.tests.pylucene._pylucene_execution_path_utils import ( + CAGRA_BUILT_HNSW_BASE_LAYER_TEST_CODEC_CLASS, + CAGRA_BUILT_HNSW_THREE_LAYER_TEST_CODEC_CLASS, + CAGRA_TEST_CODEC_CLASS, + CAGRA_TEST_QUERY_CLASS, + CAGRA_VECTOR_READER_CLASS, + CPU_HNSW_TEST_CODEC_CLASS, + HNSW_GRAPH_VERIFYING_EXACT_QUERY_CLASS, + HNSW_GRAPH_VERIFYING_QUERY_CLASS, + IndexRun, + IndexScenario, + PyLuceneContext, + QueryDocumentState, + deterministic_float32_vector, + initialize_pylucene_context, + run_index_scenario, + segment_document_id_ranges, +) + +# Python 3.14 reports one deprecation per JCC-generated PyLucene builtin type. +# This narrow third-party filter keeps native/cuVS warnings fully visible. +pytestmark = [ + pytest.mark.pylucene, + pytest.mark.filterwarnings( + "ignore:builtin type .* has no __module__ attribute:DeprecationWarning" + ), +] + +# TODO(https://github.com/NVIDIA/cuvs/issues/2407): Add multithreaded concurrency coverage. + +HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +CAGRA_CODEC = "CuVS2510GPUSearchCodec" + +CPU_HNSW_WRITER_CLASS = ( + "org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter" +) +GPU_CAGRA_BUILT_HNSW_WRITER_CLASS = ( + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsWriter" +) +GPU_CAGRA_SEARCH_WRITER_CLASS = ( + "com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter" +) + +EXPECTED_SPI_CODECS = ( + HNSW_CODEC, + CAGRA_CODEC, + "Lucene101AcceleratedHNSWBinaryQuantizedCodec", + "Lucene101AcceleratedHNSWScalarQuantizedCodec", +) + +CAGRA_GRAPH_DEGREE = 32 +CAGRA_INTERMEDIATE_GRAPH_DEGREE = 64 +# cuVS expands NN-Descent's requested degree by 1.5 and reduces it with a +# warning when that internal degree is at least the number of input vectors. +NN_DESCENT_INTERNAL_GRAPH_DEGREE = CAGRA_INTERMEDIATE_GRAPH_DEGREE * 3 // 2 +MIN_VECTORS_PER_CAGRA_BUILD = NN_DESCENT_INTERNAL_GRAPH_DEGREE + 1 +CAGRA_HNSW_M = CAGRA_GRAPH_DEGREE // 2 +MIN_VECTORS_FOR_THREE_HNSW_LAYERS = ( + MIN_VECTORS_PER_CAGRA_BUILD * CAGRA_HNSW_M**2 +) +DEFAULT_MIN_RECALL = 0.75 + + +class ExecutionPath(Enum): + CPU_HNSW = ( + "CPU HNSW build -> HNSW search", + "cpu-hnsw", + CPU_HNSW_WRITER_CLASS, + ) + GPU_CAGRA_BUILT_HNSW = ( + "GPU CAGRA build -> HNSW search", + "gpu-cagra-built-hnsw", + GPU_CAGRA_BUILT_HNSW_WRITER_CLASS, + ) + GPU_CAGRA_SEARCH = ( + "GPU CAGRA build -> CAGRA search", + "gpu-cagra-search", + GPU_CAGRA_SEARCH_WRITER_CLASS, + ) + + def __init__( + self, label: str, selector: str, expected_writer_class: str + ) -> None: + self.label = label + self.selector = selector + self.expected_writer_class = expected_writer_class + + @property + def requires_gpu(self) -> bool: + return self is not ExecutionPath.CPU_HNSW + + +EXACT_FILTER_PATH_LABELS = { + ExecutionPath.CPU_HNSW: "CPU HNSW build -> Lucene exact filtered search", + ExecutionPath.GPU_CAGRA_BUILT_HNSW: ( + "GPU CAGRA build -> Lucene exact filtered search" + ), +} + + +class DocumentSetup(Enum): + ALL_SEARCHABLE = "all-searchable" + ONE_DELETED = "one-deleted" + SINGLE_LIVE = "single-live" + + +@dataclass(frozen=True) +class SuiteSettings: + requested_document_count: int + dimensions: int + top_k: int + minimum_recall: float + + +@dataclass(frozen=True) +class DocumentConfiguration: + document_ids_without_vectors: frozenset[int] = frozenset() + document_ids_to_delete: frozenset[int] = frozenset() + additional_query_document_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class DocumentFilterConfiguration: + query_document_id: int | None = None + accepted_document_ids: frozenset[int] = frozenset() + + +@dataclass(frozen=True) +class EndToEndCase: + selector: str + scenario: IndexScenario + execution_path: ExecutionPath + expected_index_file_suffixes: tuple[str, ...] + expected_hnsw_layers: int = 0 + min_recall: float = DEFAULT_MIN_RECALL + + +def _positive_int_env(name: str, default: int) -> int: + value = int(os.environ.get(name, default)) + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + return value + + +def _minimum_recall_from_env() -> float: + value = float( + os.environ.get("CUVS_BENCH_PYLUCENE_MIN_RECALL", DEFAULT_MIN_RECALL) + ) + if not 0.0 <= value <= 1.0: + raise ValueError( + "CUVS_BENCH_PYLUCENE_MIN_RECALL must be between 0 and 1, " + f"got {value}" + ) + return value + + +def _suite_settings() -> SuiteSettings: + requested_document_count = _positive_int_env( + "CUVS_BENCH_PYLUCENE_ROWS", 2000 + ) + dimensions = _positive_int_env("CUVS_BENCH_PYLUCENE_DIMS", 32) + top_k = _positive_int_env("CUVS_BENCH_PYLUCENE_TOPK", 20) + if dimensions > 4096: + raise ValueError( + f"CUVS_BENCH_PYLUCENE_DIMS must be at most 4096, got {dimensions}" + ) + return SuiteSettings( + requested_document_count=requested_document_count, + dimensions=dimensions, + top_k=top_k, + minimum_recall=_minimum_recall_from_env(), + ) + + +def _document_configuration( + document_count: int, setup: DocumentSetup +) -> DocumentConfiguration: + middle_document_id = document_count // 2 + if setup is DocumentSetup.SINGLE_LIVE: + deleted_document_ids = frozenset( + document_id + for document_id in range(document_count) + if document_id != middle_document_id + ) + return DocumentConfiguration( + document_ids_to_delete=deleted_document_ids, + additional_query_document_ids=(0,), + ) + if setup is DocumentSetup.ONE_DELETED: + return DocumentConfiguration( + document_ids_to_delete=frozenset({middle_document_id}), + additional_query_document_ids=(middle_document_id,), + ) + return DocumentConfiguration() + + +def _minimum_document_count_for_selective_filter( + segment_count: int, top_k: int +) -> int: + return segment_count * 4 * (top_k + 1) + + +def _selective_document_filter_configuration( + document_count: int, + segment_count: int, + top_k: int, + *, + use_lucene_exact_search: bool = False, +) -> DocumentFilterConfiguration: + accepted_document_ids = set() + for segment_document_ids in segment_document_id_ranges( + document_count, segment_count + ): + # Lucene 10.2 selects exact vector scoring up front when the accepted + # document count is no greater than k. + accepted_document_count = ( + top_k + if use_lucene_exact_search + else max(top_k + 1, len(segment_document_ids) // 4) + ) + first_accepted_document_id = ( + segment_document_ids.stop - accepted_document_count + ) + accepted_document_ids.update( + range( + first_accepted_document_id, + segment_document_ids.stop, + ) + ) + + query_document_id = 0 + if query_document_id in accepted_document_ids: + raise ValueError("selective filter must reject its query document") + return DocumentFilterConfiguration( + query_document_id=query_document_id, + accepted_document_ids=frozenset(accepted_document_ids), + ) + + +def _minimum_document_count( + execution_path: ExecutionPath, + segment_count: int, + hnsw_layers: int, + document_setup: DocumentSetup, +) -> int: + if execution_path is ExecutionPath.CPU_HNSW: + required_document_count = segment_count + elif hnsw_layers == 3: + required_document_count = MIN_VECTORS_FOR_THREE_HNSW_LAYERS + else: + required_document_count = segment_count * MIN_VECTORS_PER_CAGRA_BUILD + + if document_setup is DocumentSetup.ONE_DELETED: + required_document_count = max(required_document_count, 2) + return required_document_count + + +def _cpu_hnsw_case( + selector: str, + *, + segment_count: int = 1, + force_merge_segment_count: int = 0, + document_setup: DocumentSetup = DocumentSetup.ALL_SEARCHABLE, + selective_filter: bool = False, +) -> EndToEndCase: + settings = _suite_settings() + minimum_document_count = _minimum_document_count( + ExecutionPath.CPU_HNSW, + segment_count, + 0, + document_setup, + ) + if selective_filter: + minimum_document_count = max( + minimum_document_count, + _minimum_document_count_for_selective_filter( + segment_count, settings.top_k + ), + ) + document_count = max( + settings.requested_document_count, minimum_document_count + ) + documents = _document_configuration(document_count, document_setup) + document_filter = ( + _selective_document_filter_configuration( + document_count, + segment_count, + settings.top_k, + use_lucene_exact_search=True, + ) + if selective_filter + else DocumentFilterConfiguration() + ) + scenario = IndexScenario( + name=selector, + codec_name=HNSW_CODEC, + codec_factory_class=CPU_HNSW_TEST_CODEC_CLASS, + document_count=document_count, + dimensions=settings.dimensions, + top_k=settings.top_k, + segment_count=segment_count, + force_merge_segment_count=force_merge_segment_count, + document_ids_without_vectors=documents.document_ids_without_vectors, + document_ids_to_delete=documents.document_ids_to_delete, + additional_query_document_ids=documents.additional_query_document_ids, + document_ids_accepted_by_filter=document_filter.accepted_document_ids, + filter_query_document_id=document_filter.query_document_id, + use_exact_hnsw_filter_query=selective_filter, + expected_hnsw_m=32, + ) + return EndToEndCase( + selector=selector, + scenario=scenario, + execution_path=ExecutionPath.CPU_HNSW, + expected_index_file_suffixes=(".vex", ".vem"), + min_recall=settings.minimum_recall, + ) + + +def _cagra_built_hnsw_case( + selector: str, + *, + hnsw_layers: int = 1, + segment_count: int = 1, + force_merge_segment_count: int = 0, + document_setup: DocumentSetup = DocumentSetup.ALL_SEARCHABLE, + selective_filter: bool = False, +) -> EndToEndCase: + settings = _suite_settings() + minimum_document_count = _minimum_document_count( + ExecutionPath.GPU_CAGRA_BUILT_HNSW, + segment_count, + hnsw_layers, + document_setup, + ) + if selective_filter: + minimum_document_count = max( + minimum_document_count, + _minimum_document_count_for_selective_filter( + segment_count, settings.top_k + ), + ) + document_count = max( + settings.requested_document_count, minimum_document_count + ) + documents = _document_configuration(document_count, document_setup) + document_filter = ( + _selective_document_filter_configuration( + document_count, + segment_count, + settings.top_k, + use_lucene_exact_search=True, + ) + if selective_filter + else DocumentFilterConfiguration() + ) + codec_factory_class = ( + CAGRA_BUILT_HNSW_THREE_LAYER_TEST_CODEC_CLASS + if hnsw_layers == 3 + else CAGRA_BUILT_HNSW_BASE_LAYER_TEST_CODEC_CLASS + ) + scenario = IndexScenario( + name=selector, + codec_name=HNSW_CODEC, + codec_factory_class=codec_factory_class, + document_count=document_count, + dimensions=settings.dimensions, + top_k=settings.top_k, + segment_count=segment_count, + force_merge_segment_count=force_merge_segment_count, + document_ids_without_vectors=documents.document_ids_without_vectors, + document_ids_to_delete=documents.document_ids_to_delete, + additional_query_document_ids=documents.additional_query_document_ids, + document_ids_accepted_by_filter=document_filter.accepted_document_ids, + filter_query_document_id=document_filter.query_document_id, + use_exact_hnsw_filter_query=selective_filter, + expected_hnsw_m=16, + ) + return EndToEndCase( + selector=selector, + scenario=scenario, + execution_path=ExecutionPath.GPU_CAGRA_BUILT_HNSW, + expected_index_file_suffixes=(".vex", ".vem"), + expected_hnsw_layers=hnsw_layers, + min_recall=settings.minimum_recall, + ) + + +def _cagra_search_case( + selector: str, + *, + segment_count: int = 1, + force_merge_segment_count: int = 0, + document_setup: DocumentSetup = DocumentSetup.ALL_SEARCHABLE, + search_width: int = 1, + selective_filter: bool = False, +) -> EndToEndCase: + settings = _suite_settings() + if settings.top_k > 1024: + raise ValueError( + "GPU CAGRA search cases require CUVS_BENCH_PYLUCENE_TOPK <= 1024" + ) + minimum_document_count = _minimum_document_count( + ExecutionPath.GPU_CAGRA_SEARCH, + segment_count, + 0, + document_setup, + ) + if selective_filter: + minimum_document_count = max( + minimum_document_count, + _minimum_document_count_for_selective_filter( + segment_count, settings.top_k + ), + ) + document_count = max( + settings.requested_document_count, minimum_document_count + ) + documents = _document_configuration(document_count, document_setup) + document_filter = ( + _selective_document_filter_configuration( + document_count, segment_count, settings.top_k + ) + if selective_filter + else DocumentFilterConfiguration() + ) + scenario = IndexScenario( + name=selector, + codec_name=CAGRA_CODEC, + codec_factory_class=CAGRA_TEST_CODEC_CLASS, + document_count=document_count, + dimensions=settings.dimensions, + top_k=settings.top_k, + segment_count=segment_count, + force_merge_segment_count=force_merge_segment_count, + document_ids_without_vectors=documents.document_ids_without_vectors, + document_ids_to_delete=documents.document_ids_to_delete, + additional_query_document_ids=documents.additional_query_document_ids, + document_ids_accepted_by_filter=document_filter.accepted_document_ids, + filter_query_document_id=document_filter.query_document_id, + use_cagra_search_query=True, + search_width=search_width, + i_top_k=max(64, settings.top_k), + ) + return EndToEndCase( + selector=selector, + scenario=scenario, + execution_path=ExecutionPath.GPU_CAGRA_SEARCH, + expected_index_file_suffixes=(".vcag", ".vemc"), + min_recall=settings.minimum_recall, + ) + + +SEGMENT_CASES = ( + _cpu_hnsw_case( + "cpu-hnsw-1-segment", + ), + _cpu_hnsw_case( + "cpu-hnsw-10-segments", + segment_count=10, + ), + _cagra_built_hnsw_case( + "gpu-cagra-built-hnsw-1-segment", + ), + _cagra_built_hnsw_case( + "gpu-cagra-built-hnsw-10-segments", + segment_count=10, + ), + _cagra_search_case( + "gpu-cagra-search-10-segments", + segment_count=10, + ), +) + +SINGLE_LIVE_DOCUMENT_CASES = ( + _cagra_search_case( + "gpu-cagra-search-single-live-doc", + document_setup=DocumentSetup.SINGLE_LIVE, + ), +) + +FORCE_MERGE_CASES = ( + _cpu_hnsw_case( + "cpu-hnsw-10-to-1-force-merge", + segment_count=10, + force_merge_segment_count=1, + ), + _cpu_hnsw_case( + "cpu-hnsw-100-to-10-force-merge", + segment_count=100, + force_merge_segment_count=10, + ), + _cagra_built_hnsw_case( + "gpu-cagra-built-hnsw-10-to-1-force-merge", + segment_count=10, + force_merge_segment_count=1, + ), + _cagra_built_hnsw_case( + "gpu-cagra-built-hnsw-100-to-10-force-merge", + segment_count=100, + force_merge_segment_count=10, + ), + _cagra_search_case( + "gpu-cagra-search-10-to-1-force-merge", + segment_count=10, + force_merge_segment_count=1, + ), + _cagra_search_case( + "gpu-cagra-search-100-to-10-force-merge", + segment_count=100, + force_merge_segment_count=10, + ), +) + +HNSW_LAYER_CASES = ( + _cagra_built_hnsw_case( + "gpu-cagra-built-hnsw-3-layers", + hnsw_layers=3, + ), +) + +CAGRA_SEARCH_WIDTH_CASES = ( + _cagra_search_case( + "gpu-cagra-search-1-segment", + ), + _cagra_search_case( + "gpu-cagra-search-width-16", + search_width=16, + ), + _cagra_search_case( + "gpu-cagra-search-width-32", + search_width=32, + ), +) + +DELETED_DOCUMENT_CASES = ( + _cagra_search_case( + "gpu-cagra-search-deleted-documents", + document_setup=DocumentSetup.ONE_DELETED, + ), +) + +DOCUMENT_FILTER_CASES = ( + _cpu_hnsw_case( + "cpu-hnsw-lucene-exact-filter", + selective_filter=True, + ), + _cagra_built_hnsw_case( + "gpu-cagra-built-hnsw-lucene-exact-filter", + selective_filter=True, + ), + _cagra_search_case( + "gpu-cagra-search-selective-filter-10-segments", + segment_count=10, + selective_filter=True, + ), +) + + +def _case_parameter(case: EndToEndCase) -> object: + smoke_cases = { + "cpu-hnsw-1-segment", + "gpu-cagra-built-hnsw-1-segment", + "gpu-cagra-search-1-segment", + } + marks = ( + () + if case.selector in smoke_cases + else (pytest.mark.pylucene_extended,) + ) + return pytest.param(case, id=case.selector, marks=marks) + + +def _case_parameters( + cases: tuple[EndToEndCase, ...], +) -> tuple[object, ...]: + return tuple(_case_parameter(case) for case in cases) + + +@pytest.fixture(scope="session") +def pylucene_context(pylucene_runtime_config) -> PyLuceneContext: + return initialize_pylucene_context( + pylucene_runtime_config.backend_config, + EXPECTED_SPI_CODECS, + ) + + +def _squared_distance( + left: tuple[float, ...], right: tuple[float, ...] +) -> float: + return sum( + (left_value - right_value) ** 2 + for left_value, right_value in zip(left, right) + ) + + +def _brute_force_neighbor_ids( + case: EndToEndCase, + result: IndexRun, + query_document_id: int, + candidate_document_ids: tuple[int, ...] | None = None, +) -> tuple[str, ...]: + query_vector = deterministic_float32_vector( + query_document_id, case.scenario.dimensions + ) + candidates = ( + result.searchable_vector_document_ids + if candidate_document_ids is None + else candidate_document_ids + ) + ordered_ids = sorted( + candidates, + key=lambda document_id: ( + _squared_distance( + query_vector, + deterministic_float32_vector( + document_id, case.scenario.dimensions + ), + ), + document_id, + ), + ) + return tuple( + f"doc-{document_id}" + for document_id in ordered_ids[: case.scenario.top_k] + ) + + +def _assert_index_files_and_segments( + case: EndToEndCase, result: IndexRun +) -> None: + for suffix in case.expected_index_file_suffixes: + assert any(name.endswith(suffix) for name in result.index_files), ( + f"{case.selector}: no index file ending with {suffix}; " + f"files={result.index_files}" + ) + + expected_initial_segments = min( + case.scenario.segment_count, case.scenario.document_count + ) + assert result.pre_merge_segment_count == expected_initial_segments, ( + f"{case.selector}: expected {expected_initial_segments} segments " + f"before forceMerge, got {result.pre_merge_segment_count}" + ) + expected_final_segments = ( + case.scenario.force_merge_segment_count or expected_initial_segments + ) + assert result.segment_count == expected_final_segments, ( + f"{case.selector}: expected {expected_final_segments} final segments, " + f"got {result.segment_count}" + ) + + +def _assert_index_metadata(case: EndToEndCase, result: IndexRun) -> None: + expected_live_documents = case.scenario.document_count - len( + result.deleted_document_ids + ) + assert result.live_document_count == expected_live_documents + expected_max_documents = ( + expected_live_documents + if case.scenario.force_merge_segment_count + else case.scenario.document_count + ) + assert result.max_document_count == expected_max_documents + + expected_vector_count = ( + len(result.searchable_vector_document_ids) + if case.scenario.force_merge_segment_count + else ( + case.scenario.document_count + - len(result.document_ids_without_vectors) + ) + ) + assert result.vector_count == expected_vector_count, ( + f"{case.selector}: expected {expected_vector_count} vector values, " + f"got {result.vector_count}" + ) + assert result.vector_dimensions + assert set(result.vector_dimensions) == {case.scenario.dimensions} + + +def _assert_execution_path(case: EndToEndCase, result: IndexRun) -> None: + telemetry = result.writer_telemetry + assert telemetry.get("configuredPath") == case.execution_path.selector + assert ( + telemetry.get("writerClass") + == case.execution_path.expected_writer_class + ) + + if case.execution_path is ExecutionPath.GPU_CAGRA_SEARCH: + assert set(result.vector_reader_classes) == {CAGRA_VECTOR_READER_CLASS} + assert all( + observation.query_class == CAGRA_TEST_QUERY_CLASS + for observation in result.query_observations + ) + if result.filtered_query_observation is not None: + assert ( + result.filtered_query_observation.query_class + == CAGRA_TEST_QUERY_CLASS + ) + assert not result.hnsw_layer_counts + return + + assert result.vector_reader_classes + assert all( + observation.query_class == HNSW_GRAPH_VERIFYING_QUERY_CLASS + for observation in result.query_observations + ) + if result.filtered_query_observation is not None: + expected_filtered_query_class = ( + HNSW_GRAPH_VERIFYING_EXACT_QUERY_CLASS + if case.scenario.use_exact_hnsw_filter_query + else HNSW_GRAPH_VERIFYING_QUERY_CLASS + ) + assert ( + result.filtered_query_observation.query_class + == expected_filtered_query_class + ) + assert all( + "Lucene99HnswVectorsReader" in reader_class + for reader_class in result.vector_reader_classes + ) + assert all( + not reader_class.startswith("com.nvidia.cuvs.lucene") + for reader_class in result.vector_reader_classes + ) + + +def _assert_graph_configuration(case: EndToEndCase, result: IndexRun) -> None: + if case.expected_hnsw_layers: + assert result.hnsw_layer_counts + assert set(result.hnsw_layer_counts) == {case.expected_hnsw_layers}, ( + f"{case.selector}: expected {case.expected_hnsw_layers} HNSW " + f"layers, got {result.hnsw_layer_counts}" + ) + + if not case.execution_path.requires_gpu: + return + + telemetry = result.writer_telemetry + assert telemetry.get("cagraGraphBuildAlgo") == "NN_DESCENT" + assert telemetry.get("cagraGraphDegree") == "32" + assert telemetry.get("cagraIntermediateGraphDegree") == "64" + if case.execution_path is ExecutionPath.GPU_CAGRA_SEARCH: + assert telemetry.get("cagraStrategy") == "CUSTOM" + + +def _assert_search_results( + case: EndToEndCase, result: IndexRun +) -> tuple[float, ...]: + documents_without_vectors = { + f"doc-{document_id}" + for document_id in result.document_ids_without_vectors + } + deleted_documents = { + f"doc-{document_id}" for document_id in result.deleted_document_ids + } + expected_hit_count = min( + case.scenario.top_k, len(result.searchable_vector_document_ids) + ) + + recalls = [] + for observation in result.query_observations: + hits = observation.hit_ids + queried_document = f"doc-{observation.query_document_id}" + if observation.query_document_state is QueryDocumentState.SEARCHABLE: + assert len(hits) == expected_hit_count, ( + f"{case.selector}: query {queried_document} expected " + f"{expected_hit_count} hits, got {hits}" + ) + assert hits[0] == queried_document, ( + f"{case.selector}: queried document must rank first; " + f"expected {queried_document}, got {hits}" + ) + else: + assert queried_document not in hits, ( + f"{case.selector}: {observation.query_document_state.value} " + f"document was returned: {hits}" + ) + + assert len(hits) == len(set(hits)), ( + f"{case.selector}: duplicate hits returned: {hits}" + ) + returned_vectorless_documents = tuple( + document + for document in hits + if document in documents_without_vectors + ) + assert not returned_vectorless_documents, ( + f"{case.selector}: live documents without vectors were returned: " + f"{returned_vectorless_documents}" + ) + returned_deleted_documents = tuple( + document for document in hits if document in deleted_documents + ) + assert not returned_deleted_documents, ( + f"{case.selector}: deleted documents were returned: " + f"{returned_deleted_documents}" + ) + + if ( + observation.query_document_state + is not QueryDocumentState.SEARCHABLE + ): + continue + + expected_neighbors = _brute_force_neighbor_ids( + case, result, observation.query_document_id + ) + recall = len(set(hits) & set(expected_neighbors)) / len( + expected_neighbors + ) + assert recall >= case.min_recall, ( + f"{case.selector}: query {queried_document} recall " + f"{recall:.3f} is below floor {case.min_recall:.3f}; " + f"expected={expected_neighbors}, actual={hits}" + ) + recalls.append(recall) + + return tuple(recalls) + + +def _execution_path_label(case: EndToEndCase) -> str: + return ( + EXACT_FILTER_PATH_LABELS[case.execution_path] + if case.scenario.use_exact_hnsw_filter_query + else case.execution_path.label + ) + + +def _print_result( + case: EndToEndCase, result: IndexRun, recalls: tuple[float, ...] +) -> None: + report_label = case.selector.removeprefix( + f"{case.execution_path.selector}-" + ) + segment_summary = ( + f"{result.pre_merge_segment_count}->{result.segment_count}" + if case.scenario.force_merge_segment_count + else str(result.segment_count) + ) + details = [ + f"documents={case.scenario.document_count}", + f"liveDocuments={result.live_document_count}", + f"segments={segment_summary}", + f"topK={case.scenario.top_k}", + f"recall(min/avg)={min(recalls):.3f}/" + f"{sum(recalls) / len(recalls):.3f}", + ] + if result.document_ids_without_vectors: + details.append( + f"documentsWithoutVectors={len(result.document_ids_without_vectors)}" + ) + if result.deleted_document_ids: + details.append(f"deletedDocuments={len(result.deleted_document_ids)}") + if case.execution_path is ExecutionPath.GPU_CAGRA_SEARCH: + details.append(f"searchWidth={case.scenario.search_width}") + else: + details.append(f"hnswLayers={sorted(set(result.hnsw_layer_counts))}") + details.append(f"hnswM={case.scenario.expected_hnsw_m}") + + print( + f"PASS [{_execution_path_label(case)}] {report_label}: " + + ", ".join(details) + ) + + +def _run_and_verify( + case: EndToEndCase, context: PyLuceneContext +) -> tuple[IndexRun, tuple[float, ...]]: + result = run_index_scenario(case.scenario, context) + _assert_index_files_and_segments(case, result) + _assert_index_metadata(case, result) + _assert_execution_path(case, result) + _assert_graph_configuration(case, result) + recalls = _assert_search_results(case, result) + _print_result(case, result, recalls) + return result, recalls + + +@pytest.mark.parametrize("case", _case_parameters(SEGMENT_CASES)) +def test_search_with_configured_segment_count( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Exercise each execution path across configured segment topologies.""" + result, _ = _run_and_verify(case, pylucene_context) + assert not result.document_ids_without_vectors + assert not result.deleted_document_ids + + +@pytest.mark.parametrize("case", _case_parameters(SINGLE_LIVE_DOCUMENT_CASES)) +def test_cagra_search_with_single_live_document( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Search a warning-free CAGRA index after deleting all but one document.""" + result, _ = _run_and_verify(case, pylucene_context) + assert result.live_document_count == 1 + assert len(result.searchable_vector_document_ids) == 1 + assert len(result.deleted_document_ids) == result.max_document_count - 1 + assert any( + observation.query_document_state is QueryDocumentState.DELETED + for observation in result.query_observations + ) + + +@pytest.mark.parametrize("case", _case_parameters(FORCE_MERGE_CASES)) +def test_search_after_force_merge( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Verify CPU and GPU search paths after representative force merges.""" + result, _ = _run_and_verify(case, pylucene_context) + assert not result.document_ids_without_vectors + assert not result.deleted_document_ids + + +@pytest.mark.parametrize("case", _case_parameters(HNSW_LAYER_CASES)) +def test_cagra_built_hnsw_has_expected_layer_count( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Verify CAGRA-built HNSW persists the requested three layers.""" + result, _ = _run_and_verify(case, pylucene_context) + assert set(result.hnsw_layer_counts) == {case.expected_hnsw_layers} + + +@pytest.mark.parametrize("case", _case_parameters(CAGRA_SEARCH_WIDTH_CASES)) +def test_cagra_search_with_configured_search_width( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Exercise GPU CAGRA search widths 1, 16, and 32.""" + _run_and_verify(case, pylucene_context) + + +@pytest.mark.parametrize("case", _case_parameters(DELETED_DOCUMENT_CASES)) +def test_deleted_documents_are_not_searchable( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Verify GPU CAGRA search never returns a deleted vector document.""" + result, _ = _run_and_verify(case, pylucene_context) + assert len(result.deleted_document_ids) == 1 + assert not result.document_ids_without_vectors + assert any( + observation.query_document_state is QueryDocumentState.DELETED + for observation in result.query_observations + ) + + +@pytest.mark.parametrize("case", _case_parameters(DOCUMENT_FILTER_CASES)) +def test_vector_search_honors_selective_document_filter( + pylucene_context: PyLuceneContext, case: EndToEndCase +) -> None: + """Compare exact-filter HNSW and filtered CAGRA results with brute force.""" + result, _ = _run_and_verify(case, pylucene_context) + observation = result.filtered_query_observation + assert observation is not None + + accepted_document_ids = tuple( + document_id + for document_id in result.searchable_vector_document_ids + if document_id in case.scenario.document_ids_accepted_by_filter + ) + accepted_document_id_set = set(accepted_document_ids) + accepted_hit_ids = { + f"doc-{document_id}" for document_id in accepted_document_ids + } + accepted_counts_by_segment = tuple( + sum( + document_id in accepted_document_id_set + for document_id in segment_document_ids + ) + for segment_document_ids in segment_document_id_ranges( + case.scenario.document_count, + case.scenario.segment_count, + ) + ) + query_document_id = case.scenario.filter_query_document_id + assert query_document_id is not None + queried_document = f"doc-{query_document_id}" + + assert observation.query_document_id == query_document_id + assert query_document_id not in ( + case.scenario.document_ids_accepted_by_filter + ) + assert queried_document not in observation.hit_ids + if case.scenario.use_exact_hnsw_filter_query: + assert set(accepted_counts_by_segment) == {case.scenario.top_k}, ( + f"{case.selector}: every segment must retain exactly topK=" + f"{case.scenario.top_k} accepted vectors to select Lucene exact " + f"scoring; acceptedPerSegment={accepted_counts_by_segment}" + ) + else: + assert min(accepted_counts_by_segment) > case.scenario.top_k, ( + f"{case.selector}: every segment must retain more than topK=" + f"{case.scenario.top_k} accepted vectors; " + f"acceptedPerSegment={accepted_counts_by_segment}" + ) + + rejected_hit_ids = tuple( + hit_id + for hit_id in observation.hit_ids + if hit_id not in accepted_hit_ids + ) + assert not rejected_hit_ids, ( + f"{case.selector}: filter-rejected documents were returned: " + f"{rejected_hit_ids}" + ) + expected_hit_count = min( + case.scenario.top_k, + len(accepted_document_ids), + ) + assert len(observation.hit_ids) == expected_hit_count, ( + f"{case.selector}: expected {expected_hit_count} filtered hits, " + f"got {observation.hit_ids}" + ) + assert len(observation.hit_ids) == len(set(observation.hit_ids)), ( + f"{case.selector}: duplicate filtered hits returned: " + f"{observation.hit_ids}" + ) + + expected_neighbors = _brute_force_neighbor_ids( + case, + result, + query_document_id, + accepted_document_ids, + ) + recall = len(set(observation.hit_ids) & set(expected_neighbors)) / len( + expected_neighbors + ) + assert recall >= case.min_recall, ( + f"{case.selector}: filtered query {queried_document} recall " + f"{recall:.3f} is below floor {case.min_recall:.3f}; " + f"expected={expected_neighbors}, actual={observation.hit_ids}" + ) + print( + f"FILTER [{_execution_path_label(case)}] {case.selector}: " + f"acceptedPerSegment={min(accepted_counts_by_segment)}-" + f"{max(accepted_counts_by_segment)}, " + f"recall={recall:.3f}" + ) diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_integration.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_integration.py new file mode 100644 index 0000000000..ee047d1e92 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_integration.py @@ -0,0 +1,1285 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Opt-in integration coverage for the PyLucene benchmark backend.""" + +import csv +import json +import os +import re +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytest + +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends._pylucene_java import ( + EF_CONSTRUCTION_PROPERTY, + M_PROPERTY, +) +from cuvs_bench.backends._utils import compute_recall +from cuvs_bench.backends.base import BuildResult, Dataset +from cuvs_bench.backends.pylucene import ( + _BuildCodec, + _CAGRA_PROVENANCE_FILE, + PyLuceneBackend, + _HNSW_PROVENANCE_FILE, + _PyLuceneRuntime, + _restore_java_property, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig +from cuvs_bench.tests.pylucene._pylucene_execution_path_utils import ( + HNSW_GRAPH_VERIFYING_QUERY_CLASS, + _new_hnsw_graph_verifying_query, + initialize_pylucene_context, +) +from cuvs_bench.tests.pylucene._pylucene_live_test_config import ( + PYLUCENE_TEST_CLASSES_ENV, + PYTHON_PACKAGE_ROOT, +) + +pytestmark = [ + pytest.mark.pylucene, + pytest.mark.filterwarnings( + "ignore:builtin type .* has no __module__ attribute:DeprecationWarning" + ), +] + +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_HNSW_WRITER_POLICY = "gpu-with-cpu-fallback" +_COMPOUND_FILE_POLICY = { + _HNSW_CODEC: "lucene-default", + _CAGRA_CODEC: "disabled", +} +_GPU_HNSW_WRITER = ( + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsWriter" +) +_CPU_HNSW_WRITER = ( + "org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter" +) +_WRITER_SELECTION_CODEC = "com.nvidia.cuvs.bench.PyLuceneWriterSelectionCodec" +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} +_CAGRA_BUILD_PARAMETERS = {"codec": _CAGRA_CODEC} +_CPU_FALLBACK_PROBE = Path(__file__).with_name( + "pylucene_cpu_fallback_probe.py" +) +_GRAPH_CLAMP_WARNING_FRAGMENTS = ( + "Intermediate graph degree cannot be larger", + "cannot be larger than intermediate graph degree", + "for nn-descent needs to match cagra intermediate graph degree", +) + + +@dataclass(frozen=True) +class _DatasetCase: + dataset: Dataset + query_ids: np.ndarray + squared_distances: np.ndarray + k: int + + +@dataclass(frozen=True) +class _BaselineIndex: + algo: str + codec: str + writer_policy: str + build_parameters: dict[str, object] + index_path: Path + index: IndexConfig + build_result: BuildResult + dataset_case: _DatasetCase + runtime_config: dict[str, str] + + +def _write_test_bin(path: Path, data: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as file: + write_bin_header(file, data.shape[0], data.shape[1]) + np.ascontiguousarray(data).tofile(file) + + +def _single_search_result(results): + assert len(results) == 1 + return results[0] + + +def _writer_diagnostics(java_codec): + diagnostics = str(java_codec.knnVectorsFormat()) + match = re.search( + r"writerClass=([^,)]+), fieldsWriterCalls=(\d+)", diagnostics + ) + assert match is not None, diagnostics + return match.group(1), int(match.group(2)) + + +def _configured_writer_selection_codec(runtime, m, ef_construction): + properties = { + M_PROPERTY: str(m), + EF_CONSTRUCTION_PROPERTY: str(ef_construction), + } + previous = {name: runtime.System.getProperty(name) for name in properties} + try: + for name, value in properties.items(): + runtime.System.setProperty(name, value) + reflected_codec = runtime.Class.forName( + _WRITER_SELECTION_CODEC + ).newInstance() + return runtime.Codec.cast_(reflected_codec) + finally: + for name, value in previous.items(): + _restore_java_property(runtime.System, name, value) + + +def _backend(algo, codec, runtime_config): + return PyLuceneBackend( + { + "name": f"{algo}-integration", + "algo": algo, + "codec": codec, + **runtime_config, + } + ) + + +def _index_at(baseline, index_path): + return IndexConfig( + name=baseline.index.name, + algo=baseline.algo, + build_param=baseline.build_parameters, + search_params=[{}], + file=str(index_path), + ) + + +def _copy_baseline(baseline, tmp_path): + index_path = tmp_path / f"{baseline.algo}-index" + shutil.copytree(baseline.index_path, index_path) + return index_path, _index_at(baseline, index_path) + + +def _search_copied_index(baseline, index, dataset=None): + backend = _backend(baseline.algo, baseline.codec, baseline.runtime_config) + search_dataset = ( + baseline.dataset_case.dataset if dataset is None else dataset + ) + try: + return _single_search_result( + backend.search( + search_dataset, + [index], + k=baseline.dataset_case.k, + ) + ) + finally: + backend.cleanup() + + +def _assert_cagra_verification( + metadata, expected_vector_count, expected_dimensions +): + verification = metadata["cagra_verification"] + assert verification["status"] == "cagra-only" + assert verification["segment_count"] >= 1 + assert verification["field_count"] >= 1 + assert verification["vector_count"] == expected_vector_count + assert verification["dimensions"] == expected_dimensions + + +def _assert_cagra_provenance( + metadata, + expected_vector_count, + expected_dimensions, +): + assert metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 4, + "codec": _CAGRA_CODEC, + "build_parameters": _CAGRA_BUILD_PARAMETERS, + "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "segment_count": metadata["segment_count"], + "commit_file_count": 1, + } + + +def _commit_one_deletion(backend, index_path): + runtime = backend._get_runtime() + runtime.attach_current_thread() + directory = runtime.FSDirectory.open(runtime.Paths.get(str(index_path))) + writer = None + reader = None + try: + writer_config = runtime.IndexWriterConfig() + writer_config.setOpenMode(runtime.IndexWriterConfig.OpenMode.APPEND) + writer = runtime.IndexWriter(directory, writer_config) + reader = runtime.DirectoryReader.open(writer) + assert int(writer.tryDeleteDocument(reader, 0)) >= 0 + writer.commit() + finally: + try: + if reader is not None: + reader.close() + finally: + try: + if writer is not None: + writer.close() + finally: + directory.close() + + +def _assert_hnsw_verification( + metadata, + codec, + expected_vector_count, + expected_dimensions, + expected_build_parameters, +): + verification = metadata["hnsw_verification"] + assert verification == { + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 4, + "codec": codec, + "build_parameters": expected_build_parameters, + "writer_policy": _HNSW_WRITER_POLICY, + "compound_file_policy": "lucene-default", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "segment_count": metadata["segment_count"], + "commit_file_count": 1, + } + + +@pytest.fixture(scope="module") +def integration_dataset_case(): + rng = np.random.default_rng(1907) + training_vectors = rng.standard_normal((512, 32)).astype(np.float32) + query_ids = np.asarray([0, 137, 259, 511], dtype=np.int64) + query_vectors = training_vectors[query_ids].copy() + k = 5 + + squared_distances = np.sum( + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, + axis=2, + ) + groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] + groundtruth_distances = np.take_along_axis( + squared_distances, groundtruth_neighbors, axis=1 + ) + return _DatasetCase( + dataset=Dataset( + name="pylucene-integration", + training_vectors=training_vectors, + query_vectors=query_vectors, + groundtruth_neighbors=groundtruth_neighbors.astype(np.int32), + groundtruth_distances=groundtruth_distances.astype(np.float32), + distance_metric="euclidean", + ), + query_ids=query_ids, + squared_distances=squared_distances, + k=k, + ) + + +def _build_baseline( + tmp_path_factory, + runtime_fixture, + dataset_case, + *, + algo, + codec, + writer_policy, +): + build_parameters = ( + _DEFAULT_HNSW_BUILD_PARAMETERS + if codec == _HNSW_CODEC + else _CAGRA_BUILD_PARAMETERS + ) + index_path = tmp_path_factory.mktemp(f"{algo}-baseline") / "index" + index = IndexConfig( + name=f"{algo}-integration", + algo=algo, + build_param=build_parameters, + search_params=[{}], + file=str(index_path), + ) + backend = _backend(algo, codec, runtime_fixture.backend_config) + try: + build_result = backend.build(dataset_case.dataset, [index], force=True) + finally: + backend.cleanup() + assert build_result.success, build_result.error_message + return _BaselineIndex( + algo=algo, + codec=codec, + writer_policy=writer_policy, + build_parameters=build_parameters, + index_path=index_path, + index=index, + build_result=build_result, + dataset_case=dataset_case, + runtime_config=runtime_fixture.backend_config, + ) + + +@pytest.fixture(scope="module") +def hnsw_baseline( + tmp_path_factory, pylucene_runtime_config, integration_dataset_case +): + return _build_baseline( + tmp_path_factory, + pylucene_runtime_config, + integration_dataset_case, + algo="pylucene_cuvs_hnsw", + codec=_HNSW_CODEC, + writer_policy=_HNSW_WRITER_POLICY, + ) + + +@pytest.fixture(scope="module") +def cagra_baseline( + tmp_path_factory, pylucene_runtime_config, integration_dataset_case +): + return _build_baseline( + tmp_path_factory, + pylucene_runtime_config, + integration_dataset_case, + algo="pylucene_cuvs_cagra", + codec=_CAGRA_CODEC, + writer_policy="gpu-cagra", + ) + + +def _assert_build_contract(baseline): + result = baseline.build_result + case = baseline.dataset_case + vectors = case.dataset.training_vectors + assert result.build_time_seconds > 0 + assert result.index_size_bytes > 0 + assert result.build_params == baseline.build_parameters + assert result.metadata["codec"] == baseline.codec + assert result.metadata["pylucene_version"] != "unknown" + assert result.metadata["writer_policy"] == baseline.writer_policy + assert result.metadata["segment_count"] >= 1 + assert ( + result.metadata["compound_file_policy"] + == _COMPOUND_FILE_POLICY[baseline.codec] + ) + if baseline.codec == _CAGRA_CODEC: + assert any( + path.suffix == ".vemc" for path in baseline.index_path.iterdir() + ) + assert any( + path.suffix == ".vcag" for path in baseline.index_path.iterdir() + ) + _assert_cagra_provenance( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + _assert_cagra_verification( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + assert (baseline.index_path / _CAGRA_PROVENANCE_FILE).is_file() + else: + _assert_hnsw_verification( + result.metadata, + baseline.codec, + vectors.shape[0], + vectors.shape[1], + baseline.build_parameters, + ) + assert (baseline.index_path / _HNSW_PROVENANCE_FILE).is_file() + + +def _assert_search_contract(baseline, result): + case = baseline.dataset_case + dataset = case.dataset + query_vectors = dataset.query_vectors + training_vectors = dataset.training_vectors + assert result.success, result.error_message + assert result.neighbors.shape == (query_vectors.shape[0], case.k) + assert result.distances.shape == (query_vectors.shape[0], case.k) + np.testing.assert_array_equal(result.neighbors[:, 0], case.query_ids) + assert np.all(np.isfinite(result.distances)) + recall = compute_recall( + result.neighbors, dataset.groundtruth_neighbors, case.k + ) + assert recall >= 0.75 + returned_distances = np.take_along_axis( + case.squared_distances, result.neighbors, axis=1 + ) + np.testing.assert_allclose( + result.distances, returned_distances, rtol=1e-5, atol=1e-5 + ) + assert np.all(np.diff(result.distances, axis=1) >= -1e-6) + assert result.search_time_ms > 0 + assert result.metadata["latency_seconds"] > 0 + assert result.queries_per_second > 0 + assert result.metadata["codec"] == baseline.codec + assert result.metadata["segment_count"] >= 1 + assert result.metadata["pylucene_version"] != "unknown" + assert ( + result.metadata["compound_file_policy"] + == _COMPOUND_FILE_POLICY[baseline.codec] + ) + assert result.metadata["num_batches"] == 2 + assert result.metadata["mode"] == "latency" + assert set(result.latency_percentiles) == {"p50", "p95", "p99"} + if baseline.codec == _CAGRA_CODEC: + assert result.search_params == [{}] + _assert_cagra_provenance( + result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + else: + assert result.search_params == [{"num_candidates": case.k}] + assert result.metadata["num_candidates"] == case.k + assert "cagra_verification" not in result.metadata + _assert_hnsw_verification( + result.metadata, + baseline.codec, + training_vectors.shape[0], + training_vectors.shape[1], + baseline.build_parameters, + ) + + +def test_build_hnsw_with_real_pylucene_runtime(hnsw_baseline): + _assert_build_contract(hnsw_baseline) + + +def test_build_cagra_with_real_pylucene_runtime(cagra_baseline): + _assert_build_contract(cagra_baseline) + + +def _assert_reuse_contract(baseline_index): + backend = _backend( + baseline_index.algo, + baseline_index.codec, + baseline_index.runtime_config, + ) + try: + result = backend.build( + baseline_index.dataset_case.dataset, + [baseline_index.index], + force=False, + ) + finally: + backend.cleanup() + assert result.success, result.error_message + assert result.build_params == baseline_index.build_parameters + assert result.metadata["skipped"] is True + assert result.metadata["segment_count"] >= 1 + assert ( + result.metadata["compound_file_policy"] + == _COMPOUND_FILE_POLICY[baseline_index.codec] + ) + vectors = baseline_index.dataset_case.dataset.training_vectors + if baseline_index.codec == _CAGRA_CODEC: + _assert_cagra_provenance( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + _assert_cagra_verification( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + else: + _assert_hnsw_verification( + result.metadata, + baseline_index.codec, + vectors.shape[0], + vectors.shape[1], + baseline_index.build_parameters, + ) + + +def test_reuse_real_hnsw_index(hnsw_baseline): + _assert_reuse_contract(hnsw_baseline) + + +def test_reuse_real_cagra_index(cagra_baseline): + _assert_reuse_contract(cagra_baseline) + + +def _assert_baseline_search(baseline_index): + backend = _backend( + baseline_index.algo, + baseline_index.codec, + baseline_index.runtime_config, + ) + try: + result = _single_search_result( + backend.search( + baseline_index.dataset_case.dataset, + [baseline_index.index], + k=baseline_index.dataset_case.k, + batch_size=2, + ) + ) + finally: + backend.cleanup() + _assert_search_contract(baseline_index, result) + + +def test_search_real_hnsw_index(hnsw_baseline): + _assert_baseline_search(hnsw_baseline) + + +def test_search_real_cagra_index(cagra_baseline): + _assert_baseline_search(cagra_baseline) + + +def test_cagra_rejects_dataset_vector_count_mismatch(tmp_path, cagra_baseline): + _, index = _copy_baseline(cagra_baseline, tmp_path) + case = cagra_baseline.dataset_case + dataset = Dataset( + name="pylucene-integration", + training_vectors=case.dataset.training_vectors[:-1], + query_vectors=case.dataset.query_vectors, + distance_metric="euclidean", + ) + backend = _backend( + cagra_baseline.algo, + cagra_baseline.codec, + cagra_baseline.runtime_config, + ) + try: + build_result = backend.build(dataset, [index], force=False) + search_result = _single_search_result( + backend.search(dataset, [index], k=case.k) + ) + finally: + backend.cleanup() + assert not build_result.success + assert "vector count does not match the dataset: 512 != 511" in ( + build_result.error_message + ) + assert not search_result.success + assert "vector count does not match the dataset: 512 != 511" in ( + search_result.error_message + ) + + +def test_cagra_rejects_missing_vector_data(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + next(index_path.glob("*.vcag")).unlink() + result = _search_copied_index(cagra_baseline, index) + assert not result.success + assert "cannot read" in result.error_message + + +def test_cagra_rejects_truncated_vector_data(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + data_path = next(index_path.glob("*.vcag")) + data = data_path.read_bytes() + data_path.write_bytes(data[:-1]) + result = _search_copied_index(cagra_baseline, index) + assert not result.success + assert "do not exactly cover" in result.error_message + + +def test_cagra_rejects_corrupted_vector_data(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + data_path = next(index_path.glob("*.vcag")) + data = bytearray(data_path.read_bytes()) + data[len(data) // 2] ^= 0xFF + original_stat = data_path.stat() + data_path.write_bytes(data) + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + result = _search_copied_index(cagra_baseline, index) + assert not result.success + assert "checksum" in result.error_message.lower() + + +def test_cagra_rejects_committed_deletion(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + backend = _backend( + cagra_baseline.algo, + cagra_baseline.codec, + cagra_baseline.runtime_config, + ) + try: + _commit_one_deletion(backend, index_path) + with pytest.raises(RuntimeError, match="committed deletions"): + backend._get_runtime().verify_cagra_index(index_path) + result = _single_search_result( + backend.search( + cagra_baseline.dataset_case.dataset, + [index], + k=cagra_baseline.dataset_case.k, + ) + ) + finally: + backend.cleanup() + assert not result.success + assert "does not match the current Lucene commit" in result.error_message + + +def test_cagra_rejects_corrupted_metadata(tmp_path, cagra_baseline): + index_path, _ = _copy_baseline(cagra_baseline, tmp_path) + metadata_path = next(index_path.glob("*.vemc")) + contents = bytearray(metadata_path.read_bytes()) + contents[-1] ^= 0xFF + metadata_path.write_bytes(contents) + runtime = _PyLuceneRuntime.create(cagra_baseline.runtime_config) + with pytest.raises(RuntimeError, match="metadata format v0"): + runtime.verify_cagra_index(index_path) + + +def test_hnsw_rejects_missing_provenance(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + (index_path / _HNSW_PROVENANCE_FILE).unlink() + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "provenance manifest is missing" in result.error_message + + +def test_hnsw_rejects_unreadable_provenance(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + (index_path / _HNSW_PROVENANCE_FILE).write_text("{") + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "provenance manifest cannot be read" in result.error_message + + +def test_hnsw_rejects_changed_lucene_commit(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + segment_path = next(index_path.glob("segments_*")) + segment_path.write_bytes(segment_path.read_bytes() + b"stale") + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "does not match the current Lucene commit" in result.error_message + + +def test_hnsw_rejects_wrong_codec_provenance(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + manifest_path = index_path / _HNSW_PROVENANCE_FILE + payload = json.loads(manifest_path.read_bytes()) + payload["codec"] = _CAGRA_CODEC + manifest_path.write_text(json.dumps(payload)) + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "does not match the requested codec" in result.error_message + + +def _build_and_force_merge(runtime, index_path, vectors, writer_config): + directory = runtime.FSDirectory.open(runtime.Paths.get(str(index_path))) + writer = None + reader = None + try: + writer = runtime.IndexWriter(directory, writer_config) + for document_id, vector in enumerate(vectors): + writer.addDocument(runtime._vector_document(document_id, vector)) + writer.commit() + + reader = runtime.DirectoryReader.open(writer) + initial_segment_count = int(reader.leaves().size()) + initial_segment_document_counts = tuple( + int(leaf.reader().maxDoc()) for leaf in reader.leaves() + ) + reader.close() + reader = None + + writer.forceMerge(1) + writer.commit() + reader = runtime.DirectoryReader.open(writer) + final_segment_count = int(reader.leaves().size()) + final_segment_document_counts = tuple( + int(leaf.reader().maxDoc()) for leaf in reader.leaves() + ) + finally: + try: + if reader is not None: + reader.close() + finally: + try: + if writer is not None: + writer.close() + finally: + directory.close() + return ( + initial_segment_count, + initial_segment_document_counts, + final_segment_count, + final_segment_document_counts, + ) + + +def _assert_no_graph_clamp_warnings(captured_output): + normalized_output = captured_output.lower() + for fragment in _GRAPH_CLAMP_WARNING_FRAGMENTS: + assert fragment.lower() not in normalized_output, captured_output + + +def _assert_persisted_hnsw_m_and_traversal( + context, + index_path, + target, + *, + expected_document_id, + expected_m, + top_k, +): + runtime = context.runtime + directory = runtime.FSDirectory.open(runtime.Paths.get(str(index_path))) + reader = None + try: + reader = runtime.DirectoryReader.open(directory) + searcher = runtime.IndexSearcher(reader) + query = _new_hnsw_graph_verifying_query( + context, + tuple(float(value) for value in target), + top_k, + expected_m, + ) + hits = searcher.search(query, top_k).scoreDocs + assert str(query.getClass().getName()) == ( + HNSW_GRAPH_VERIFYING_QUERY_CLASS + ) + assert len(hits) == top_k + stored_id = searcher.storedFields().document(hits[0].doc).get("id") + assert stored_id is not None + assert int(stored_id) == expected_document_id + finally: + try: + if reader is not None: + reader.close() + finally: + directory.close() + + +def _assert_default_merge_scheduler(writer_config): + assert ( + str(writer_config.getMergeScheduler().getClass().getName()) + == "org.apache.lucene.index.ConcurrentMergeScheduler" + ) + + +def _committed_segment_compound_flags(runtime, index_path): + verifier = runtime._cagra_verifier + runtime.attach_current_thread() + directory = verifier.FSDirectory.open(verifier.Paths.get(str(index_path))) + try: + segment_infos = verifier.SegmentInfos.readLatestCommit(directory) + return [ + bool( + verifier.SegmentCommitInfo.cast_( + raw_segment + ).info.getUseCompoundFile() + ) + for raw_segment in segment_infos + ] + finally: + directory.close() + + +def test_configured_hnsw_codec_selects_gpu_writer( + tmp_path, pylucene_runtime_config, capfd +): + """Exercise public HNSW M/efConstruction through GPU flushes and merge.""" + context = initialize_pylucene_context( + pylucene_runtime_config.backend_config, (_HNSW_CODEC,) + ) + runtime = context.runtime + runtime.attach_current_thread() + m = 16 + ef_construction = 64 + java_codec = _configured_writer_selection_codec( + runtime, m, ef_construction + ) + assert str(java_codec) == ( + "PyLuceneWriterSelectionCodec(" + "PyLuceneConfiguredHnswCodec(m=16, efConstruction=64))" + ) + index_path = tmp_path / "hnsw-writer-selection-index" + index_path.mkdir() + vectors = ( + np.random.default_rng(1907) + .standard_normal((1024, 32)) + .astype(np.float32) + ) + default_config = runtime.IndexWriterConfig() + writer_config = runtime._new_index_writer_config( + _BuildCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec, + writer_policy=_HNSW_WRITER_POLICY, + build_parameters={ + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": m, + "ef_construction": ef_construction, + }, + ), + vector_count=vectors.shape[0], + ) + assert bool(writer_config.getUseCompoundFile()) == bool( + default_config.getUseCompoundFile() + ) + assert float( + writer_config.getMergePolicy().getNoCFSRatio() + ) == pytest.approx(float(default_config.getMergePolicy().getNoCFSRatio())) + _assert_default_merge_scheduler(writer_config) + writer_config.setMaxBufferedDocs(256) + + capfd.readouterr() + ( + initial_segments, + initial_segment_document_counts, + final_segments, + final_segment_document_counts, + ) = _build_and_force_merge(runtime, index_path, vectors, writer_config) + captured = capfd.readouterr() + _assert_no_graph_clamp_warnings(captured.out + captured.err) + assert initial_segments == 4 + assert initial_segment_document_counts == (256, 256, 256, 256) + assert final_segments == 1 + assert final_segment_document_counts == (vectors.shape[0],) + + writer_class, writer_calls = _writer_diagnostics(java_codec) + assert writer_class == _GPU_HNSW_WRITER + assert writer_calls >= 5 + search = runtime.search_index( + index_path, + vectors[[3]], + k=2, + batch_size=1, + num_candidates=64, + ) + assert search.hits[0][0].document_id == 3 + _assert_persisted_hnsw_m_and_traversal( + context, + index_path, + vectors[3], + expected_document_id=3, + expected_m=m, + top_k=16, + ) + + +def test_configured_hnsw_codecs_retain_sequential_parameters( + pylucene_runtime_config, +): + """Resolve independent configured codecs in the process-wide JVM.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + + first = runtime.resolve_configured_hnsw_codec(16, 48) + assert str(first) == "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48)" + + second = runtime.resolve_configured_hnsw_codec(24, 96) + assert ( + str(second) == "PyLuceneConfiguredHnswCodec(m=24, efConstruction=96)" + ) + assert str(first) == "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48)" + assert str(first.getName()) == _HNSW_CODEC + assert str(second.getName()) == _HNSW_CODEC + + +def test_configured_hnsw_direct_single_segment_and_num_candidates( + tmp_path, pylucene_runtime_config, integration_dataset_case +): + """Build one committed leaf and over-fetch without changing top-k.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + build_parameters = { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 16, + "ef_construction": 48, + "direct_single_segment": True, + } + java_codec = runtime.resolve_configured_hnsw_codec(16, 48) + index_path = tmp_path / "configured-hnsw-single-segment" + index_path.mkdir() + vectors = integration_dataset_case.dataset.training_vectors + + topology = runtime.build_index( + index_path, + vectors, + _BuildCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec, + writer_policy=_HNSW_WRITER_POLICY, + build_parameters=build_parameters, + ), + ) + + expected_vector_count = vectors.shape[0] + assert topology.segment_count == 1 + assert topology.segment_document_counts == (expected_vector_count,) + assert topology.segment_vector_counts == (expected_vector_count,) + + top_k = 5 + num_candidates = 11 + search = runtime.search_index( + index_path, + vectors[[137]], + k=top_k, + batch_size=1, + num_candidates=num_candidates, + ) + hits = search.hits[0] + document_ids = [hit.document_id for hit in hits] + assert len(hits) == top_k + assert document_ids[0] == 137 + assert len(set(document_ids)) == top_k + assert all( + 0 <= document_id < expected_vector_count + for document_id in document_ids + ) + + +def test_real_cagra_codec_merges_without_compound_files( + tmp_path, pylucene_runtime_config, integration_dataset_case, capfd +): + """Exercise CAGRA flush and merge layout through production configuration.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + java_codec = runtime.resolve_codec(_CAGRA_CODEC) + index_path = tmp_path / "cagra-merge-index" + index_path.mkdir() + vectors = integration_dataset_case.dataset.training_vectors + writer_config = runtime._new_index_writer_config( + _BuildCodec( + codec_name=_CAGRA_CODEC, + java_codec=java_codec, + writer_policy="gpu-cagra", + build_parameters=_CAGRA_BUILD_PARAMETERS, + ), + vector_count=vectors.shape[0], + ) + assert not bool(writer_config.getUseCompoundFile()) + assert float(writer_config.getMergePolicy().getNoCFSRatio()) == 0.0 + _assert_default_merge_scheduler(writer_config) + writer_config.setMaxBufferedDocs(256) + + capfd.readouterr() + ( + initial_segments, + initial_segment_document_counts, + final_segments, + final_segment_document_counts, + ) = _build_and_force_merge(runtime, index_path, vectors, writer_config) + captured = capfd.readouterr() + _assert_no_graph_clamp_warnings(captured.out + captured.err) + assert initial_segments == 2 + assert initial_segment_document_counts == (256, 256) + assert final_segments == 1 + assert final_segment_document_counts == (vectors.shape[0],) + assert _committed_segment_compound_flags(runtime, index_path) == [False] + + suffixes = {path.suffix for path in index_path.iterdir()} + assert ".vemc" in suffixes + assert ".vcag" in suffixes + assert suffixes.isdisjoint({".cfs", ".cfe"}) + verification = runtime.verify_cagra_index( + index_path, + expected_vector_count=vectors.shape[0], + expected_dimensions=vectors.shape[1], + ) + assert verification.segment_count == 1 + search = runtime.search_index( + index_path, + vectors[[137]], + k=1, + batch_size=1, + num_candidates=1, + ) + assert search.hits[0][0].document_id == 137 + + +def test_real_hnsw_codec_falls_back_to_cpu_in_fresh_process( + pylucene_runtime_config, +): + """Prove production fallback when CUDA devices are hidden at startup.""" + environment = os.environ.copy() + source_root = str(PYTHON_PACKAGE_ROOT) + environment.update( + { + "CUDA_VISIBLE_DEVICES": "", + "CUVS_LUCENE_CUVS_JAVA_JAR": ( + pylucene_runtime_config.backend_config["cuvs_java_jar"] + ), + "CUVS_LUCENE_JAR": pylucene_runtime_config.backend_config[ + "cuvs_lucene_jar" + ], + "JAVA_LIBRARY_PATH": pylucene_runtime_config.backend_config[ + "java_library_path" + ], + PYLUCENE_TEST_CLASSES_ENV: str( + pylucene_runtime_config.test_classes + ), + "PYTHONPATH": os.pathsep.join( + value + for value in (source_root, environment.get("PYTHONPATH")) + if value + ), + } + ) + completed = subprocess.run( + [sys.executable, str(_CPU_FALLBACK_PROBE)], + env=environment, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert completed.returncode == 0, ( + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + assert f"writerClass={_CPU_HNSW_WRITER}" in completed.stdout + assert "firstHit=2" in completed.stdout + + +def test_real_verifier_rejects_cagra_brute_force_fallback( + tmp_path, pylucene_runtime_config +): + """Reject cuVS-Lucene's real one-vector fallback before search.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + vectors = np.ones((1, 32), dtype=np.float32) + index_path = tmp_path / "cagra-fallback-index" + index_path.mkdir() + java_codec = runtime.resolve_codec(_CAGRA_CODEC) + runtime.build_index( + index_path, + vectors, + _BuildCodec( + codec_name=_CAGRA_CODEC, + java_codec=java_codec, + writer_policy="gpu-cagra", + build_parameters=_CAGRA_BUILD_PARAMETERS, + ), + ) + + with pytest.raises(RuntimeError, match="persisted brute-force"): + runtime.verify_cagra_index(index_path, expected_vector_count=1) + + dataset = Dataset( + name="pylucene-cagra-fallback", + training_vectors=vectors, + query_vectors=vectors.copy(), + distance_metric="euclidean", + ) + index = IndexConfig( + name="pylucene-cagra-fallback", + algo="pylucene_cuvs_cagra", + build_param={"codec": _CAGRA_CODEC}, + search_params=[{}], + file=str(index_path), + ) + backend = PyLuceneBackend( + { + "name": index.name, + "algo": index.algo, + "codec": _CAGRA_CODEC, + **pylucene_runtime_config.backend_config, + } + ) + backend._runtime = runtime + try: + result = _single_search_result(backend.search(dataset, [index], k=1)) + assert not result.success + assert "CAGRA provenance manifest is missing" in result.error_message + finally: + backend.cleanup() + + +def test_cli_build_and_search_with_real_pylucene_runtime( + tmp_path, pylucene_runtime_config +): + """Exercise a build/search parameter sweep in a fresh PyLucene process.""" + rng = np.random.default_rng(174) + training_vectors = rng.standard_normal((512, 32)).astype(np.float32) + query_ids = np.asarray([0, 137, 259, 511], dtype=np.int64) + query_vectors = training_vectors[query_ids].copy() + k = 5 + squared_distances = np.sum( + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, + axis=2, + ) + groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] + groundtruth_distances = np.take_along_axis( + squared_distances, groundtruth_neighbors, axis=1 + ) + + dataset_name = "pylucene-cli-integration" + dataset_path = tmp_path / "datasets" + dataset_dir = dataset_path / dataset_name + _write_test_bin(dataset_dir / "base.fbin", training_vectors) + _write_test_bin(dataset_dir / "query.fbin", query_vectors) + _write_test_bin( + dataset_dir / "groundtruth.neighbors.ibin", + groundtruth_neighbors.astype(np.int32), + ) + _write_test_bin( + dataset_dir / "groundtruth.distances.fbin", + groundtruth_distances.astype(np.float32), + ) + + dataset_config = tmp_path / "datasets.yaml" + dataset_config.write_text( + json.dumps( + [ + { + "name": dataset_name, + "base_file": f"{dataset_name}/base.fbin", + "query_file": f"{dataset_name}/query.fbin", + "groundtruth_neighbors_file": ( + f"{dataset_name}/groundtruth.neighbors.ibin" + ), + "groundtruth_distances_file": ( + f"{dataset_name}/groundtruth.distances.fbin" + ), + "distance": "euclidean", + "dims": training_vectors.shape[1], + } + ] + ) + ) + backend_config = tmp_path / "pylucene-backend.yaml" + backend_config.write_text( + json.dumps( + { + "backend": "pylucene", + **pylucene_runtime_config.backend_config, + } + ) + ) + algorithm_config = tmp_path / "pylucene-hnsw-sweep.yaml" + algorithm_config.write_text( + json.dumps( + { + "name": "pylucene_cuvs_hnsw", + "groups": { + "test": { + "build": { + "codec": [_HNSW_CODEC], + "m": [16, 24], + "ef_construction": [32], + "direct_single_segment": [False], + }, + "search": {"num_candidates": [k, 11]}, + } + }, + } + ) + ) + + environment = os.environ.copy() + source_root = str(PYTHON_PACKAGE_ROOT) + environment["PYTHONPATH"] = os.pathsep.join( + value + for value in (source_root, environment.get("PYTHONPATH")) + if value + ) + completed = subprocess.run( + [ + sys.executable, + "-m", + "cuvs_bench.run", + "--backend-config", + str(backend_config), + "--dataset-configuration", + str(dataset_config), + "--configuration", + str(algorithm_config), + "--dataset", + dataset_name, + "--dataset-path", + str(dataset_path), + "--algorithms", + "pylucene_cuvs_hnsw", + "--groups", + "test", + "--batch-size", + "2", + "-k", + str(k), + "-m", + "latency", + "--build", + "--search", + "--force", + ], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert completed.returncode == 0, ( + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + + codec = _HNSW_CODEC + index_names = { + ( + f"pylucene_cuvs_hnsw[group=test][codec={codec}]" + f"[m={m}][ef_construction=32][direct_single_segment=false]" + ) + for m in (16, 24) + } + for index_name in index_names: + index_path = dataset_dir / "index" / index_name + assert any(index_path.glob("segments_*")) + assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + + result_path = dataset_dir / "result" + build_csv = result_path / "build" / "pylucene_cuvs_hnsw,test.csv" + search_stem = f"pylucene_cuvs_hnsw,test,k{k},bs2" + with build_csv.open(newline="") as file: + build_rows = list(csv.DictReader(file)) + assert len(build_rows) == 2 + assert {row["index_name"] for row in build_rows} == index_names + assert {int(row["m"]) for row in build_rows} == {16, 24} + for build_row in build_rows: + assert float(build_row["time"]) > 0 + assert build_row["codec"] == codec + assert build_row["writer_policy"] == _HNSW_WRITER_POLICY + assert build_row["compound_file_policy"] == "lucene-default" + + raw_csv = result_path / "search" / f"{search_stem},raw.csv" + with raw_csv.open(newline="") as file: + csv_rows = list(csv.DictReader(file)) + assert len(csv_rows) == 4 + assert {row["index_name"] for row in csv_rows} == index_names + assert {int(row["num_candidates"]) for row in csv_rows} == {k, 11} + assert { + (int(row["m"]), int(row["num_candidates"])) for row in csv_rows + } == { + (16, k), + (16, 11), + (24, k), + (24, 11), + } + for csv_row in csv_rows: + assert float(csv_row["recall"]) >= 0.75 + assert float(csv_row["throughput"]) > 0 + assert float(csv_row["latency"]) > 0 + assert float(csv_row["build time"]) > 0 + assert float(csv_row["p50"]) > 0 + assert float(csv_row["p95"]) > 0 + assert float(csv_row["p99"]) > 0 + assert csv_row["codec"] == codec + assert csv_row["compound_file_policy"] == "lucene-default" + assert (result_path / "search" / f"{search_stem},latency.csv").is_file() + assert (result_path / "search" / f"{search_stem},throughput.csv").is_file() diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_package_resources.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_package_resources.py new file mode 100644 index 0000000000..c5a91bf8cc --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_package_resources.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Package-resource checks for the PyLucene backend.""" + +from importlib.resources import files +from pathlib import PurePosixPath + + +_REQUIRED_PYLUCENE_RESOURCES = ( + PurePosixPath("backends/_java/PyLuceneConfiguredHnswCodec.java"), + PurePosixPath("config/algos/pylucene_cuvs_cagra.yaml"), + PurePosixPath("config/algos/pylucene_cuvs_hnsw.yaml"), +) + + +def _resource_files(root, relative=PurePosixPath()): + for child in root.iterdir(): + child_relative = relative / child.name + if child.is_dir(): + yield from _resource_files(child, child_relative) + elif child.is_file(): + yield child_relative + + +def test_pylucene_production_resources_are_packaged(): + package_root = files("cuvs_bench") + + for relative in _REQUIRED_PYLUCENE_RESOURCES: + resource = package_root.joinpath(*relative.parts) + assert resource.is_file(), f"missing package resource: {relative}" + assert resource.read_bytes(), f"empty package resource: {relative}" + + +def test_pylucene_test_java_and_classes_are_not_packaged(): + payload = tuple(_resource_files(files("cuvs_bench"))) + test_java = tuple( + relative + for relative in payload + if relative.parts[:2] == ("tests", "java") + ) + class_files = tuple( + relative for relative in payload if relative.suffix == ".class" + ) + + assert not test_java, ( + f"test-only Java resources were packaged: {test_java}" + ) + assert not class_files, ( + f"compiled Java classes were packaged: {class_files}" + ) diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_provenance.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_provenance.py new file mode 100644 index 0000000000..a0a6b265af --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_provenance.py @@ -0,0 +1,352 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Index provenance and training-shape tests for PyLucene.""" + +from __future__ import annotations + +import json +import stat + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.base import Dataset +from cuvs_bench.tests.pylucene._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, + _write_test_bin, +) + + +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} + + +def test_hnsw_provenance_round_trip(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + + verification = pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_vector_count=10, + expected_dimensions=4, + ) + + assert verification.to_metadata() == { + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 4, + "codec": _HNSW_CODEC, + "build_parameters": _DEFAULT_HNSW_BUILD_PARAMETERS, + "writer_policy": "gpu-with-cpu-fallback", + "compound_file_policy": "lucene-default", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + payload = json.loads(manifest_path.read_text()) + assert payload["build_parameters"] == _DEFAULT_HNSW_BUILD_PARAMETERS + assert payload["commit_fingerprints"] == [ + { + "name": "segments_1", + "sha256": ( + "1bc04b5291c26a46d918139138b992d2de976d6851d0893b0476b85bfbdfc6e6" + ), + } + ] + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + assert list(index_path.glob(f"{manifest_path.name}.*.tmp")) == [] + + +def test_cagra_provenance_round_trip(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_cagra_index(index_path) + + verification = pylucene_backend._verify_cagra_provenance( + index_path, + expected_vector_count=10, + expected_dimensions=4, + ) + + assert verification.to_metadata() == { + "status": "gpu-cagra-provenance", + "schema_version": 4, + "codec": _CAGRA_CODEC, + "build_parameters": {"codec": _CAGRA_CODEC}, + "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + + +def test_hnsw_provenance_rejects_build_parameter_mismatch(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + with pytest.raises(RuntimeError, match="build parameter"): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_build_parameters={ + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 24, + }, + ) + + +def test_hnsw_provenance_rejects_inconsistent_direct_segment_count(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + build_parameters = { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "direct_single_segment": True, + } + payload = json.loads(manifest_path.read_text()) + payload["build_parameters"] = build_parameters + payload["segment_count"] = 2 + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="must record exactly one segment"): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_build_parameters=build_parameters, + ) + + +@pytest.mark.parametrize( + "malformed_build_parameters", + [ + { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": True, + }, + { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "ef_construction": 513, + }, + { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "direct_single_segment": 0, + }, + ], +) +def test_hnsw_provenance_rejects_malformed_build_parameters( + tmp_path, malformed_build_parameters +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + payload["build_parameters"] = malformed_build_parameters + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="build parameter"): + pylucene_backend._verify_hnsw_provenance(index_path, _HNSW_CODEC) + + +@pytest.mark.parametrize( + ("expected_vector_count", "expected_dimensions", "error"), + [ + (11, 4, "vector count"), + (10, 8, "dimensions"), + ], +) +def test_hnsw_provenance_rejects_dataset_shape_mismatch( + tmp_path, expected_vector_count, expected_dimensions, error +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + + +@pytest.mark.parametrize( + ("manifest_state", "error"), + [ + ("wrong-policy", "expected writer policy"), + ("wrong-compound-policy", "expected compound-file policy"), + ("boolean-count", "positive integer"), + ("malformed-fingerprint", "malformed Lucene commit fingerprint"), + ("duplicate-fingerprint", "must be unique and sorted"), + ], +) +def test_hnsw_provenance_rejects_malformed_manifest_fields( + tmp_path, manifest_state, error +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + + if manifest_state == "wrong-policy": + payload["writer_policy"] = "gpu-cagra" + elif manifest_state == "wrong-compound-policy": + payload["compound_file_policy"] = "disabled" + elif manifest_state == "boolean-count": + payload["vector_count"] = True + elif manifest_state == "malformed-fingerprint": + payload["commit_fingerprints"] = [{"name": "segments_1"}] + else: + payload["commit_fingerprints"] *= 2 + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + ) + + +@pytest.mark.parametrize( + ("field_name", "field_value", "error"), + [ + ("name", 1, "commit filename must be a string"), + ("name", "commit_1", "must start with 'segments_'"), + ("name", "segments_1/child", "must not contain a path"), + ("sha256", 1, "commit SHA-256 must be a string"), + ("sha256", "0" * 63, "must contain 64 characters"), + ("sha256", "A" * 64, "must be lowercase hexadecimal"), + ], +) +def test_hnsw_provenance_reports_invalid_commit_fingerprint_field( + tmp_path, field_name, field_value, error +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + payload["commit_fingerprints"][0][field_name] = field_value + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance(index_path, _HNSW_CODEC) + + +def test_expected_training_shape_prefers_explicit_vectors_over_base_file( + tmp_path, +): + base_file = tmp_path / "different.fbin" + _write_test_bin(base_file, np.zeros((3, 8), dtype=np.float32)) + dataset = _dataset(n_base=10, dimensions=4) + dataset.base_file = str(base_file) + + assert pylucene_backend._expected_training_shape(dataset) == (10, 4) + + +def test_expected_training_shape_does_not_validate_unused_file_metadata( + tmp_path, +): + dataset = _dataset(n_base=10, dimensions=4) + dataset.base_file = str(tmp_path / "missing.ibin") + dataset.metadata["subset_size"] = True + + assert pylucene_backend._expected_training_shape(dataset) == (10, 4) + + +def test_expected_training_shape_returns_none_without_training_source(): + dataset = Dataset( + name="query-only", + query_vectors=np.zeros((2, 4), dtype=np.float32), + distance_metric="euclidean", + ) + + assert pylucene_backend._expected_training_shape(dataset) is None + + +def test_expected_training_shape_clamps_file_to_valid_subset(tmp_path): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": 3}, + ) + + assert pylucene_backend._expected_training_shape(dataset) == (3, 4) + assert dataset.loaded_training_vectors is None + + +def test_build_reuses_file_backed_subset_without_loading_vectors(tmp_path): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": 3}, + ) + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path, vector_count=3, dimensions=4) + backend = _backend() + + result = backend.build(dataset, [_index(index_path)], force=False) + + assert result.success + assert result.metadata["skipped"] is True + assert result.metadata["hnsw_verification"]["vector_count"] == 3 + assert dataset.loaded_training_vectors is None + assert backend._runtime is None + + +@pytest.mark.parametrize("subset_size", [0, -1, True, "3"]) +def test_expected_training_shape_rejects_invalid_file_subset( + tmp_path, subset_size +): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": subset_size}, + ) + + with pytest.raises(ValueError, match="positive integer"): + pylucene_backend._expected_training_shape(dataset) + + +def test_reused_index_rejects_non_float32_base_file(tmp_path): + base_file = tmp_path / "base.ibin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.int32)) + dataset = Dataset( + name="integer-data", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + ) + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + result = _backend(_FakeRuntime()).build( + dataset, [_index(index_path)], force=False + ) + + assert not result.success + assert "must use float32 values, got int32" in result.error_message diff --git a/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_runtime.py new file mode 100644 index 0000000000..b24dcb66dd --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene/test_pylucene_runtime.py @@ -0,0 +1,1064 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyLucene JVM and Java-boundary unit tests.""" + +from __future__ import annotations + +import zipfile +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuvs_bench.backends._pylucene_java as pylucene_java +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _BuildCodec, _IndexTopology +from cuvs_bench.tests.pylucene._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, +) + + +def _build_codec( + java_codec=None, + codec_name=_HNSW_CODEC, + *, + build_parameters=None, +) -> _BuildCodec: + if build_parameters is None: + build_parameters = ( + {"codec": _CAGRA_CODEC} + if codec_name == _CAGRA_CODEC + else { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + } + ) + return _BuildCodec( + codec_name=codec_name, + java_codec=java_codec if java_codec is not None else object(), + writer_policy=( + "gpu-cagra" + if codec_name == _CAGRA_CODEC + else "gpu-with-cpu-fallback" + ), + build_parameters=build_parameters, + ) + + +class _FakeVMEnvironment: + def __init__(self): + self.attach_count = 0 + + def attachCurrentThread(self): + self.attach_count += 1 + + +class _FakeLuceneModule: + CLASSPATH = "/pylucene/lucene-core.jar" + VERSION = pylucene_backend._REQUIRED_PYLUCENE_VERSION + + def __init__(self): + self.environment = None + self.init_calls = [] + + def getVMEnv(self): + return self.environment + + def initVM(self, **kwargs): + self.init_calls.append(kwargs) + self.environment = _FakeVMEnvironment() + return self.environment + + +class _FakeIndexWriter: + def __init__(self, error_at=None): + self.error_at = error_at + self.documents = [] + self.committed = False + self.rollback_called = False + self.close_called = False + self.force_merge_calls = [] + + def addDocument(self, document): + if self.error_at in {"interrupt", "interrupt-rollback"}: + raise KeyboardInterrupt + if self.error_at in {"add", "rollback"}: + raise RuntimeError("add failed") + self.documents.append(document) + + def commit(self): + if self.error_at == "commit": + raise RuntimeError("commit failed") + self.committed = True + + def rollback(self): + self.rollback_called = True + if self.error_at in {"rollback", "interrupt-rollback"}: + raise RuntimeError("rollback failed") + + def close(self): + self.close_called = True + if self.error_at == "close": + raise RuntimeError("close failed") + + def forceMerge(self, segment_count): + self.force_merge_calls.append(segment_count) + + +class _FakeMergePolicy: + def __init__(self): + self.no_cfs_ratio = None + + def setNoCFSRatio(self, ratio): + self.no_cfs_ratio = ratio + + +class _FakeIndexWriterConfig: + OpenMode = SimpleNamespace(CREATE=object()) + DISABLE_AUTO_FLUSH = -1 + + def __init__(self): + self.use_compound_file = None + self.merge_policy = _FakeMergePolicy() + self.max_buffered_docs = None + self.ram_buffer_size_mb = None + + def setOpenMode(self, _mode): + pass + + def setCodec(self, _codec): + pass + + def setUseCompoundFile(self, enabled): + self.use_compound_file = enabled + + def getMergePolicy(self): + return self.merge_policy + + def setMaxBufferedDocs(self, max_buffered_docs): + self.max_buffered_docs = max_buffered_docs + + def setRAMBufferSizeMB(self, ram_buffer_size_mb): + self.ram_buffer_size_mb = ram_buffer_size_mb + + def setMergePolicy(self, merge_policy): + self.merge_policy = merge_policy + + def setMergeScheduler(self, _scheduler): + raise AssertionError("PyLucene must use Lucene's default scheduler") + + +class _FakeAvailableCodecs: + def __init__(self, names): + self.names = names + + def contains(self, name): + return name in self.names + + def __iter__(self): + return iter(self.names) + + +def _fake_codec_runtime(codec, available_names=(_HNSW_CODEC,)): + registry = SimpleNamespace( + calls=[], + availableCodecs=lambda: _FakeAvailableCodecs(available_names), + ) + + def for_name(codec_name): + registry.calls.append(codec_name) + return codec + + registry.forName = for_name + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.Codec = registry + runtime._codec_cache = {} + return runtime, registry + + +def _fake_configured_codec_runtime(diagnostics, initial_properties=None): + properties = dict(initial_properties or {}) + property_snapshots = [] + class_names = [] + + class _System: + @staticmethod + def getProperty(name): + return properties.get(name) + + @staticmethod + def setProperty(name, value): + properties[name] = value + + @staticmethod + def clearProperty(name): + properties.pop(name, None) + + class _Codec: + @staticmethod + def getName(): + return _HNSW_CODEC + + @staticmethod + def knnVectorsFormat(): + return object() + + def __str__(self): + return diagnostics + + codec = _Codec() + + def new_instance(): + property_snapshots.append(dict(properties)) + return codec + + def for_name(class_name): + class_names.append(class_name) + return SimpleNamespace(newInstance=new_instance) + + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.System = _System + runtime.Class = SimpleNamespace(forName=for_name) + runtime.Codec = SimpleNamespace(cast_=lambda reflected: reflected) + return runtime, properties, property_snapshots, class_names, codec + + +def _fake_index_writer_runtime(error_at=None, directory_close_error=None): + writer = _FakeIndexWriter(error_at=error_at) + directory = SimpleNamespace(closed=False) + + def close_directory(): + directory.closed = True + if directory_close_error is not None: + raise directory_close_error + + directory.close = close_directory + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.resolve_codec = lambda _codec_name: object() + runtime.Paths = SimpleNamespace(get=lambda path: path) + runtime.FSDirectory = SimpleNamespace(open=lambda _path: directory) + runtime.IndexWriterConfig = _FakeIndexWriterConfig + runtime.NoMergePolicy = SimpleNamespace(INSTANCE=object()) + + def create_writer(_directory, config): + runtime._test_writer_config = config + return writer + + runtime.IndexWriter = create_writer + runtime._vector_document = lambda document_id, vector: ( + document_id, + vector.copy(), + ) + runtime._test_writer = writer + runtime._test_directory = directory + runtime._test_codec = object() + runtime._index_topology = lambda _directory: _IndexTopology( + segment_document_counts=(len(writer.documents),), + segment_vector_counts=(len(writer.documents),), + ) + return runtime + + +@pytest.fixture(autouse=True) +def _reset_jvm_tracking(monkeypatch): + monkeypatch.setattr(pylucene_backend, "_INITIALIZED_CLASSPATH", None) + monkeypatch.setattr(pylucene_backend, "_INITIALIZED_VMARGS", None) + monkeypatch.setattr( + pylucene_backend, + "configured_codec_classes_path", + lambda *_args: "/configured-codec-classes", + ) + + +def test_initialize_pylucene_uses_verified_classpath_and_vmargs( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + returned = pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/native", + "jvm_args": ["-Xms1g"], + } + ) + + assert returned is fake_lucene + assert len(fake_lucene.init_calls) == 1 + init_call = fake_lucene.init_calls[0] + assert init_call["classpath"].split(":") == [ + "/configured-codec-classes", + str(cuvs_java), + str(cuvs_lucene), + fake_lucene.CLASSPATH, + ] + assert init_call["vmargs"] == [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + "-Djava.library.path=/native", + "-Xms1g", + ] + assert fake_lucene.environment.attach_count == 1 + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/native", + "jvm_args": ["-Xms1g"], + } + ) + assert len(fake_lucene.init_calls) == 1 + assert fake_lucene.environment.attach_count == 2 + + +def test_configured_codec_compile_error_identifies_required_cuvs_lucene_api( + monkeypatch, +): + completed = SimpleNamespace( + returncode=1, + stderr="cannot find symbol: withHnswHeuristicType", + stdout="", + ) + monkeypatch.setattr(pylucene_java, "_find_javac", lambda: "/jdk/bin/javac") + monkeypatch.setattr( + pylucene_java.subprocess, + "run", + lambda *_args, **_kwargs: completed, + ) + + with pytest.raises( + RuntimeError, + match="PyLucene 10.2 codec support.*HNSW heuristic API", + ): + pylucene_java._compile("/dependencies") + + +@pytest.mark.parametrize( + "lucene", + [SimpleNamespace(VERSION="10.0.0"), SimpleNamespace()], + ids=["mismatched", "missing"], +) +def test_validate_pylucene_version_rejects_incompatible_binding(lucene): + with pytest.raises( + RuntimeError, + match="expected 10[.]2[.]0, found .*Activate a matching PyLucene build", + ): + pylucene_backend._validate_pylucene_version(lucene) + + +def test_initialize_pylucene_checks_version_before_starting_jvm(monkeypatch): + fake_lucene = _FakeLuceneModule() + fake_lucene.VERSION = "10.0.0" + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="expected 10[.]2[.]0"): + pylucene_backend._initialize_pylucene({}) + + assert fake_lucene.init_calls == [] + + +@pytest.mark.parametrize( + ("library_paths", "expected_library_path"), + [ + ( + { + "JAVA_LIBRARY_PATH": "/java-native", + "LD_LIBRARY_PATH": "/ld-native", + }, + "/java-native", + ), + ({"LD_LIBRARY_PATH": "/ld-native"}, "/ld-native"), + ], + ids=["java-library-path-precedence", "ld-library-path-fallback"], +) +def test_initialize_pylucene_uses_environment_runtime_config( + library_paths, + expected_library_path, + tmp_path, + monkeypatch, +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + monkeypatch.setenv("CUVS_LUCENE_CUVS_JAVA_JAR", str(cuvs_java)) + monkeypatch.setenv("CUVS_LUCENE_JAR", str(cuvs_lucene)) + monkeypatch.delenv("JAVA_LIBRARY_PATH", raising=False) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + for name, value in library_paths.items(): + monkeypatch.setenv(name, value) + + pylucene_backend._initialize_pylucene({}) + + init_call = fake_lucene.init_calls[0] + assert init_call["classpath"].split(":") == [ + "/configured-codec-classes", + str(cuvs_java), + str(cuvs_lucene), + fake_lucene.CLASSPATH, + ] + assert init_call["vmargs"] == [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + f"-Djava.library.path={expected_library_path}", + ] + + +def test_initialize_pylucene_rejects_fat_cuvs_lucene_jar_before_init( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene-jar-with-dependencies.jar" + cuvs_java.touch() + with zipfile.ZipFile(cuvs_lucene, "w") as archive: + archive.writestr( + pylucene_backend._LUCENE_CORE_CLASS, + b"bundled Lucene bytecode", + ) + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="standard thin cuvs-lucene JAR"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + } + ) + + assert fake_lucene.init_calls == [] + + +def test_initialize_pylucene_rejects_externally_started_jvm( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + fake_lucene.environment = _FakeVMEnvironment() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="initialized before"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + } + ) + + +def test_initialize_pylucene_rejects_different_jars_after_start( + tmp_path, monkeypatch +): + first_java = tmp_path / "first-java.jar" + first_lucene = tmp_path / "first-lucene.jar" + second_java = tmp_path / "second-java.jar" + for path in (first_java, first_lucene, second_java): + path.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": first_java, + "cuvs_lucene_jar": first_lucene, + } + ) + + with pytest.raises(RuntimeError, match="different"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": second_java, + "cuvs_lucene_jar": first_lucene, + } + ) + + +def test_initialize_pylucene_rejects_different_vmargs_after_start( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/first", + } + ) + + with pytest.raises(RuntimeError, match="different JVM arguments"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/second", + } + ) + + +def test_initialize_pylucene_reports_missing_binding(monkeypatch): + def missing_import(_name): + raise ImportError("missing") + + monkeypatch.setattr( + pylucene_backend.importlib, "import_module", missing_import + ) + + with pytest.raises(ImportError, match="must be built"): + pylucene_backend._initialize_pylucene({}) + + +def test_initialize_pylucene_reports_missing_jar(monkeypatch): + monkeypatch.delenv("CUVS_LUCENE_CUVS_JAVA_JAR", raising=False) + monkeypatch.delenv("CUVS_LUCENE_JAR", raising=False) + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: _FakeLuceneModule(), + ) + + with pytest.raises(RuntimeError, match="cuvs_java_jar"): + pylucene_backend._initialize_pylucene({}) + + +def test_resolve_codec_validates_and_caches_initialized_vector_format(): + vectors_format = object() + codec = SimpleNamespace( + getName=lambda: _HNSW_CODEC, + knnVectorsFormat=lambda: vectors_format, + ) + runtime, registry = _fake_codec_runtime(codec) + + assert runtime.resolve_codec(_HNSW_CODEC) is codec + assert runtime.resolve_codec(_HNSW_CODEC) is codec + assert registry.calls == [_HNSW_CODEC] + + +@pytest.mark.parametrize( + ("available_names", "returned_name", "vectors_format", "error"), + [ + ((), _HNSW_CODEC, object(), "was not advertised by Lucene SPI"), + ((_HNSW_CODEC,), "DifferentCodec", object(), "Requested codec"), + ( + (_HNSW_CODEC,), + _HNSW_CODEC, + None, + "did not initialize a Lucene vector format", + ), + ], + ids=["missing-spi", "wrong-name", "missing-vector-format"], +) +def test_resolve_codec_rejects_unusable_codec( + available_names, returned_name, vectors_format, error +): + codec = SimpleNamespace( + getName=lambda: returned_name, + knnVectorsFormat=lambda: vectors_format, + ) + runtime, _ = _fake_codec_runtime(codec, available_names) + + with pytest.raises(RuntimeError, match=error): + runtime.resolve_codec(_HNSW_CODEC) + + +def test_resolve_configured_hnsw_codec_passes_parameters_and_restores_state(): + diagnostics = "PyLuceneConfiguredHnswCodec(m=24, efConstruction=96)" + runtime, properties, snapshots, class_names, codec = ( + _fake_configured_codec_runtime( + diagnostics, + {pylucene_backend.M_PROPERTY: "previous-m"}, + ) + ) + + assert runtime.resolve_configured_hnsw_codec(24, 96) is codec + assert snapshots == [ + { + pylucene_backend.M_PROPERTY: "24", + pylucene_backend.EF_CONSTRUCTION_PROPERTY: "96", + } + ] + assert properties == {pylucene_backend.M_PROPERTY: "previous-m"} + assert class_names == [pylucene_backend.CONFIGURED_HNSW_CODEC_CLASS] + + +def test_resolve_configured_hnsw_codec_rejects_diagnostic_mismatch_and_restores_state(): + runtime, properties, snapshots, _, _ = _fake_configured_codec_runtime( + "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48)", + {pylucene_backend.EF_CONSTRUCTION_PROPERTY: "previous-ef"}, + ) + + with pytest.raises(RuntimeError, match="did not retain the requested"): + runtime.resolve_configured_hnsw_codec(24, 96) + + assert snapshots == [ + { + pylucene_backend.M_PROPERTY: "24", + pylucene_backend.EF_CONSTRUCTION_PROPERTY: "96", + } + ] + assert properties == { + pylucene_backend.EF_CONSTRUCTION_PROPERTY: "previous-ef" + } + + +@pytest.mark.parametrize("jvm_args", ["-Xmx1g", ["-Xmx1g", 1]]) +def test_initialize_pylucene_rejects_invalid_jvm_args( + jvm_args, tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: _FakeLuceneModule(), + ) + + with pytest.raises(TypeError, match="jvm_args"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "jvm_args": jvm_args, + } + ) + + +def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): + runtime = _fake_index_writer_runtime() + vectors = np.zeros((2, 4), dtype=np.float32) + + result = runtime.build_index( + tmp_path, + vectors, + _build_codec(runtime._test_codec), + ) + + assert result == _IndexTopology( + segment_document_counts=(2,), + segment_vector_counts=(2,), + ) + assert len(runtime._test_writer.documents) == 2 + assert runtime._test_writer.committed is True + assert runtime._test_writer.rollback_called is False + assert runtime._test_writer.close_called is True + assert runtime._test_directory.closed is True + assert runtime._test_writer_config.use_compound_file is None + assert runtime._test_writer_config.merge_policy.no_cfs_ratio is None + assert runtime._test_writer.force_merge_calls == [] + + +@pytest.mark.parametrize( + ("codec_name", "use_compound_file", "no_cfs_ratio"), + [ + (_HNSW_CODEC, None, None), + (_CAGRA_CODEC, False, 0.0), + ], +) +def test_runtime_writer_config_applies_codec_compound_file_policy( + codec_name, use_compound_file, no_cfs_ratio +): + runtime = _fake_index_writer_runtime() + + config = runtime._new_index_writer_config( + _build_codec(runtime._test_codec, codec_name), + vector_count=10, + ) + + assert config.use_compound_file is use_compound_file + assert config.merge_policy.no_cfs_ratio == no_cfs_ratio + + +def test_runtime_writer_config_builds_one_segment_without_force_merge(): + runtime = _fake_index_writer_runtime() + build_parameters = { + "codec": _HNSW_CODEC, + "m": 24, + "ef_construction": 96, + "direct_single_segment": True, + } + + config = runtime._new_index_writer_config( + _build_codec( + runtime._test_codec, + build_parameters=build_parameters, + ), + vector_count=1_000_000, + ) + + assert config.max_buffered_docs == 1_000_001 + assert config.ram_buffer_size_mb == config.DISABLE_AUTO_FLUSH + assert type(config.ram_buffer_size_mb) is float + assert config.merge_policy is runtime.NoMergePolicy.INSTANCE + + +def test_runtime_direct_single_segment_build_does_not_force_merge(tmp_path): + runtime = _fake_index_writer_runtime() + vectors = np.zeros((2, 4), dtype=np.float32) + build_parameters = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": True, + } + + topology = runtime.build_index( + tmp_path, + vectors, + _build_codec( + runtime._test_codec, + build_parameters=build_parameters, + ), + ) + + assert topology.segment_count == 1 + assert runtime._test_writer.force_merge_calls == [] + assert runtime._test_writer.close_called is True + assert runtime._test_directory.closed is True + + +def test_runtime_reads_committed_segment_topology_and_closes_reader(): + reader = SimpleNamespace(closed=False) + + def close_reader(): + reader.closed = True + + def leaf(document_count, vector_count): + leaf_reader = SimpleNamespace( + numDocs=lambda: document_count, + getFloatVectorValues=lambda _field: SimpleNamespace( + size=lambda: vector_count + ), + ) + return SimpleNamespace(reader=lambda: leaf_reader) + + reader.close = close_reader + reader.leaves = lambda: [leaf(4, 4), leaf(6, 6)] + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.DirectoryReader = SimpleNamespace(open=lambda _directory: reader) + + topology = runtime._index_topology(object()) + + assert topology == _IndexTopology( + segment_document_counts=(4, 6), + segment_vector_counts=(4, 6), + ) + assert reader.closed is True + + +def test_runtime_topology_rejects_segment_without_vectors_and_closes_reader(): + leaf_reader = SimpleNamespace( + numDocs=lambda: 2, + getFloatVectorValues=lambda _field: None, + ) + reader = SimpleNamespace( + leaves=lambda: [SimpleNamespace(reader=lambda: leaf_reader)], + closed=False, + ) + + def close_reader(): + reader.closed = True + + reader.close = close_reader + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.DirectoryReader = SimpleNamespace(open=lambda _directory: reader) + + with pytest.raises(RuntimeError, match="contains no vector values"): + runtime._index_topology(object()) + + assert reader.closed is True + + +@pytest.mark.parametrize( + ("topology", "direct_single_segment", "error"), + [ + (_IndexTopology((), ()), False, "no committed segments"), + ( + _IndexTopology((1,), (1,)), + False, + "document counts do not match", + ), + ( + _IndexTopology((2,), (1,)), + False, + "document and vector counts do not match", + ), + ( + _IndexTopology((1, 1), (1, 1)), + True, + "requested one committed Lucene segment", + ), + ], + ids=[ + "no-segments", + "wrong-document-count", + "wrong-vector-count", + "multiple-direct-segments", + ], +) +def test_runtime_build_rejects_invalid_topology_and_closes_resources( + tmp_path, topology, direct_single_segment, error +): + runtime = _fake_index_writer_runtime() + runtime._index_topology = lambda _directory: topology + vectors = np.zeros((2, 4), dtype=np.float32) + build_parameters = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": direct_single_segment, + } + + with pytest.raises(RuntimeError, match=error): + runtime.build_index( + tmp_path, + vectors, + _build_codec( + runtime._test_codec, + build_parameters=build_parameters, + ), + ) + + assert runtime._test_writer.close_called is True + assert runtime._test_directory.closed is True + + +@pytest.mark.parametrize( + ("error_at", "expected_error", "expect_rollback", "expect_close"), + [ + ("add", "add failed", True, False), + ("commit", "commit failed", True, False), + ("rollback", "rollback failed", True, False), + ("close", "close failed", False, True), + ], +) +def test_runtime_build_index_preserves_transaction_cleanup( + tmp_path, error_at, expected_error, expect_rollback, expect_close +): + runtime = _fake_index_writer_runtime(error_at=error_at) + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(RuntimeError, match=expected_error): + runtime.build_index( + tmp_path, + vectors, + _build_codec(runtime._test_codec), + ) + + assert runtime._test_writer.rollback_called is expect_rollback + assert runtime._test_writer.close_called is expect_close + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_rolls_back_on_interrupt(tmp_path): + runtime = _fake_index_writer_runtime(error_at="interrupt") + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt): + runtime.build_index( + tmp_path, + vectors, + _build_codec(runtime._test_codec), + ) + + assert runtime._test_writer.rollback_called is True + assert runtime._test_writer.close_called is False + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_preserves_interrupt_when_rollback_fails(tmp_path): + runtime = _fake_index_writer_runtime(error_at="interrupt-rollback") + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt) as exc_info: + runtime.build_index( + tmp_path, + vectors, + _build_codec(runtime._test_codec), + ) + + assert exc_info.value.__notes__ == [ + "IndexWriter rollback also failed: RuntimeError: rollback failed" + ] + assert runtime._test_writer.rollback_called is True + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_preserves_interrupt_when_directory_close_fails( + tmp_path, +): + runtime = _fake_index_writer_runtime( + error_at="interrupt", + directory_close_error=RuntimeError("directory close failed"), + ) + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt) as exc_info: + runtime.build_index( + tmp_path, + vectors, + _build_codec(runtime._test_codec), + ) + + assert exc_info.value.__notes__ == [ + "Failed to close Lucene directory: RuntimeError: directory close failed" + ] + assert runtime._test_writer.rollback_called is True + assert runtime._test_directory.closed is True + + +def test_runtime_search_uses_candidates_for_query_and_top_k_for_results(): + query_calls = [] + search_calls = [] + query = object() + + def new_query(field, vector, num_candidates): + query_calls.append((field, vector, num_candidates)) + return query + + def search(received_query, top_k): + search_calls.append((received_query, top_k)) + return SimpleNamespace(scoreDocs=[SimpleNamespace(doc=4, score=0.5)]) + + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.KnnFloatVectorQuery = new_query + runtime._java_float_array = lambda vector: tuple(vector.tolist()) + searcher = SimpleNamespace(search=search) + stored_fields = SimpleNamespace( + document=lambda _document_id: SimpleNamespace(get=lambda _field: "17") + ) + + hits = runtime._search_vector( + searcher, + stored_fields, + np.array([1.0, 2.0], dtype=np.float32), + k=150, + num_candidates=300, + ) + + assert query_calls == [(pylucene_backend._VECTOR_FIELD, (1.0, 2.0), 300)] + assert search_calls == [(query, 150)] + assert hits == [pylucene_backend._SearchHit(document_id=17, score=0.5)] + + +def test_search_cleanup_attempts_directory_after_reader_close_fails(): + close_calls = [] + + def fail_reader_close(): + close_calls.append("reader") + raise RuntimeError("reader close failed") + + def fail_directory_close(): + close_calls.append("directory") + raise RuntimeError("directory close failed") + + reader = SimpleNamespace(close=fail_reader_close) + directory = SimpleNamespace(close=fail_directory_close) + + with pytest.raises(RuntimeError, match="reader close failed") as exc_info: + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + cleanups.add("close Lucene index reader", reader.close) + + assert close_calls == ["reader", "directory"] + assert exc_info.value.__notes__ == [ + "Failed to close Lucene directory: RuntimeError: directory close failed" + ] + + +def test_cleanup_stack_ignores_unrelated_handled_exception(): + def fail_close(): + raise RuntimeError("close failed") + + unrelated_error = None + try: + raise ValueError("unrelated") + except ValueError as exc: + unrelated_error = exc + with pytest.raises(RuntimeError, match="close failed"): + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close test resource", fail_close) + + assert getattr(unrelated_error, "__notes__", []) == [] + + +def test_cleanup_stack_prioritizes_cleanup_interrupt_over_operation_error(): + def interrupt_close(): + raise KeyboardInterrupt("stop") + + with pytest.raises(KeyboardInterrupt, match="stop") as exc_info: + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close test resource", interrupt_close) + raise RuntimeError("operation failed") + + assert exc_info.value.__notes__ == [ + "Raised while attempting to close test resource; prior failure: " + "RuntimeError: operation failed" + ] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cli.py b/python/cuvs_bench/cuvs_bench/tests/test_cli.py index fa2a63d041..a9dd27e76d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cli.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cli.py @@ -1,13 +1,15 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from pathlib import Path +from types import SimpleNamespace import pandas as pd import pytest from click.testing import CliRunner +from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.get_dataset.__main__ import main @@ -519,11 +521,247 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path): ) -# FIXME: Tests below use --dry-run to verify CLI flag parsing and orchestrator -# routing without requiring actual benchmark execution. Tune mode (--mode tune) -# requires Optuna and actual search results, so only flag acceptance is tested -# here. End-to-end tests for tune mode and non-C++ backends should be added -# when those features are exercised in integration testing. +# The mocked-result tests below isolate CLI outcome and export semantics. +# Later tests use --dry-run to verify flag parsing and orchestrator routing +# without requiring native benchmark execution. + + +def _invoke_run_with_results( + monkeypatch, + tmp_path, + mode, + results, + *, + dry_run=False, + backend_type="fake", +): + from cuvs_bench.run import __main__ as run_module + + orchestrator = SimpleNamespace( + run_benchmark=lambda **_kwargs: results, + ) + monkeypatch.setattr( + run_module, + "BenchmarkOrchestrator", + lambda backend_type: orchestrator, + ) + exported = [] + monkeypatch.setattr( + run_module, + "write_results_to_csv", + lambda *args, **kwargs: exported.append((args, kwargs)), + ) + backend_config = tmp_path / "backend.yaml" + backend_config.write_text(f"backend: {backend_type}\n") + + args = [ + "--dataset", + "test-data", + "--dataset-path", + str(tmp_path), + "--algorithms", + "fake", + "--groups", + "base", + "--batch-size", + "1", + "-k", + "1", + "-m", + "latency", + "--mode", + mode, + "--backend-config", + str(backend_config), + ] + if dry_run: + args.append("--dry-run") + + result = CliRunner().invoke( + run_module.main, + args, + ) + return result, exported + + +def _build_result(success=True, error_message=None): + return BuildResult( + index_path="test-index", + build_time_seconds=1.0, + index_size_bytes=1, + algorithm="fake", + build_params={}, + metadata={"group": "base", "index_name": "test-index"}, + success=success, + error_message=error_message, + ) + + +def _search_result(success=True, error_message=None): + return SearchResult( + neighbors=None, + distances=None, + search_time_ms=1.0, + queries_per_second=1.0, + recall=1.0, + success=success, + algorithm="fake", + search_params=[{}], + metadata={"group": "base", "index_name": "test-index"}, + error_message=error_message, + ) + + +def test_tune_mixed_trial_results_exit_zero_and_export(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "tune", + [ + _search_result(False, "trial failed"), + _search_result(), + ], + ) + + assert result.exit_code == 0, result.output + assert len(exported) == 1 + assert exported[0][1]["search_requested"] is True + + +@pytest.mark.parametrize( + ("results", "expected_message"), + [ + ([_search_result(False, "trial failed")], "trial failed"), + ([_build_result()], "no successful search result"), + ([], "tune mode produced no benchmark results"), + ], +) +def test_tune_without_successful_results_exits_nonzero( + monkeypatch, tmp_path, results, expected_message +): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "tune", results + ) + + assert result.exit_code != 0 + assert expected_message in result.output + assert len(exported) == 1 + + +def test_tune_all_constraint_pruned_measurements_exit_zero_and_export( + monkeypatch, tmp_path +): + # _run_tune retains successful measurements even when Optuna prunes every + # trial for violating a hard constraint. + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "tune", + [_search_result(), _search_result()], + ) + + assert result.exit_code == 0, result.output + assert len(exported) == 1 + + +def test_sweep_mixed_results_exit_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [ + _build_result(False, "benchmark failed"), + _search_result(), + ], + ) + + assert result.exit_code != 0 + assert "benchmark failed" in result.output + assert len(exported) == 1 + + +def test_sweep_without_results_exits_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "sweep", [] + ) + + assert result.exit_code != 0 + assert "sweep mode produced no benchmark results" in result.output + assert len(exported) == 1 + + +def test_dry_run_allows_native_backend_without_result_objects( + monkeypatch, tmp_path +): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [], + dry_run=True, + backend_type="cpp_gbench", + ) + + assert result.exit_code == 0, result.output + assert exported == [] + + +def test_dry_run_rejects_python_backend_without_result_objects( + monkeypatch, tmp_path +): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "sweep", [], dry_run=True + ) + + assert result.exit_code != 0 + assert "sweep mode produced no benchmark results" in result.output + assert exported == [] + + +def test_dry_run_reports_explicit_backend_failures(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [_build_result(False, "invalid configuration")], + dry_run=True, + ) + + assert result.exit_code != 0 + assert "invalid configuration" in result.output + assert exported == [] + + +@pytest.mark.parametrize( + "dataset", ["..", "../escape", "nested/dataset", "/absolute/dataset"] +) +def test_data_export_rejects_invalid_dataset_name(tmp_path, dataset): + from cuvs_bench.run.__main__ import main as run_main + + runner = CliRunner() + result = runner.invoke( + run_main, + [ + "--data-export", + "--dataset", + dataset, + "--dataset-path", + str(tmp_path), + "--count", + "10", + "--batch-size", + "100", + "--algorithms", + "cuvs_cagra", + "--groups", + "base", + "--search-mode", + "latency", + ], + ) + + assert result.exit_code == 2 + assert "Invalid benchmark dataset name" in result.output def test_run_with_mode_sweep(temp_datasets_dir): diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py index 721d03c5ac..df43613131 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd +import pytest from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.orchestrator.config_loaders import ( @@ -14,7 +15,10 @@ ) from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator from cuvs_bench.plot.__main__ import load_all_results -from cuvs_bench.run.data_export import write_results_to_csv +from cuvs_bench.run.data_export import ( + validate_dataset_name, + write_results_to_csv, +) def test_python_backend_csv_is_plot_compatible(tmp_path): @@ -23,12 +27,12 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): index_name = "test-index" results = [ BuildResult( - index_path=index_name, + index_path="/tmp/physical-index-path", build_time_seconds=1.5, index_size_bytes=1024, algorithm=algorithm, build_params={"m": 16}, - metadata={"group": "base"}, + metadata={"group": "base", "index_name": index_name}, ), SearchResult( neighbors=np.empty((0, 2), dtype=np.int64), @@ -38,6 +42,7 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): recall=0.5, algorithm=algorithm, search_params=[{"ef_search": 50}], + latency_percentiles={"p50": 8.0, "p95": 12.0, "p99": 15.0}, metadata={ "group": "base", "index_name": index_name, @@ -76,6 +81,7 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): ] assert raw["ef_search"].tolist() == [50, 100] assert raw["build time"].tolist() == [1.5, 1.5] + assert raw.loc[0, ["p50", "p95", "p99"]].tolist() == [8.0, 12.0, 15.0] assert not list(result_path.rglob("*.json")) write_results_to_csv( @@ -197,3 +203,196 @@ def search(self, dataset, indexes, k, **kwargs): assert [type(result) for result in results] == [BuildResult, SearchResult] assert results[1].recall == 1.0 + + +def test_subset_scopes_use_independent_result_files(tmp_path): + dataset = "test-dataset" + algorithm = "pylucene_cuvs_hnsw" + + for scope in ("subset100", "subset200"): + index_name = f"test-index-{scope}" + write_results_to_csv( + [ + BuildResult( + index_path=f"/tmp/{index_name}", + build_time_seconds=1.0, + index_size_bytes=100, + algorithm=algorithm, + build_params={}, + metadata={ + "group": "base", + "index_name": index_name, + "result_scope": scope, + }, + ), + SearchResult( + neighbors=np.empty((0, 1), dtype=np.int64), + distances=np.empty((0, 1), dtype=np.float32), + search_time_ms=1.0, + queries_per_second=10.0, + recall=1.0, + algorithm=algorithm, + search_params=[{}], + metadata={ + "group": "base", + "index_name": index_name, + "result_scope": scope, + }, + ), + ], + dataset, + str(tmp_path), + count=1, + batch_size=1, + ) + + result_path = tmp_path / dataset / "result" + for scope in ("subset100", "subset200"): + assert ( + result_path / "build" / f"{algorithm},base,{scope}.csv" + ).is_file() + assert ( + result_path / "search" / f"{algorithm},base,{scope},k1,bs1,raw.csv" + ).is_file() + + plotted = load_all_results( + str(tmp_path / dataset), + algorithms=[algorithm], + groups=["base"], + algo_groups=[], + k=1, + batch_size=1, + method="search", + index_key="algo", + raw=True, + mode="latency", + time_unit="s", + ) + assert set(plotted) == { + f"{algorithm}[scope=subset100]", + f"{algorithm}[scope=subset200]", + } + + +def test_algorithm_and_group_names_remain_distinct_in_plot_series(tmp_path): + dataset = "test-dataset" + identities = ( + ("a_b", "c", "first-index"), + ("a", "b_c", "second-index"), + ) + + for algorithm, group, index_name in identities: + write_results_to_csv( + [ + SearchResult( + neighbors=np.empty((0, 1), dtype=np.int64), + distances=np.empty((0, 1), dtype=np.float32), + search_time_ms=1.0, + queries_per_second=10.0, + recall=1.0, + algorithm=algorithm, + search_params=[{}], + metadata={"group": group, "index_name": index_name}, + ) + ], + dataset, + str(tmp_path), + count=1, + batch_size=1, + ) + + plotted = load_all_results( + str(tmp_path / dataset), + algorithms=["a_b", "a"], + groups=["c", "b_c"], + algo_groups=[], + k=1, + batch_size=1, + method="search", + index_key="algo", + raw=True, + mode="latency", + time_unit="s", + ) + + assert set(plotted) == {"a_b[group=c]", "a[group=b_c]"} + + +def test_failed_build_removes_only_matching_stale_search_files(tmp_path): + dataset = "test-dataset" + algorithm = "pylucene_cuvs_hnsw" + search_dir = tmp_path / dataset / "result" / "search" + search_dir.mkdir(parents=True) + + matching_stem = f"{algorithm},base,subset100,k1,bs1" + other_stem = f"{algorithm},base,subset200,k1,bs1" + suffixes = (",raw.csv", ",throughput.csv", ",latency.csv") + for stem in (matching_stem, other_stem): + for suffix in suffixes: + (search_dir / f"{stem}{suffix}").write_text("stale\n") + + write_results_to_csv( + [ + BuildResult( + index_path="/tmp/failed-index", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=algorithm, + build_params={}, + metadata={ + "group": "base", + "index_name": "failed-index", + "result_scope": "subset100", + }, + success=False, + error_message="build failed", + ) + ], + dataset, + str(tmp_path), + count=1, + batch_size=1, + search_requested=True, + ) + + assert not list(search_dir.glob(f"{matching_stem}*")) + assert len(list(search_dir.glob(f"{other_stem}*"))) == 3 + + +@pytest.mark.parametrize("dataset", ["", ".", "..", "nested/dataset"]) +def test_dataset_name_must_be_a_safe_path_component(dataset): + with pytest.raises(ValueError, match="Invalid benchmark dataset name"): + validate_dataset_name(dataset) + + +@pytest.mark.parametrize( + ("algorithm", "metadata", "match"), + [ + ("../pylucene", {"group": "base"}, "benchmark algorithm name"), + ("pylucene,other", {"group": "base"}, "benchmark algorithm name"), + ("pylucene[other]", {"group": "base"}, "benchmark algorithm name"), + ("pylucene", {"group": "../base"}, "benchmark group name"), + ("pylucene", {"group": "base\nother"}, "benchmark group name"), + ( + "pylucene", + {"group": "base", "result_scope": "../subset100"}, + "benchmark result scope", + ), + ], +) +def test_result_identity_must_use_safe_path_components( + tmp_path, algorithm, metadata, match +): + result = BuildResult( + index_path="index", + build_time_seconds=1.0, + index_size_bytes=1, + algorithm=algorithm, + build_params={}, + metadata=metadata, + ) + + with pytest.raises(ValueError, match=match): + write_results_to_csv( + [result], "test-dataset", str(tmp_path), count=1, batch_size=1 + ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_utils.py b/python/cuvs_bench/cuvs_bench/tests/test_utils.py index 01dd803a07..8fddb356b0 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_utils.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_utils.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -528,7 +528,9 @@ def test_lazy_load_training_vectors(self, tmp_path): _write_test_bin(path, data) dataset = Dataset(name="test", base_file=path) + assert dataset.loaded_training_vectors is None np.testing.assert_array_equal(dataset.training_vectors, data) + np.testing.assert_array_equal(dataset.loaded_training_vectors, data) def test_lazy_load_query_vectors(self, tmp_path): """Test that query vectors are loaded from file on first access.""" @@ -614,7 +616,7 @@ def test_file_path_access_does_not_trigger_loading(self, tmp_path): _ = dataset.name _ = dataset.distance_metric - assert dataset._training_vectors.size == 0 + assert dataset.loaded_training_vectors is None def test_dims_and_counts(self, tmp_path): """Test dims, n_base, and n_queries properties.""" diff --git a/python/cuvs_bench/java-format-pom.xml b/python/cuvs_bench/java-format-pom.xml new file mode 100644 index 0000000000..b771f45e77 --- /dev/null +++ b/python/cuvs_bench/java-format-pom.xml @@ -0,0 +1,41 @@ + + + + 4.0.0 + + com.nvidia.cuvs + cuvs-bench-java-format + 1.0.0 + pom + + + + ${project.basedir}/../../build/cuvs-bench-java-format + + + com.diffplug.spotless + spotless-maven-plugin + 2.44.5 + + + + cuvs_bench/backends/_java/**/*.java + tests/java/**/*.java + + + 1.27.0 + + true + false + + + + + + + diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 24e54ff28f..43ed84d039 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -21,6 +21,7 @@ requires-python = ">=3.11" dependencies = [ "click", "cuvs==26.10.*,>=0.0.0a0", + "h5py>=3.8.0", "matplotlib>=3.9", "pandas", "pyyaml", @@ -68,6 +69,8 @@ elastic = "cuvs_bench.backends.elasticsearch:register" [tool.pytest.ini_options] markers = [ "opensearch: tests that require a live OpenSearch node (run with '-m opensearch')", + "pylucene: real PyLucene/JVM/cuVS tests (set CUVS_BENCH_PYLUCENE_INTEGRATION=1 and run with '-m pylucene')", + "pylucene_extended: extended PyLucene codec-path, topology, merge, deletion, filter, and searchWidth coverage", ] [tool.isort] diff --git a/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java new file mode 100644 index 0000000000..53f2f0e78e --- /dev/null +++ b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java @@ -0,0 +1,98 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.bench; + +import java.io.IOException; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** Test-only codec that reports the writer selected by the configured production HNSW codec. */ +public final class PyLuceneWriterSelectionCodec extends FilterCodec { + + private static final String CONFIGURED_CODEC_CLASS = + "com.nvidia.cuvs.bench.PyLuceneConfiguredHnswCodec"; + + private final KnnVectorsFormat knnVectorsFormat; + private final String configuredCodecDiagnostics; + + public PyLuceneWriterSelectionCodec() throws Exception { + this(configuredCodec()); + } + + private PyLuceneWriterSelectionCodec(Codec delegate) { + super(delegate.getName(), delegate); + knnVectorsFormat = new WriterSelectionFormat(delegate.knnVectorsFormat()); + configuredCodecDiagnostics = delegate.toString(); + } + + private static Codec configuredCodec() throws Exception { + return (Codec) Class.forName(CONFIGURED_CODEC_CLASS).getConstructor().newInstance(); + } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return knnVectorsFormat; + } + + @Override + public String toString() { + return getClass().getSimpleName() + "(" + configuredCodecDiagnostics + ")"; + } + + private static final class WriterSelectionFormat extends KnnVectorsFormat { + + private static final String NOT_SELECTED = "not-selected"; + + private final KnnVectorsFormat delegate; + private String writerClass = NOT_SELECTED; + private int fieldsWriterCalls; + + private WriterSelectionFormat(KnnVectorsFormat delegate) { + super(delegate.getName()); + this.delegate = delegate; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + KnnVectorsWriter writer = delegate.fieldsWriter(state); + recordWriter(writer.getClass().getName()); + return writer; + } + + private synchronized void recordWriter(String selectedClass) { + if (!NOT_SELECTED.equals(writerClass) && !writerClass.equals(selectedClass)) { + throw new AssertionError( + "Vector writer selection changed from " + writerClass + " to " + selectedClass); + } + writerClass = selectedClass; + fieldsWriterCalls++; + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return delegate.fieldsReader(state); + } + + @Override + public int getMaxDimensions(String fieldName) { + return delegate.getMaxDimensions(fieldName); + } + + @Override + public synchronized String toString() { + return getName() + + "(writerClass=" + + writerClass + + ", fieldsWriterCalls=" + + fieldsWriterCalls + + ")"; + } + } +} diff --git a/python/cuvs_bench/tests/java/com/nvidia/cuvs/lucene/PyLuceneTestSupport.java b/python/cuvs_bench/tests/java/com/nvidia/cuvs/lucene/PyLuceneTestSupport.java new file mode 100644 index 0000000000..2121c6e929 --- /dev/null +++ b/python/cuvs_bench/tests/java/com/nvidia/cuvs/lucene/PyLuceneTestSupport.java @@ -0,0 +1,596 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import java.io.IOException; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.Map; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.FilterLeafReader; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.QueryTimeout; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.knn.KnnCollectorManager; +import org.apache.lucene.util.Bits; + +/** + * Test-only adapters that PyLucene creates reflectively. Public no-argument constructors read query + * configuration from Java system properties. + */ +public final class PyLuceneTestSupport { + + public static final String QUERY_FIELD_PROPERTY = "cuvs.lucene.pylucene.query.field"; + public static final String QUERY_TARGET_PROPERTY = "cuvs.lucene.pylucene.query.target"; + public static final String QUERY_K_PROPERTY = "cuvs.lucene.pylucene.query.k"; + public static final String QUERY_I_TOP_K_PROPERTY = "cuvs.lucene.pylucene.query.iTopK"; + public static final String QUERY_SEARCH_WIDTH_PROPERTY = "cuvs.lucene.pylucene.query.searchWidth"; + public static final String QUERY_EXPECTED_HNSW_M_PROPERTY = + "cuvs.lucene.pylucene.query.expectedHnswM"; + public static final String QUERY_FILTER_FIELD_PROPERTY = "cuvs.lucene.pylucene.query.filterField"; + public static final String QUERY_FILTER_VALUE_PROPERTY = "cuvs.lucene.pylucene.query.filterValue"; + + private static final int CAGRA_GRAPH_DEGREE = 32; + private static final int CAGRA_INTERMEDIATE_GRAPH_DEGREE = 64; + private static final int CPU_HNSW_MAX_CONN = 32; + private static final int CPU_HNSW_BEAM_WIDTH = 32; + + private PyLuceneTestSupport() {} + + /** Stock Lucene CPU HNSW codec with max connections 32. */ + public static final class CpuHnswCodec extends Lucene101AcceleratedHNSWCodec { + + public CpuHnswCodec() throws Exception { + super(); + setKnnFormat( + diagnosticFormat( + new Lucene99HnswVectorsFormat(CPU_HNSW_MAX_CONN, CPU_HNSW_BEAM_WIDTH), + "configuredPath=cpu-hnsw;hnswM=" + CPU_HNSW_MAX_CONN)); + } + } + + /** CAGRA codec with graph degree 32 and intermediate graph degree 64. */ + public static final class CagraSearchCodec extends CuVS2510GPUSearchCodec { + + public CagraSearchCodec() throws Exception { + this(cagraSearchParams()); + } + + private CagraSearchCodec(GPUSearchParams params) throws Exception { + super(params); + setKnnFormat( + diagnosticFormat( + super.knnVectorsFormat(), + "configuredPath=gpu-cagra-search;" + cagraDiagnostics(params))); + } + } + + /** CAGRA-built HNSW codec that persists only the base layer. */ + public static final class CagraBuiltHnswBaseLayerCodec extends Lucene101AcceleratedHNSWCodec { + + public CagraBuiltHnswBaseLayerCodec() throws Exception { + this(cagraBuiltHnswParams(1)); + } + + private CagraBuiltHnswBaseLayerCodec(AcceleratedHNSWParams params) throws Exception { + super(params); + setKnnFormat( + diagnosticFormat( + super.knnVectorsFormat(), + "configuredPath=gpu-cagra-built-hnsw;" + hnswDiagnostics(params))); + } + } + + /** CAGRA-built HNSW codec that persists exactly three layers. */ + public static final class CagraBuiltHnswThreeLayerCodec extends Lucene101AcceleratedHNSWCodec { + + public CagraBuiltHnswThreeLayerCodec() throws Exception { + this(cagraBuiltHnswParams(3)); + } + + private CagraBuiltHnswThreeLayerCodec(AcceleratedHNSWParams params) throws Exception { + super(params); + setKnnFormat( + diagnosticFormat( + super.knnVectorsFormat(), + "configuredPath=gpu-cagra-built-hnsw;" + hnswDiagnostics(params))); + } + } + + private static GPUSearchParams cagraSearchParams() { + return new GPUSearchParams.Builder() + .withStrategy(GPUSearchParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(CAGRA_GRAPH_DEGREE) + .withIntermediateGraphDegree(CAGRA_INTERMEDIATE_GRAPH_DEGREE) + .build(); + } + + private static AcceleratedHNSWParams cagraBuiltHnswParams(int hnswLayers) { + return new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(CAGRA_GRAPH_DEGREE) + .withIntermediateGraphDegree(CAGRA_INTERMEDIATE_GRAPH_DEGREE) + .withHNSWLayer(hnswLayers) + .build(); + } + + private static String cagraDiagnostics(GPUSearchParams params) { + return "cagraStrategy=" + + params.getStrategy().name() + + ";cagraGraphBuildAlgo=" + + params.getCagraGraphBuildAlgo().name() + + ";cagraGraphDegree=" + + params.getGraphdegree() + + ";cagraIntermediateGraphDegree=" + + params.getIntermediateGraphDegree(); + } + + private static String hnswDiagnostics(AcceleratedHNSWParams params) { + return "hnswLayers=" + params.getHnswLayers() + ";" + cagraDiagnostics(params); + } + + private static String cagraDiagnostics(AcceleratedHNSWParams params) { + return "cagraStrategy=" + + params.getStrategy().name() + + ";cagraGraphBuildAlgo=" + + params.getCagraGraphBuildAlgo().name() + + ";cagraGraphDegree=" + + params.getGraphdegree() + + ";cagraIntermediateGraphDegree=" + + params.getIntermediateGraphDegree(); + } + + private static KnnVectorsFormat diagnosticFormat( + KnnVectorsFormat delegate, String configuration) { + return new DiagnosticKnnVectorsFormat(delegate, configuration); + } + + /** Records the concrete writer selected by a test codec without probing or owning resources. */ + private static final class DiagnosticKnnVectorsFormat extends KnnVectorsFormat { + + private static final String NOT_SELECTED = "not-selected"; + + private final KnnVectorsFormat delegate; + private final String configuration; + private volatile String writerClass = NOT_SELECTED; + + private DiagnosticKnnVectorsFormat(KnnVectorsFormat delegate, String configuration) { + super(delegate.getName()); + this.delegate = delegate; + this.configuration = configuration; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + KnnVectorsWriter writer = delegate.fieldsWriter(state); + String selectedClass = writer.getClass().getName(); + String previousClass = writerClass; + if (!NOT_SELECTED.equals(previousClass) && !previousClass.equals(selectedClass)) { + throw new AssertionError( + "Vector writer selection changed from " + previousClass + " to " + selectedClass); + } + writerClass = selectedClass; + return writer; + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return delegate.fieldsReader(state); + } + + @Override + public int getMaxDimensions(String fieldName) { + return delegate.getMaxDimensions(fieldName); + } + + @Override + public String toString() { + return getName() + "(" + configuration + ";writerClass=" + writerClass + ")"; + } + } + + public static final class HnswGraphVerifyingQuery extends KnnFloatVectorQuery { + + private final int expectedM; + + public HnswGraphVerifyingQuery() { + this(QueryProperties.fromSystemPropertiesForHnsw()); + } + + private HnswGraphVerifyingQuery(QueryProperties properties) { + super(properties.field, properties.target, properties.k, properties.filter); + expectedM = parsePositiveInt(QUERY_EXPECTED_HNSW_M_PROPERTY); + } + + @Override + protected TopDocs approximateSearch( + LeafReaderContext context, + Bits acceptDocs, + int visitedLimit, + KnnCollectorManager knnCollectorManager) + throws IOException { + requireExpectedHnswGraph(context, getField(), expectedM); + return super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager); + } + + @Override + protected TopDocs exactSearch( + LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout) + throws IOException { + throw new AssertionError("HNSW graph-verifying query must not use exact vector scoring"); + } + } + + public static final class HnswGraphVerifyingExactQuery extends KnnFloatVectorQuery { + + private final int expectedM; + + public HnswGraphVerifyingExactQuery() { + this(QueryProperties.fromSystemPropertiesForHnsw()); + } + + private HnswGraphVerifyingExactQuery(QueryProperties properties) { + super(properties.field, properties.target, properties.k, properties.filter); + expectedM = parsePositiveInt(QUERY_EXPECTED_HNSW_M_PROPERTY); + } + + @Override + protected TopDocs approximateSearch( + LeafReaderContext context, + Bits acceptDocs, + int visitedLimit, + KnnCollectorManager knnCollectorManager) { + throw new AssertionError( + "HNSW exact-filter graph-verifying query must not use approximate traversal"); + } + + @Override + protected TopDocs exactSearch( + LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout) + throws IOException { + requireExpectedHnswGraph(context, getField(), expectedM); + return super.exactSearch(context, acceptIterator, queryTimeout); + } + } + + private static void requireExpectedHnswGraph( + LeafReaderContext context, String fieldName, int expectedM) { + KnnVectorsReader vectorsReader = vectorReaderForField(context, fieldName); + if (!(vectorsReader instanceof HnswGraphProvider)) { + throw new AssertionError( + "HNSW graph-verifying query requires HnswGraphProvider for field '" + + fieldName + + "', got " + + className(vectorsReader)); + } + + var fieldInfo = context.reader().getFieldInfos().fieldInfo(fieldName); + if (fieldInfo == null) { + throw new AssertionError("HNSW field metadata has no field info for: " + fieldName); + } + + int actualM = persistedHnswM(vectorsReader, fieldName, fieldInfo.number); + if (actualM != expectedM) { + throw new AssertionError( + "HNSW graph-verifying query expected persisted M " + + expectedM + + " for field '" + + fieldName + + "', got " + + actualM); + } + } + + private static int persistedHnswM( + KnnVectorsReader vectorsReader, String fieldName, int fieldNumber) { + // The PyLucene bindings do not expose persisted M through their wrapped HNSW graph API. + try { + Field fieldsField = vectorsReader.getClass().getDeclaredField("fields"); + fieldsField.setAccessible(true); + Object fieldsValue = fieldsField.get(vectorsReader); + Object fieldEntry; + if (fieldsValue instanceof Map fields) { + fieldEntry = fields.get(fieldName); + } else { + Method getAccessor = fieldsValue.getClass().getMethod("get", int.class); + fieldEntry = getAccessor.invoke(fieldsValue, fieldNumber); + } + if (fieldEntry == null) { + throw new AssertionError("Persisted HNSW metadata has no entry for field: " + fieldName); + } + Method mAccessor = fieldEntry.getClass().getDeclaredMethod("M"); + mAccessor.setAccessible(true); + Object value = mAccessor.invoke(fieldEntry); + if (!(value instanceof Integer persistedM)) { + throw new AssertionError("Unexpected persisted HNSW M value: " + value); + } + return persistedM; + } catch (ReflectiveOperationException | RuntimeException exception) { + throw new AssertionError( + "Unable to inspect persisted HNSW M for field '" + fieldName + "'", exception); + } + } + + public static final class CagraSearchQuery extends GPUKnnFloatVectorQuery { + + public CagraSearchQuery() { + this(QueryProperties.fromSystemProperties()); + } + + private CagraSearchQuery(QueryProperties properties) { + super( + properties.field, + properties.target, + properties.k, + properties.filter, + properties.iTopK, + properties.searchWidth); + requireConfiguredSearchWidth(properties.searchWidth); + } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + for (LeafReaderContext context : indexSearcher.getIndexReader().leaves()) { + requireCagraOnlyIndex(context); + } + return super.rewrite(indexSearcher); + } + + @Override + protected TopDocs approximateSearch( + LeafReaderContext context, + Bits acceptDocs, + int visitedLimit, + KnnCollectorManager knnCollectorManager) + throws IOException { + requireCagraOnlyIndex(context); + return super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager); + } + + @Override + protected TopDocs exactSearch( + LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout) { + throw new AssertionError("CAGRA search query must not use Lucene exact vector scoring"); + } + + private void requireConfiguredSearchWidth(int expectedSearchWidth) { + try { + Field searchWidthField = GPUKnnFloatVectorQuery.class.getDeclaredField("searchWidth"); + searchWidthField.setAccessible(true); + int actualSearchWidth = searchWidthField.getInt(this); + if (actualSearchWidth != expectedSearchWidth) { + throw new AssertionError( + "CAGRA search query expected retained searchWidth " + + expectedSearchWidth + + ", got " + + actualSearchWidth); + } + } catch (ReflectiveOperationException | RuntimeException exception) { + throw new AssertionError("Unable to inspect retained CAGRA searchWidth", exception); + } + } + + private void requireCagraOnlyIndex(LeafReaderContext context) { + KnnVectorsReader vectorsReader = vectorReaderForField(context, getField()); + if (!(vectorsReader instanceof CuVS2510GPUVectorsReader cuvsReader)) { + throw new AssertionError( + "CAGRA search query requires CuVS2510GPUVectorsReader for field '" + + getField() + + "', got " + + className(vectorsReader)); + } + + var fieldInfo = cuvsReader.getFieldInfos().fieldInfo(getField()); + if (fieldInfo == null) { + throw new AssertionError( + "CAGRA search field is absent from the cuVS reader: " + getField()); + } + var cuvsIndexes = cuvsReader.getCuvsIndexes(); + if (cuvsIndexes == null) { + throw new AssertionError( + "CAGRA search found no loaded GPU indexes for field: " + getField()); + } + GPUIndex gpuIndex = cuvsIndexes.get(fieldInfo.number); + if (gpuIndex == null) { + throw new AssertionError("CAGRA search found no GPU index for field: " + getField()); + } + var cagraIndex = gpuIndex.getCagraIndex(); + if (cagraIndex == null) { + throw new AssertionError( + "CAGRA search requires a CAGRA index for field '" + + getField() + + "'; a brute-force or CPU fallback is not allowed"); + } + if (gpuIndex.getBruteforceIndex() != null) { + throw new AssertionError( + "CAGRA search requires a CAGRA-only index for field '" + + getField() + + "'; a brute-force index was also loaded"); + } + + long actualGraphDegree = cagraIndex.getGraphDegree(); + if (actualGraphDegree != CAGRA_GRAPH_DEGREE) { + throw new AssertionError( + "CAGRA search expected actual graph degree " + + CAGRA_GRAPH_DEGREE + + " for field '" + + getField() + + "', got " + + actualGraphDegree); + } + } + } + + private static KnnVectorsReader vectorReaderForField( + LeafReaderContext context, String fieldName) { + LeafReader leafReader = FilterLeafReader.unwrap(context.reader()); + if (!(leafReader instanceof CodecReader codecReader)) { + throw new AssertionError( + "Graph-verifying query requires a CodecReader leaf, got " + + leafReader.getClass().getName()); + } + + KnnVectorsReader vectorsReader = codecReader.getVectorReader(); + if (vectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + vectorsReader = fieldsReader.getFieldReader(fieldName); + } + return vectorsReader; + } + + private static String className(Object value) { + return value == null ? "" : value.getClass().getName(); + } + + private static final class QueryProperties { + + private final String field; + private final float[] target; + private final int k; + private final int iTopK; + private final int searchWidth; + private final Query filter; + + private QueryProperties( + String field, float[] target, int k, int iTopK, int searchWidth, Query filter) { + this.field = field; + this.target = target; + this.k = k; + this.iTopK = iTopK; + this.searchWidth = searchWidth; + this.filter = filter; + } + + private static QueryProperties fromSystemProperties() { + String field = requiredProperty(QUERY_FIELD_PROPERTY); + float[] target = parseTarget(requiredProperty(QUERY_TARGET_PROPERTY)); + int k = parsePositiveInt(QUERY_K_PROPERTY); + int iTopK = parsePositiveInt(QUERY_I_TOP_K_PROPERTY); + int searchWidth = parsePositiveInt(QUERY_SEARCH_WIDTH_PROPERTY); + Query filter = optionalFilterFromSystemProperties(); + if (iTopK < k) { + throw new IllegalArgumentException( + "Java system property " + + QUERY_I_TOP_K_PROPERTY + + " must be greater than or equal to " + + QUERY_K_PROPERTY + + " (" + + k + + "), got: " + + iTopK); + } + return new QueryProperties(field, target, k, iTopK, searchWidth, filter); + } + + private static QueryProperties fromSystemPropertiesForHnsw() { + String field = requiredProperty(QUERY_FIELD_PROPERTY); + float[] target = parseTarget(requiredProperty(QUERY_TARGET_PROPERTY)); + int k = parsePositiveInt(QUERY_K_PROPERTY); + Query filter = optionalFilterFromSystemProperties(); + return new QueryProperties(field, target, k, k, 1, filter); + } + } + + private static Query optionalFilterFromSystemProperties() { + String field = System.getProperty(QUERY_FILTER_FIELD_PROPERTY); + String value = System.getProperty(QUERY_FILTER_VALUE_PROPERTY); + if (field == null && value == null) { + return null; + } + if (field == null || value == null) { + throw new IllegalArgumentException( + "Java system properties " + + QUERY_FILTER_FIELD_PROPERTY + + " and " + + QUERY_FILTER_VALUE_PROPERTY + + " must be set together"); + } + return new TermQuery( + new Term( + requiredProperty(QUERY_FILTER_FIELD_PROPERTY), + requiredProperty(QUERY_FILTER_VALUE_PROPERTY))); + } + + private static String requiredProperty(String propertyName) { + String value = System.getProperty(propertyName); + if (value == null || value.trim().isEmpty()) { + throw new IllegalArgumentException("Missing required Java system property: " + propertyName); + } + return value.trim(); + } + + private static int parsePositiveInt(String propertyName) { + String value = requiredProperty(propertyName); + final int parsed; + try { + parsed = Integer.parseInt(value); + } catch (NumberFormatException exception) { + throw new IllegalArgumentException( + "Java system property " + + propertyName + + " must be a positive integer, got: '" + + value + + "'", + exception); + } + if (parsed <= 0) { + throw new IllegalArgumentException( + "Java system property " + propertyName + " must be a positive integer, got: " + parsed); + } + return parsed; + } + + private static float[] parseTarget(String value) { + String[] values = value.split(",", -1); + float[] target = new float[values.length]; + for (int index = 0; index < values.length; index++) { + String component = values[index].trim(); + if (component.isEmpty()) { + throw malformedTarget(value, index, null); + } + try { + target[index] = Float.parseFloat(component); + } catch (NumberFormatException exception) { + throw malformedTarget(value, index, exception); + } + if (!Float.isFinite(target[index])) { + throw malformedTarget(value, index, null); + } + } + return target; + } + + private static IllegalArgumentException malformedTarget( + String value, int componentIndex, NumberFormatException cause) { + String message = + "Java system property " + + QUERY_TARGET_PROPERTY + + " must be a comma-separated list of finite floats; invalid component " + + componentIndex + + " in: '" + + value + + "'"; + return cause == null + ? new IllegalArgumentException(message) + : new IllegalArgumentException(message, cause); + } +}