From a9aeb4407e3b39956254a076e48d1017535a033f Mon Sep 17 00:00:00 2001 From: vic Date: Thu, 1 Oct 2026 14:27:19 +0200 Subject: [PATCH 1/2] Fix flaky accelerated HNSW gaps test --- .../lucene/TestCuVSAcceleratedHNSWGaps.java | 41 ++++--------------- 1 file changed, 9 insertions(+), 32 deletions(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java index fa0936e67e..7120b227f2 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java @@ -8,10 +8,7 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import java.io.IOException; -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; @@ -116,7 +113,8 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + int vectorCount = (datasetSize + 1) / 2; + int topK = Math.min(random.nextInt(5, TOP_K_LIMIT), vectorCount); Query query = new KnnFloatVectorQuery("vector", queryVector, topK); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; @@ -129,17 +127,16 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { String docId = reader.storedFields().document(hit.doc).get("id"); int id = Integer.parseInt(docId); assertEquals("All results should be even-numbered (have vectors)", 0, id % 2); + float expectedScore = VectorSimilarityFunction.EUCLIDEAN.compare(queryVector, dataset[id]); + assertEquals( + "Score should match the vector for document " + id, + expectedScore, + hit.score, + expectedScore * 1e-4f); log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); } - // Verify the results match expected top-k based on Euclidean distance - List expectedIds = calculateExpectedTopK(queryVector, topK, dataset); - for (int i = 0; i < hits.length; i++) { - String docId = reader.storedFields().document(hits[i].doc).get("id"); - int id = Integer.parseInt(docId); - assertTrue("Result " + id + " should be in expected top-k results", expectedIds.contains(id)); - } - + // HNSW search is approximate, so exact Euclidean top-k membership is not guaranteed. log.log(Level.FINE, "Alternating document test passed with " + hits.length + " results"); } @@ -168,24 +165,4 @@ public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOExcepti Level.FINE, "Filtered alternating document test passed with " + filteredHits.length + " results"); } - - public static List calculateExpectedTopK(float[] query, int topK, float[][] dataset) { - Map distances = new TreeMap<>(); - - // Calculate distances only for documents that have vectors (even-numbered) - for (int i = 0; i < dataset.length; i += 2) { - double distance = 0; - for (int j = 0; j < dataset[0].length; j++) { - distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); - } - distances.put(i, distance); - } - - // Sort by distance and return top-k - return distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .limit(topK) - .toList(); - } } From 4f756080c7beec06c698d8938638b31115a49e2b Mon Sep 17 00:00:00 2001 From: vic Date: Tue, 6 Oct 2026 15:45:12 +0200 Subject: [PATCH 2/2] address feedback --- .../lucene/TestCuVSAcceleratedHNSWGaps.java | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java index 7120b227f2..cac5cc237f 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java @@ -8,7 +8,10 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import java.io.IOException; +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; @@ -118,6 +121,7 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { Query query = new KnnFloatVectorQuery("vector", queryVector, topK); ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + List acceptableIds = calculateAcceptableNeighbors(queryVector, topK, dataset); // Verify we get exactly TOP_K results assertEquals("Should return exactly " + topK + " results", topK, hits.length); @@ -133,10 +137,12 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { expectedScore, hit.score, expectedScore * 1e-4f); + assertTrue( + "Result " + id + " was not among the closest " + (topK * 3) + " documents", + acceptableIds.contains(id)); log.log(Level.FINE, "Document ID: " + id + ", Score: " + hit.score); } - // HNSW search is approximate, so exact Euclidean top-k membership is not guaranteed. log.log(Level.FINE, "Alternating document test passed with " + hits.length + " results"); } @@ -165,4 +171,25 @@ public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOExcepti Level.FINE, "Filtered alternating document test passed with " + filteredHits.length + " results"); } + + private static List calculateAcceptableNeighbors( + float[] query, int topK, float[][] dataset) { + Map distances = new TreeMap<>(); + + // Only even-numbered documents have vectors. + for (int i = 0; i < dataset.length; i += 2) { + double distance = 0; + for (int j = 0; j < query.length; j++) { + distance += (query[j] - dataset[i][j]) * (query[j] - dataset[i][j]); + } + distances.put(i, distance); + } + + // HNSW search is approximate, so accept results within three times the requested K. + return distances.entrySet().stream() + .sorted(Map.Entry.comparingByValue()) + .limit(topK * 3L) + .map(Map.Entry::getKey) + .toList(); + } }