diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java index e09457539..4d41feef3 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java @@ -199,11 +199,14 @@ public GraphSearcher build() { /** * @param scoreProvider provides functions to return the similarity of a given node to the query vector - * @param topK the number of results to look for. With threshold=0, the search will continue until at least - * `topK` results have been found, or until the entire graph has been searched. + * @param topK the number of results to look for. With threshold=Float.NEGATIVE_INFINITY, the search will + * continue until at least `topK` results have been found, or until the entire graph has been + * searched. * @param rerankK the number of (approximately-scored) results to rerank before returning the best `topK`. - * @param threshold the minimum similarity (0..1) to accept; 0 will accept everything. May be used - * with a large topK to find (approximately) all nodes above the given threshold. + * @param threshold the minimum similarity (0..1) to accept; 0 will accept everything for EUCLIDEAN and COSINE. + * DOT_PRODUCT is not bounded unless callers pre-normalize vectors to unit length and + * requires Float.NEGATIVE_INFINITY rather than 0. May be used with a large topK to find + * (approximately) all nodes above the given threshold. * If threshold > 0 then the search will stop when it is probabilistically unlikely * to find more nodes above the threshold, even if `topK` results have not yet been found. * @param rerankFloor (Experimental!) Candidates whose approximate similarity is at least this value @@ -251,8 +254,10 @@ public SearchResult search(SearchScoreProvider scoreProvider, * @param topK the number of results to look for. With threshold=0, the search will continue until at least * `topK` results have been found, or until the entire graph has been searched. * @param rerankK the number of (approximately-scored) results to rerank before returning the best `topK`. - * @param threshold the minimum similarity (0..1) to accept; 0 will accept everything. May be used - * with a large topK to find (approximately) all nodes above the given threshold. + * @param threshold the minimum similarity (0..1) to accept; 0 will accept everything for EUCLIDEAN and COSINE. + * DOT_PRODUCT is not bounded unless callers pre-normalize vectors to unit length and + * requires Float.NEGATIVE_INFINITY rather than 0. May be used with a large topK to find + * (approximately) all nodes above the given threshold. * If threshold > 0 then the search will stop when it is probabilistically unlikely * to find more nodes above the threshold, even if `topK` results have not yet been found. * @param acceptOrds a Bits instance indicating which nodes are acceptable results. @@ -271,8 +276,13 @@ protected void internalSearch(SearchScoreProvider scoreProvider, // Move downward from entry.level to 1 for (int lvl = entry.level; lvl > 0; lvl--) { - // Search this layer with minimal parameters since we just want the best candidate - searchOneLayer(scoreProvider, 1, 0.0f, lvl, Bits.ALL); + // Search this layer with minimal parameters since we just want the best candidate. + // Threshold must be NEGATIVE_INFINITY, not 0.0f: 0.0f is only a valid "accept everything" + // floor for similarity functions bounded to (0, 1] (COSINE, EUCLIDEAN). DOT_PRODUCT with + // non-unit-normalized vectors can legitimately score below 0, and a 0.0f floor here would + // silently discard the only reachable candidate(s) in a hierarchy layer, leaving + // approximateResults empty and tripping the assertion below. + searchOneLayer(scoreProvider, 1, Float.NEGATIVE_INFINITY, lvl, Bits.ALL); assert approximateResults.size() == 1 : approximateResults.size(); setEntryPointsFromPreviousLayer(); } @@ -285,8 +295,10 @@ protected void internalSearch(SearchScoreProvider scoreProvider, * @param scoreProvider provides functions to return the similarity of a given node to the query vector * @param topK the number of results to look for. With threshold=0, the search will continue until at least * `topK` results have been found, or until the entire graph has been searched. - * @param threshold the minimum similarity (0..1) to accept; 0 will accept everything. May be used - * with a large topK to find (approximately) all nodes above the given threshold. + * @param threshold the minimum similarity (0..1) to accept; 0 will accept everything for EUCLIDEAN and COSINE. + * DOT_PRODUCT is not bounded unless callers pre-normalize vectors to unit length and + * requires Float.NEGATIVE_INFINITY rather than 0. May be used with a large topK to find + * (approximately) all nodes above the given threshold. * If threshold > 0 then the search will stop when it is probabilistically unlikely * to find more nodes above the threshold, even if `topK` results have not yet been found. * @param acceptOrds a Bits instance indicating which nodes are acceptable results. @@ -317,7 +329,7 @@ public SearchResult search(SearchScoreProvider scoreProvider, int topK, Bits acceptOrds) { - return search(scoreProvider, topK, 0.0f, acceptOrds); + return search(scoreProvider, topK, Float.NEGATIVE_INFINITY, acceptOrds); } @Experimental diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestHierarchyDotProductSearch.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestHierarchyDotProductSearch.java new file mode 100644 index 000000000..edf5ac395 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestHierarchyDotProductSearch.java @@ -0,0 +1,90 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; + +/** + * Reproduces cndb#18237: hierarchical + * search with {@link VectorSimilarityFunction#DOT_PRODUCT} and non-unit-normalized vectors can + * throw {@code AssertionError: 0} from {@link GraphSearcher#internalSearch}. + * + *

{@code DOT_PRODUCT} similarity is {@code (1 + dotProduct(v1, v2)) / 2}, which is only + * guaranteed to land in {@code (0, 1]} when both vectors are unit-length. With larger-magnitude, + * non-normalized vectors (the deliberate "true dot product" use case this test mirrors, taken + * from Cassandra's {@code VectorDotProductWithLengthTest#testTrueDotproduct}), the raw score can + * legitimately be negative. {@code internalSearch}'s hierarchy-descent loop searches each upper + * layer for a single best candidate using a hardcoded {@code 0.0f} threshold; if every candidate + * reachable from the entry point in that layer scores below zero for a given query, all of them + * are filtered out, leaving {@code approximateResults} empty and tripping the + * {@code assert approximateResults.size() == 1} a few lines later. + * + *

This only manifests with hierarchy enabled: layer-0 search uses a much larger candidate + * pool and generally still finds enough positively-scored neighbors, but the hierarchy-descent + * layers search with {@code rerankK=1} from a single fixed entry point, making it far more + * likely that every candidate examined happens to score negative for an unlucky query direction. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestHierarchyDotProductSearch extends LuceneTestCase { + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + // Mirrors Cassandra's create2DVector(): low dimensionality and a wide magnitude range so + // dot products routinely fall well outside [-1, 1], unlike unit-normalized test vectors. + private VectorFloat randomUnnormalizedVector() { + float x = getRandom().nextFloat() * 200f - 100f; + float y = getRandom().nextFloat() * 200f - 100f; + return vts.createFloatVector(new float[]{x, y}); + } + + @Test + public void testHierarchySearchWithNonUnitDotProductVectors() { + int nDoc = 2000; + int dimension = 2; + + List> vectors = new ArrayList<>(nDoc); + for (int i = 0; i < nDoc; i++) { + vectors.add(randomUnnormalizedVector()); + } + var ravv = new ListRandomAccessVectorValues(vectors, dimension); + + // addHierarchy=true is required to reach the buggy code path; M/beamWidth mirror + // typical defaults, not anything specific to reproduction. + var builder = new GraphIndexBuilder(ravv, VectorSimilarityFunction.DOT_PRODUCT, + 16, 100, 1.2f, 1.2f, true); + var graph = TestUtil.buildSequentially(builder, ravv); + + // Enough queries that, pre-fix, at least one is very likely to hit a hierarchy layer + // where every reachable candidate scores negative against it. + for (int q = 0; q < 200; q++) { + VectorFloat query = randomUnnormalizedVector(); + // Pre-fix, this throws AssertionError(0) from GraphSearcher.internalSearch (line 263) + // via searchOneLayer's hardcoded 0.0f threshold during hierarchy descent. + GraphSearcher.search(query, 10, ravv.copy(), VectorSimilarityFunction.DOT_PRODUCT, graph, Bits.ALL); + } + } +}