Apache
/
lucene
mirror of https://github.com/apache/lucene.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Allow implementers of AbstractKnnVectorQuery to access final topK results (#12590) 

* Make AbstractKnnVectorQuery#createRewrittenQuery protected

* Add explicit method to merge per-segment results

* Add javadocs explaining the new method

* Run gradlew tidy

---------

Co-authored-by: Kaival Parikh <kaivalnp@amazon.com>
This commit is contained in:
Kaival Parikh 2023-10-06 00:44:59 +05:30 committed by GitHub
parent 91019d685a
commit 2474940bff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
View File

@ -88,7 +88,7 @@ abstract class AbstractKnnVectorQuery extends Query {
TopDocs[] perLeafResults = taskExecutor.invokeAll(tasks).toArray(TopDocs[]::new);
// Merge sort the results
TopDocs topK = TopDocs.merge(k, perLeafResults);
TopDocs topK = mergeLeafResults(perLeafResults);
if (topK.scoreDocs.length == 0) {
return new MatchNoDocsQuery();
}
@ -200,6 +200,22 @@ abstract class AbstractKnnVectorQuery extends Query {
return new TopDocs(totalHits, topScoreDocs);
}
/**
* Merges all segment-level kNN results to get the index-level kNN results.
*
* <p>The default implementation delegates to {@link TopDocs#merge(int, TopDocs[])} to find the
* overall top {@link #k}, which requires input results to be sorted.
*
* <p>This method is useful for reading and / or modifying the final results as needed.
*
* @param perLeafResults array of segment-level kNN results.
* @return index-level kNN results (no constraint on their ordering).
* @lucene.experimental
*/
protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
return TopDocs.merge(k, perLeafResults);
}
private Query createRewrittenQuery(IndexReader reader, TopDocs topK) {
int len = topK.scoreDocs.length;