[lucene] branch main updated: Allow implementers of AbstractKnnVectorQuery to access final topK results (#12590)

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

benwtrent pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git


The following commit(s) were added to refs/heads/main by this push:
     new 2474940bffe Allow implementers of AbstractKnnVectorQuery to access final topK results (#12590)
2474940bffe is described below

commit 2474940bffe6118ed31ceb717fd49705d819e1fc
Author: Kaival Parikh <46...@users.noreply.github.com>
AuthorDate: Fri Oct 6 00:44:59 2023 +0530

    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 <ka...@amazon.com>
---
 .../apache/lucene/search/AbstractKnnVectorQuery.java   | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
index d81ab2806fe..f0de3ea7a39 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -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;