Expand PyLucene CAGRA search coverage

Author: nvzm123

README.md

@@ -84,7 +84,7 @@ directly against existing artifacts instead:
 ```sh
 CUVS_LUCENE_JAR=/path/to/cuvs-lucene.jar \
 CUVS_LUCENE_CUVS_JAVA_JAR=/path/to/cuvs-java.jar \
-python3 -m pytest -q -s examples/Python/test_pylucene_smoke.py
+python3 -m pytest -q -s src/test/python/test_pylucene_smoke.py
 ```
 
 To run an expanded GPU end-to-end pytest suite through CPU HNSW,
@@ -94,9 +94,10 @@ CAGRA-to-HNSW, and CAGRA search paths:
 ./test_pylucene.sh --gpu-e2e
 ```
 
-The expanded suite runs the `gpu-basic`, `gpu-segments`, `cpu-hnsw`, and
-`cagra-hnsw` case groups. The basic cases cover `hnsw`, `cagra`, `hnsw-single`,
-and `cagra-single`. The segment cases cover 1-segment indexes, 10-segment
+The expanded suite runs the `gpu-basic`, `gpu-segments`, `cagra-search`,
+`cpu-hnsw`, and `cagra-hnsw` case groups. The basic cases cover `hnsw`,
+`cagra-search-default`, `hnsw-single`, and `cagra-search-single`. The segment cases
+cover 1-segment indexes, 10-segment
 indexes, 10 segments force-merged to 1, and 100 segments force-merged to 10 for
 both HNSW and CAGRA. The CPU HNSW cases force the accelerated HNSW codec through
 its Lucene CPU fallback path in the same run, including 10 segments force-merged
@@ -107,13 +108,31 @@ documents and 32 dimensions; high-segment cases use at least 257 rows per
 segment to avoid expected cuVS graph-degree clamps on tiny per-segment datasets.
 The suite checks Lucene SPI discovery, jar packaging, index file suffixes
 (`.vex`/`.vem` for HNSW and `.vcag`/`.vemc` for CAGRA), indexed vector metadata,
-unfiltered KNN, filtered KNN, missing-vector documents, deletions, and force
+unfiltered KNN, filtered KNN, exact self-match at rank 1, duplicate result
+checks, missing-vector documents, deletions, and force
 merge behavior. To run a subset or resize the test:
 
 ```sh
 ./test_pylucene.sh --gpu-e2e --cases=gpu-segments --rows=5000 --dims=64 --topk=20
 ```
 
+To verify that the selected pytest cases can fail quickly without initializing
+the JVM or touching the GPU:
+
+```sh
+MAKE_FAIL=true ./test_pylucene.sh --gpu-e2e --cases=all --no-build
+```
+
+The PyLucene tests live under `src/test/python` rather than `examples` so they
+follow the same test-vs-example separation used in cuVS/RAPIDS. cuVS keeps
+language-specific tests close to their owning language package or build system:
+Java integration tests live under `java/cuvs-java/src/test/java` and run through
+Maven/Failsafe, Python tests live under `python/cuvs/cuvs/tests` and run through
+pytest, and C/C++ tests live under `c/tests` and `cpp/tests`. This repository is
+Java-first, so the PyLucene pytest suite is treated as test code under
+`src/test/python`, while `test_pylucene.sh` owns the Maven build, jar validation,
+classpath setup, and pytest invocation.
+
 ### Running Tests
 
 ```sh

examples/Python/pylucene_smoke.py src/test/python/pylucene_smoke.pyexamples/Python/pylucene_smoke.py renamed to src/test/python/pylucene_smoke.py

@@ -7,10 +7,8 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-import lucene
 
-
-REPO_ROOT = Path(__file__).resolve().parents[2]
+REPO_ROOT = Path(__file__).resolve().parents[3]
 HNSW_CODEC = "Lucene101AcceleratedHNSWCodec"
 CAGRA_HNSW_BASE_LAYER_CODEC = "Lucene101AcceleratedHNSWBaseLayerCodec"
 CAGRA_HNSW_MULTI_LAYER_CODEC = "Lucene101AcceleratedHNSWMultiLayerCodec"
@@ -28,16 +26,29 @@
     "Lucene101AcceleratedHNSWBinaryQuantizedCodec",
     "Lucene101AcceleratedHNSWScalarQuantizedCodec",
 )
-BASIC_GPU_CASES = ("hnsw", "cagra", "hnsw-single", "cagra-single")
+CAGRA_SEARCH_CASES = (
+    "cagra-search-default",
+    "cagra-search-single",
+    "cagra-search-1seg",
+    "cagra-search-10seg",
+    "cagra-search-10seg-force-1",
+    "cagra-search-100seg-force-10",
+)
+BASIC_GPU_CASES = (
+    "hnsw",
+    "cagra-search-default",
+    "hnsw-single",
+    "cagra-search-single",
+)
 SEGMENT_GPU_CASES = (
     "hnsw-1seg",
-    "cagra-1seg",
+    "cagra-search-1seg",
     "hnsw-10seg",
-    "cagra-10seg",
+    "cagra-search-10seg",
     "hnsw-10seg-force-1",
-    "cagra-10seg-force-1",
+    "cagra-search-10seg-force-1",
     "hnsw-100seg-force-10",
-    "cagra-100seg-force-10",
+    "cagra-search-100seg-force-10",
 )
 CPU_HNSW_CASES = (
     "hnsw-cpu",
@@ -55,13 +66,15 @@
     "hnsw-cpu",
     "cagra-hnsw-1layer",
     "cagra-hnsw-3layer",
-    "cagra",
+    "cagra-search-default",
 )
 CASE_GROUPS = {
     "gpu": BASIC_GPU_CASES,
     "gpu-basic": BASIC_GPU_CASES,
     "gpu-segments": SEGMENT_GPU_CASES,
     "segments": SEGMENT_GPU_CASES,
+    "cagra-search": CAGRA_SEARCH_CASES,
+    "gpu-cagra-search": CAGRA_SEARCH_CASES,
     "cpu-hnsw": CPU_HNSW_CASES,
     "cagra-hnsw": CAGRA_HNSW_CASES,
     "algorithm-matrix": ALGORITHM_MATRIX_CASES,
@@ -206,6 +219,8 @@ def fvec(jarray, values):
 
 
 def init_vm(cuvs_java_jar, cuvs_lucene_jar):
+    import lucene
+
     java_library_path = os.environ.get("JAVA_LIBRARY_PATH") or os.environ.get(
         "LD_LIBRARY_PATH"
     )
@@ -385,9 +400,9 @@ def build_named_case(name):
             segment_count=3,
             force_merge_target=(1 if force_merge else 0),
         )
-    if name == "cagra":
+    if name in {"cagra", "cagra-search", "cagra-search-default"}:
         return matrix_case(
-            name="cagra",
+            name="cagra-search-default",
             codec_name=CAGRA_CODEC,
             expected_suffixes=(".vcag", ".vemc"),
             require_cuvs=True,
@@ -405,9 +420,9 @@ def build_named_case(name):
             expected_suffixes=(".vex", ".vem"),
             require_cuvs=require_cuvs,
         )
-    if name in {"cagra-single", "single-cagra"}:
+    if name in {"cagra-single", "single-cagra", "cagra-search-single"}:
         return SmokeCase(
-            name="cagra-single",
+            name="cagra-search-single",
             codec_name=CAGRA_CODEC,
             row_count=1,
             dims=matrix_dims,
@@ -417,20 +432,33 @@ def build_named_case(name):
         )
     if name == "hnsw-1seg":
         return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 1, 0)
-    if name == "cagra-1seg":
-        return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 1, 0)
+    if name in {"cagra-1seg", "cagra-search-1seg"}:
+        return build_segment_case(
+            "cagra-search-1seg", CAGRA_CODEC, (".vcag", ".vemc"), True, 1, 0
+        )
     if name == "hnsw-10seg":
         return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 10, 0)
-    if name == "cagra-10seg":
-        return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 10, 0)
+    if name in {"cagra-10seg", "cagra-search-10seg"}:
+        return build_segment_case(
+            "cagra-search-10seg", CAGRA_CODEC, (".vcag", ".vemc"), True, 10, 0
+        )
     if name == "hnsw-10seg-force-1":
         return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 10, 1)
-    if name == "cagra-10seg-force-1":
-        return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 10, 1)
+    if name in {"cagra-10seg-force-1", "cagra-search-10seg-force-1"}:
+        return build_segment_case(
+            "cagra-search-10seg-force-1", CAGRA_CODEC, (".vcag", ".vemc"), True, 10, 1
+        )
     if name == "hnsw-100seg-force-10":
         return build_segment_case(name, HNSW_CODEC, (".vex", ".vem"), require_cuvs, 100, 10)
-    if name == "cagra-100seg-force-10":
-        return build_segment_case(name, CAGRA_CODEC, (".vcag", ".vemc"), True, 100, 10)
+    if name in {"cagra-100seg-force-10", "cagra-search-100seg-force-10"}:
+        return build_segment_case(
+            "cagra-search-100seg-force-10",
+            CAGRA_CODEC,
+            (".vcag", ".vemc"),
+            True,
+            100,
+            10,
+        )
     if name in {"cagra-hnsw-1layer", "cagra-hnsw-base", "cagra-hnsw-base-layer"}:
         return matrix_case(
             name="cagra-hnsw-1layer",
@@ -767,22 +795,54 @@ def hit_ids(stored_fields, hits):
     return [stored_fields.document(hit.doc).get(ID_FIELD) for hit in hits]
 
 
-def assert_search_results(searcher, jarray, case, query_ids, inactive_ids):
+def squared_l2(left, right):
+    return sum((a - b) * (a - b) for a, b in zip(left, right))
+
+
+def exact_neighbor_doc_names(query_id, active_vector_ids, dims, limit):
+    query_vector = vector_for(query_id, dims)
+    ranked = sorted(
+        active_vector_ids,
+        key=lambda doc_id: (squared_l2(query_vector, vector_for(doc_id, dims)), doc_id),
+    )
+    return [f"doc-{doc_id}" for doc_id in ranked[:limit]]
+
+
+def assert_search_results(searcher, jarray, case, query_ids, active_vector_ids, inactive_ids):
     from org.apache.lucene.search import KnnFloatVectorQuery
 
     stored_fields = searcher.storedFields()
-    top_k = min(case.top_k, max(1, case.row_count - len(inactive_ids)))
+    top_k = min(case.top_k, max(1, len(active_vector_ids)))
     for query_id in query_ids:
         query = KnnFloatVectorQuery(
             VECTOR_FIELD, fvec(jarray, vector_for(query_id, case.dims)), top_k
         )
         ids = hit_ids(stored_fields, searcher.search(query, top_k).scoreDocs)
         expected = f"doc-{query_id}"
-        if expected not in ids:
-            raise AssertionError(f"{case.name}: expected {expected} in top {top_k}, got {ids}")
+        if len(ids) != top_k:
+            raise AssertionError(
+                f"{case.name}: expected {top_k} search results, got {len(ids)}: {ids}"
+            )
+        if len(ids) != len(set(ids)):
+            raise AssertionError(f"{case.name}: duplicate docs returned: {ids}")
+        if ids[0] != expected:
+            raise AssertionError(f"{case.name}: expected {expected} at rank 1, got {ids}")
         bad_ids = [doc_id for doc_id in ids if doc_id in inactive_ids]
         if bad_ids:
             raise AssertionError(f"{case.name}: inactive docs returned: {bad_ids}")
+        expected_candidates = set(
+            exact_neighbor_doc_names(
+                query_id, active_vector_ids, case.dims, min(len(active_vector_ids), top_k * 2 + 10)
+            )
+        )
+        unexpected_ids = [
+            doc_id for doc_id in ids[: min(5, len(ids))] if doc_id not in expected_candidates
+        ]
+        if unexpected_ids:
+            raise AssertionError(
+                f"{case.name}: nearest hits {unexpected_ids} were outside the exact "
+                f"top {len(expected_candidates)} candidates for {expected}"
+            )
 
 
 def assert_filtered_search(searcher, jarray, case, query_id):
@@ -906,7 +966,9 @@ def run_case(case, codec_class, codec_cache, jarray):
             assert_vector_metadata(reader, VECTOR_FIELD, expected_vector_count, case.dims)
             assert_segment_topology(reader, case)
             searcher = IndexSearcher(reader)
-            assert_search_results(searcher, jarray, case, query_ids, inactive_doc_names)
+            assert_search_results(
+                searcher, jarray, case, query_ids, active_vector_ids, inactive_doc_names
+            )
             assert_filtered_search(searcher, jarray, case, query_ids[0])
             telemetry = writer_telemetry(codec)
             writer_path = observed_writer_path(case, telemetry)

examples/Python/test_pylucene_smoke.py src/test/python/test_pylucene_smoke.pyexamples/Python/test_pylucene_smoke.py renamed to src/test/python/test_pylucene_smoke.py

@@ -1,11 +1,17 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+
 import pytest
 
 from pylucene_smoke import SELECTED_CASES, initialize_pylucene_context, run_case
 
 
+def make_fail_requested():
+    return os.environ.get("MAKE_FAIL", "false").lower() in {"1", "true", "yes", "on"}
+
+
 @pytest.fixture(scope="session")
 def pylucene_context(request):
     context = initialize_pylucene_context()
@@ -20,7 +26,11 @@ def pylucene_context(request):
 
 
 @pytest.mark.parametrize("case", SELECTED_CASES, ids=lambda case: case.name)
-def test_pylucene_smoke_case(pylucene_context, case):
+def test_pylucene_smoke_case(request, case):
+    if make_fail_requested():
+        pytest.fail(f"MAKE_FAIL=true requested failure for {case.name}")
+
+    pylucene_context = request.getfixturevalue("pylucene_context")
     run_case(
         case,
         pylucene_context.codec_class,

test_pylucene.sh

@@ -17,6 +17,13 @@ PYLUCENE_ROWS="${CUVS_LUCENE_PYLUCENE_ROWS:-}"
 PYLUCENE_DIMS="${CUVS_LUCENE_PYLUCENE_DIMS:-}"
 PYLUCENE_TOPK="${CUVS_LUCENE_PYLUCENE_TOPK:-}"
 CUVS_LUCENE_JAR_PATH="${CUVS_LUCENE_JAR:-}"
+MAKE_FAIL_VALUE="${MAKE_FAIL:-false}"
+MAKE_FAIL_ENABLED=0
+case "${MAKE_FAIL_VALUE}" in
+  1|true|TRUE|True|yes|YES|Yes|on|ON|On)
+    MAKE_FAIL_ENABLED=1
+    ;;
+esac
 
 for arg in "$@"; do
   case "${arg}" in
@@ -44,13 +51,16 @@ for arg in "$@"; do
       echo "Builds and checks the standard cuvs-lucene jar, then runs the PyLucene pytest suite."
       echo "Set CUVS_LUCENE_JAR to test an existing cuvs-lucene jar."
       echo "Set CUVS_LUCENE_CUVS_JAVA_JAR to the base cuvs-java jar if it is not in ~/.m2."
+      echo "Set MAKE_FAIL=true to make every selected pytest case fail before JVM setup."
       echo "Set PYTHON or MVN to override the Python or Maven executable."
       echo
-      echo "Case groups: gpu-basic, gpu-segments, cpu-hnsw, cagra-hnsw, algorithm-matrix, all."
-      echo "Core cases: smoke, hnsw, cagra, hnsw-single, cagra-single."
-      echo "Segment cases: hnsw-1seg, cagra-1seg, hnsw-10seg, cagra-10seg,"
-      echo "  hnsw-10seg-force-1, cagra-10seg-force-1,"
-      echo "  hnsw-100seg-force-10, cagra-100seg-force-10."
+      echo "Case groups: gpu-basic, gpu-segments, cagra-search, cpu-hnsw,"
+      echo "  cagra-hnsw, algorithm-matrix, all."
+      echo "Core cases: smoke, hnsw, cagra-search-default, hnsw-single, cagra-search-single."
+      echo "Segment cases: hnsw-1seg, cagra-search-1seg, hnsw-10seg,"
+      echo "  cagra-search-10seg, hnsw-10seg-force-1,"
+      echo "  cagra-search-10seg-force-1, hnsw-100seg-force-10,"
+      echo "  cagra-search-100seg-force-10."
       echo "CPU HNSW cases: hnsw-cpu, hnsw-cpu-single, hnsw-cpu-1seg,"
       echo "  hnsw-cpu-10seg, hnsw-cpu-10seg-force-1,"
       echo "  hnsw-cpu-100seg-force-10."
@@ -259,11 +269,13 @@ for provider in \
   }
 done
 
-"${PYTHON_BIN}" -c "import lucene" >/dev/null 2>&1 || {
-  echo "Python cannot import PyLucene's lucene module." >&2
-  echo "Activate or install a PyLucene environment compatible with this project's Lucene version." >&2
-  exit 1
-}
+if [[ "${MAKE_FAIL_ENABLED}" -eq 0 ]]; then
+  "${PYTHON_BIN}" -c "import lucene" >/dev/null 2>&1 || {
+    echo "Python cannot import PyLucene's lucene module." >&2
+    echo "Activate or install a PyLucene environment compatible with this project's Lucene version." >&2
+    exit 1
+  }
+fi
 
 "${PYTHON_BIN}" -m pytest --version >/dev/null 2>&1 || {
   echo "Python cannot run pytest." >&2
@@ -274,6 +286,7 @@ done
 smoke_env=(
   "CUVS_LUCENE_JAR=${cuvs_lucene_jar_abs}"
   "CUVS_LUCENE_CUVS_JAVA_JAR=${cuvs_java_jar}"
+  "MAKE_FAIL=${MAKE_FAIL_VALUE}"
 )
 
 if [[ "${GPU_E2E}" -eq 1 ]]; then
@@ -301,4 +314,4 @@ else
   smoke_env+=("CUVS_LUCENE_VERIFY_ALL_CODECS=${CUVS_LUCENE_VERIFY_ALL_CODECS:-1}")
 fi
 
-env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s examples/Python/test_pylucene_smoke.py
+env "${smoke_env[@]}" "${PYTHON_BIN}" -m pytest -q -s src/test/python/test_pylucene_smoke.py