kaivalnp commented on code in PR #15784:
URL: https://github.com/apache/lucene/pull/15784#discussion_r2983356253
##########
lucene/core/src/java/org/apache/lucene/search/VectorSimilarityCollector.java:
##########
@@ -20,49 +20,80 @@
import java.util.List;
/**
- * Perform a similarity-based graph search.
+ * Perform a similarity-based graph search to find all (approximate) vectors
above a similarity
+ * threshold.
+ *
+ * <p>The buffer for graph traversal is adaptive: starts with a high value,
and decays towards
+ * scores of nodes traversed but not collected, with a provided factor. The
decay factor should lie
+ * in {@code [0, 1]}; with higher values producing better recall using more
graph exploration.
+ *
+ * <p>Note: Some functions of this class deviate from {@link KnnCollector},
and should be used with
+ * queries that are aware of the differences (like {@link
ByteVectorSimilarityQuery} and {@link
+ * FloatVectorSimilarityQuery}).
+ *
+ * <ul>
+ * <li>{@link #k()} does NOT provide a good estimate of the number of
collected results.
+ * <li>{@link #topDocs()} does NOT return docs sorted in descending order of
scores.
+ * <li>{@link #collect(int, float)} does NOT return true / false based on
whether the document was
+ * collected.
+ * </ul>
*
* @lucene.experimental
*/
class VectorSimilarityCollector extends AbstractKnnCollector {
- private final float traversalSimilarity, resultSimilarity;
- private float maxSimilarity;
+ private final float resultSimilarity, decay;
private final List<ScoreDoc> scoreDocList;
+ private float minCompetitiveSimilarity;
/**
- * Perform a similarity-based graph search. The graph is traversed till
better scoring nodes are
- * available, or the best candidate is below {@link #traversalSimilarity}.
All traversed nodes
- * above {@link #resultSimilarity} are collected.
+ * Perform a similarity-based graph search.
*
- * @param traversalSimilarity (lower) similarity score for graph traversal.
- * @param resultSimilarity (higher) similarity score for result collection.
+ * @param resultSimilarity similarity score for result collection.
+ * @param decay decay factor for graph traversal buffer.
* @param visitLimit limit on number of nodes to visit.
*/
- public VectorSimilarityCollector(
- float traversalSimilarity, float resultSimilarity, long visitLimit) {
+ public VectorSimilarityCollector(float resultSimilarity, float decay, long
visitLimit) {
// TODO: add search strategy support
super(1, visitLimit, AbstractVectorSimilarityQuery.DEFAULT_STRATEGY);
- if (traversalSimilarity > resultSimilarity) {
- throw new IllegalArgumentException("traversalSimilarity should be <=
resultSimilarity");
+
+ if (decay < 0 || decay > 1) {
+ throw new IllegalArgumentException("decay must lie in range [0,1]; got "
+ decay);
}
- this.traversalSimilarity = traversalSimilarity;
+
this.resultSimilarity = resultSimilarity;
- this.maxSimilarity = Float.NEGATIVE_INFINITY;
+ this.decay = decay;
this.scoreDocList = new ArrayList<>();
+ this.minCompetitiveSimilarity = Float.NEGATIVE_INFINITY;
}
@Override
public boolean collect(int docId, float similarity) {
- maxSimilarity = Math.max(maxSimilarity, similarity);
+ // Returns true / false based on whether minCompetitiveSimilarity has been
updated
if (similarity >= resultSimilarity) {
scoreDocList.add(new ScoreDoc(docId, similarity));
+ return false;
}
+
+ if (decay == 1) {
+ // maximum exploration
+ return false;
+
+ } else if (minCompetitiveSimilarity == Float.NEGATIVE_INFINITY) {
+ // start with a large buffer
+ minCompetitiveSimilarity = -Float.MAX_VALUE;
Review Comment:
This is to be able to collect results with a similarity of
`Float.NEGATIVE_INFINITY` with the query vector, which seems unlikely, but
good-to-have for completeness.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
