vigyasharma commented on code in PR #15171:
URL: https://github.com/apache/lucene/pull/15171#discussion_r2335602427
##########
lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java:
##########
@@ -309,21 +309,24 @@ protected TopDocs exactSearch(
TotalHits.Relation relation = TotalHits.Relation.EQUAL_TO;
ScoreDoc topDoc = queue.top();
DocIdSetIterator vectorIterator = vectorScorer.iterator();
- DocIdSetIterator conjunction =
- ConjunctionDISI.createConjunction(List.of(vectorIterator,
acceptIterator), List.of());
- int doc;
- while ((doc = conjunction.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
+ DocAndFloatFeatureBuffer buffer = new DocAndFloatFeatureBuffer();
+ VectorScorer.Bulk bulkScorer = vectorScorer.bulk(acceptIterator);
+ while (vectorIterator.docID() != DocIdSetIterator.NO_MORE_DOCS) {
// Mark results as partial if timeout is met
if (queryTimeout != null && queryTimeout.shouldExit()) {
relation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
break;
}
- assert vectorIterator.docID() == doc;
- float score = vectorScorer.score();
- if (score > topDoc.score) {
- topDoc.score = score;
- topDoc.doc = doc;
- topDoc = queue.updateTop();
+ // iterator already takes live docs into account
+ bulkScorer.nextDocsAndScores(64, null, buffer);
Review Comment:
nit: move to a final constant?
##########
lucene/core/src/java/org/apache/lucene/search/VectorScorer.java:
##########
@@ -38,4 +40,37 @@ public interface VectorScorer {
* @return a {@link DocIdSetIterator} over the documents.
*/
DocIdSetIterator iterator();
+
+ /**
+ * An optional bulk scorer implementation that allows bulk scoring over the
provided matching docs
+ */
+ default Bulk bulk(DocIdSetIterator matchingDocs) throws IOException {
+ final DocIdSetIterator iterator =
+ ConjunctionUtils.createConjunction(List.of(matchingDocs, iterator()),
List.of());
+ if (iterator.docID() == -1) {
+ iterator.nextDoc();
+ }
+ return (nextCount, liveDocs, buffer) -> {
+ buffer.growNoCopy(nextCount);
+ int size = 0;
+ for (int doc = iterator.docID();
+ doc != DocIdSetIterator.NO_MORE_DOCS && size < nextCount;
+ doc = iterator.nextDoc()) {
+ if (liveDocs == null || liveDocs.get(doc)) {
+ buffer.docs[size] = doc;
+ buffer.features[size] = score();
+ ++size;
+ }
+ }
+ buffer.size = size;
+ if (liveDocs != null) {
+ buffer.apply(liveDocs);
+ }
+ };
+ }
+
+ interface Bulk {
+ void nextDocsAndScores(int nextCount, Bits liveDocs,
DocAndFloatFeatureBuffer buffer)
Review Comment:
Let's add a docstring to describe what these params do, and that the final
doc id and score will be stored in the `buffer`?
--
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]
