Skip to content

Use a recall-based check in TestCuVSRandomizedHNSWVectorSearch #2824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 1 commit into from
Oct 14, 2026
Merged

Use a recall-based check in TestCuVSRandomizedHNSWVectorSearch #2824

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 17 additions & 60 deletions .../cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,13 @@
*/
package com.nvidia.cuvs.lucene;

import static com.nvidia.cuvs.lucene.TestUtils.assertSearchFindsNearestNeighbors;
import static com.nvidia.cuvs.lucene.TestUtils.generateDataset;
import static com.nvidia.cuvs.lucene.TestUtils.generateQueries;
import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.TreeMap;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.apache.lucene.codecs.Codec;
Expand Down Expand Up @@ -57,7 +53,9 @@ public class TestCuVSRandomizedHNSWVectorSearch extends LuceneTestCase {
static int DATASET_SIZE_LIMIT = 1000;
static int DIMENSIONS_LIMIT = 2048;
static int NUM_QUERIES_LIMIT = 10;
static int TOP_K_LIMIT = 64; // TODO This fails beyond 64
// A wider search is more accurate, not less: in the measurements described in testVectorSearch,
// no result was outside the closest 3 * topK for any topK from 8 to 200.
static int TOP_K_LIMIT = 128;
Comment on lines +56 to +58

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The old TODO said "This fails beyond 64" (main). I couldn't find a failure there: in the #2819 sweep, topK = 89, 128 and 200 had 0 hits outside 3 * topK, mean recall was 1.000, and the worst per-query recall was 0.989 (comment). testVectorSearchWithFilter also uses this limit, but its filter matches a single document, so a larger topK doesn't change what it checks (L163-L177).

Question for reviewers: the TODO doesn't say what failed beyond 64. If anyone remembers the original failure, please check whether this change would bring it back.

static float[][] dataset;
// Whether each document of the dataset was indexed with a vector; only those can be hits.
static boolean[] hasVector;
Expand Down Expand Up @@ -121,66 +119,25 @@ public void testVectorSearch() throws IOException {
int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1;
int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length);

if (dataset.length < topK) topK = dataset.length;

float[][] queries = generateQueries(random, dataset[0].length, numQueries);
List<List<Integer>> expected = generateExpectedResults(topK, dataset, queries);

log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length);
log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length);
log.log(Level.FINE, "TopK: " + topK);

Query query = new KnnFloatVectorQuery("vector", queries[0], topK);
int correct[] = new int[topK];
for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i);

ScoreDoc[] hits = searcher.search(query, topK).scoreDocs;
log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits));
log.log(Level.FINE, "EXPECTD: " + expected.get(0));

for (ScoreDoc hit : hits) {
log.log(
Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score);
}

for (ScoreDoc hit : hits) {
int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id"));
assertTrue("Result returned was not in topK*3: " + doc, expected.get(0).contains(doc));
}
}

private static List<List<Integer>> generateExpectedResults(
int topK, float[][] dataset, float[][] queries) {
List<List<Integer>> neighborsResult = new ArrayList<>();
int dimensions = dataset[0].length;

for (float[] query : queries) {
Map<Integer, Double> distances = new TreeMap<>();
for (int j = 0; j < dataset.length; j++) {
if (hasVector[j] == false) {
continue;
}
double distance = 0;
for (int k = 0; k < dimensions; k++) {
distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]);
}
distances.put(j, (distance));
}

Map<Integer, Double> sorted = new TreeMap<Integer, Double>(distances);
log.log(Level.FINE, "EXPECTED: " + sorted);

// Sort by distance and select the topK nearest neighbors
List<Integer> neighbors =
distances.entrySet().stream()
.sorted(Map.Entry.comparingByValue())
.map(Map.Entry::getKey)
.toList();
neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, neighbors.size())));
}

log.log(Level.FINE, "Expected results generated successfully.");
return neighborsResult;
// The search is approximate. Lucene searches each segment's HNSW graph with a beam as wide as
// topK, and in this data (uniform random vectors with up to 2048 dimensions, so all distances
// are within about 10% of each other) a narrow beam sometimes stops just short of the true
// neighbors. Over 27,003 random indexes from this test's distribution, with 10 queries each, a
// result fell outside the closest 3 * topK in 1.7% of queries at topK = 1, 0.21% at topK = 2,
// 0.03% at topK = 3, 3 times in 270,000 queries at topK = 5, and never at topK >= 8. Those
// results were at most rank 17, and typically only 1.3% farther away than the nearest
// neighbor. On the same segments, Lucene's own HNSW format missed in 22% of queries at
// topK = 1. Allowing 2 misses across all the queries keeps the false failure rate below 1e-5
// per run.
int allowedMisses = 2;
assertSearchFindsNearestNeighbors(
searcher, "vector", "id", dataset, hasVector, queries, topK, allowedMisses);
Comment on lines +128 to +140

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This replaces the per-hit check on query 0 only (main) with a budget over all the queries. The numbers in the comment come from the #2819 sweeps (before, after). The new check didn't fail on any of the 16,553 sweep indexes, but it still failed on Lucene's own Lucene99HnswVectorsFormat built on the same segments (32 times with default parameters, 29 with maxConn = beamWidth = 32). So it can still tell a weaker graph from this one. #2775 uses the same helper with allowedMisses = 4 and 10 in the two TestCagraToHnswSerializationAndSearch* tests (example).

Question for reviewers: is this recall threshold appropriate? It allows 2 hits outside the closest 3 * topK, counted across all 1-10 queries of a run, no matter what topK is. I calibrated the budget for the small-topK cases where misses happen. A fixed count is lenient at large topK, where there are up to 128 results per query and the sweep never saw a miss. The alternatives I see are a mean-recall floor, or a budget that scales with numQueries * topK. The "below 1e-5 per run" figure is extrapolated from the per-query miss rates above. The direct measurement is 0 failures in 16,553 sweep indexes and 2166 test runs.

}

@Test
Expand Down
62 changes: 57 additions & 5 deletions java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,20 +90,62 @@ public static void assertSearchFindsNearestNeighbors(
int topK,
int allowedMisses)
throws IOException {
assertSearchFindsNearestNeighbors(
searcher, vectorField, idField, dataset, null, queries, topK, allowedMisses);
}

/**
* Like {@link #assertSearchFindsNearestNeighbors(IndexSearcher, String, String, float[][],
* float[][], int, int)}, for an index where only some rows of {@code dataset} were indexed with a
* vector: {@code indexed[i]} says whether row {@code i} was, and {@code null} means all of them.
* The closest vectors are ranked among the indexed rows only, since the others can never be
* returned.
*
* <p>Besides the approximate check, every query must return {@code min(topK, indexed rows)}
* distinct indexed documents in order of decreasing score, each scored by its own vector.
*/
public static void assertSearchFindsNearestNeighbors(
IndexSearcher searcher,
String vectorField,
String idField,
float[][] dataset,
boolean[] indexed,
float[][] queries,
int topK,
int allowedMisses)
throws IOException {
int indexedCount = dataset.length;
if (indexed != null) {
indexedCount = 0;
for (boolean isIndexed : indexed) {
if (isIndexed) indexedCount++;
}
}
int expectedResults = Math.min(topK, indexedCount);
Comment on lines +107 to +124

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test leaves the vector out of about 40% of the documents when the dataset has 100 or more rows (L88-L96), but the existing helper assumes every row is indexed. This overload takes the hasVector mask and ranks only the rows that have a vector, as the old generateExpectedResults did (main). The existing 7-argument overload delegates with null, so its callers in the TestCagraToHnswSerializationAndSearch* tests behave as before.

The expected count becomes min(topK, indexed rows) instead of topK, because a small dataset can have fewer indexed rows than topK. The old test didn't check the count at all.

int acceptableRank = topK * 3;
int misses = 0;
StringBuilder found = new StringBuilder();
StoredFields storedFields = searcher.storedFields();
for (int q = 0; q < queries.length; q++) {
KnnFloatVectorQuery query = new KnnFloatVectorQuery(vectorField, queries[q], topK);
TopDocs results = searcher.search(query, topK);
assertEquals("results for query " + q, topK, results.scoreDocs.length);
assertEquals("results for query " + q, expectedResults, results.scoreDocs.length);

int[] ranks = distanceRanks(dataset, queries[q]);
int[] ids = new int[topK];
int[] resultRanks = new int[topK];
for (int i = 0; i < topK; i++) {
int[] ranks = distanceRanks(dataset, indexed, queries[q]);
int[] ids = new int[expectedResults];
int[] resultRanks = new int[expectedResults];
Set<Integer> seen = new HashSet<>();
for (int i = 0; i < expectedResults; i++) {
ids[i] = Integer.parseInt(storedFields.document(results.scoreDocs[i].doc).get(idField));
assertTrue(
"id " + ids[i] + " for query " + q + " has no vector",
indexed == null || indexed[ids[i]]);
assertTrue("id " + ids[i] + " returned twice for query " + q, seen.add(ids[i]));
if (i > 0) {
assertTrue(
"results for query " + q + " are not sorted by score at position " + i,
results.scoreDocs[i].score <= results.scoreDocs[i - 1].score);
}
float expectedScore =
VectorSimilarityFunction.EUCLIDEAN.compare(queries[q], dataset[ids[i]]);
assertEquals(
Comment on lines +140 to 151

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are the hard checks, and they don't count against the miss budget: every hit has a vector, there are no duplicates, scores don't increase, and each score equals the exact distance to that document's own vector. An ordinal-to-docID mapping error or a corrupted graph would break one of them, while a legitimate ANN miss wouldn't. All of them held in every #2819 sweep (0 violations, comment).

Expand Down Expand Up @@ -142,6 +184,14 @@ public static void assertSearchFindsNearestNeighbors(
* @return the 1-based rank of each vector, indexed like {@code dataset}
*/
public static int[] distanceRanks(float[][] dataset, float[] query) {
return distanceRanks(dataset, null, query);
}

/**
* Like {@link #distanceRanks(float[][], float[])}, ranking only the rows where {@code indexed} is
* true ({@code null} means all rows). The other rows get rank {@link Integer#MAX_VALUE}.
*/
public static int[] distanceRanks(float[][] dataset, boolean[] indexed, float[] query) {
double[] distances = new double[dataset.length];
for (int i = 0; i < dataset.length; i++) {
for (int j = 0; j < query.length; j++) {
Expand All @@ -151,11 +201,13 @@ public static int[] distanceRanks(float[][] dataset, float[] query) {
}
int[] byDistance =
IntStream.range(0, dataset.length)
.filter(i -> indexed == null || indexed[i])
.boxed()
.sorted(Comparator.comparingDouble(i -> distances[i]))
.mapToInt(Integer::intValue)
.toArray();
int[] ranks = new int[dataset.length];
Arrays.fill(ranks, Integer.MAX_VALUE);
for (int rank = 0; rank < byDistance.length; rank++) {
ranks[byDistance[rank]] = rank + 1;
}
Expand Down
Loading