MinScore implementation in Linear retriever

server/src/main/java/org/elasticsearch/index/query/RankDocsQueryBuilder.java

@@ -32,11 +32,13 @@ public class RankDocsQueryBuilder extends AbstractQueryBuilder<RankDocsQueryBuil
     private final RankDoc[] rankDocs;
     private final QueryBuilder[] queryBuilders;
     private final boolean onlyRankDocs;
+    private final float minScore;
 
-    public RankDocsQueryBuilder(RankDoc[] rankDocs, QueryBuilder[] queryBuilders, boolean onlyRankDocs) {
+    public RankDocsQueryBuilder(RankDoc[] rankDocs, QueryBuilder[] queryBuilders, boolean onlyRankDocs, float minScore) {
         this.rankDocs = rankDocs;
         this.queryBuilders = queryBuilders;
         this.onlyRankDocs = onlyRankDocs;
+        this.minScore = minScore;
     }
 
     public RankDocsQueryBuilder(StreamInput in) throws IOException {
@@ -45,9 +47,11 @@ public RankDocsQueryBuilder(StreamInput in) throws IOException {
         if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_16_0)) {
             this.queryBuilders = in.readOptionalArray(c -> c.readNamedWriteable(QueryBuilder.class), QueryBuilder[]::new);
             this.onlyRankDocs = in.readBoolean();
+            this.minScore = in.readFloat();
         } else {
             this.queryBuilders = null;
             this.onlyRankDocs = false;
+            this.minScore = Float.MIN_VALUE;
         }
     }
 
@@ -70,7 +74,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
                 changed |= newQueryBuilders[i] != queryBuilders[i];
             }
             if (changed) {
-                RankDocsQueryBuilder clone = new RankDocsQueryBuilder(rankDocs, newQueryBuilders, onlyRankDocs);
+                RankDocsQueryBuilder clone = new RankDocsQueryBuilder(rankDocs, newQueryBuilders, onlyRankDocs, minScore);
                 clone.queryName(queryName());
                 return clone;
             }
@@ -88,6 +92,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_16_0)) {
             out.writeOptionalArray(StreamOutput::writeNamedWriteable, queryBuilders);
             out.writeBoolean(onlyRankDocs);
+            out.writeFloat(minScore);
         }
     }
 
@@ -115,7 +120,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
             queries = new Query[0];
             queryNames = Strings.EMPTY_ARRAY;
         }
-        return new RankDocsQuery(reader, shardRankDocs, queries, queryNames, onlyRankDocs);
+        return new RankDocsQuery(reader, shardRankDocs, queries, queryNames, onlyRankDocs, minScore);
     }
 
     @Override
@@ -135,12 +140,13 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
     protected boolean doEquals(RankDocsQueryBuilder other) {
         return Arrays.equals(rankDocs, other.rankDocs)
             && Arrays.equals(queryBuilders, other.queryBuilders)
-            && onlyRankDocs == other.onlyRankDocs;
+            && onlyRankDocs == other.onlyRankDocs
+            && minScore == other.minScore;
     }
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(Arrays.hashCode(rankDocs), Arrays.hashCode(queryBuilders), onlyRankDocs);
+        return Objects.hash(Arrays.hashCode(rankDocs), Arrays.hashCode(queryBuilders), onlyRankDocs, minScore);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java

@@ -201,7 +201,7 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
     public QueryBuilder topDocsQuery() {
         assert queryVector != null : "query vector must be materialized at this point";
         assert rankDocs != null : "rankDocs should have been materialized by now";
-        var rankDocsQuery = new RankDocsQueryBuilder(rankDocs, null, true);
+        var rankDocsQuery = new RankDocsQueryBuilder(rankDocs, null, true, Float.MIN_VALUE);
         if (preFilterQueryBuilders.isEmpty()) {
             return rankDocsQuery.queryName(retrieverName);
         }
@@ -217,7 +217,8 @@ public QueryBuilder explainQuery() {
         var rankDocsQuery = new RankDocsQueryBuilder(
             rankDocs,
             new QueryBuilder[] { new ExactKnnQueryBuilder(VectorData.fromFloats(queryVector.get()), field, similarity) },
-            true
+            false,
+            Float.MIN_VALUE
         );
         if (preFilterQueryBuilders.isEmpty()) {
             return rankDocsQuery.queryName(retrieverName);

server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java

@@ -93,7 +93,8 @@ public QueryBuilder explainQuery() {
         var explainQuery = new RankDocsQueryBuilder(
             rankDocs.get(),
             sources.stream().map(RetrieverBuilder::explainQuery).toArray(QueryBuilder[]::new),
-            true
+            true,
+            Float.MIN_VALUE
         );
         explainQuery.queryName(retrieverName());
         return explainQuery;
@@ -113,17 +114,19 @@ public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder
                 rankQuery = new RankDocsQueryBuilder(
                     rankDocResults,
                     sources.stream().map(RetrieverBuilder::topDocsQuery).toArray(QueryBuilder[]::new),
-                    false
+                    false,
+                    Float.MIN_VALUE
                 );
             } else {
                 rankQuery = new RankDocsQueryBuilder(
                     rankDocResults,
                     sources.stream().map(RetrieverBuilder::explainQuery).toArray(QueryBuilder[]::new),
-                    false
+                    false,
+                    Float.MIN_VALUE
                 );
             }
         } else {
-            rankQuery = new RankDocsQueryBuilder(rankDocResults, null, false);
+            rankQuery = new RankDocsQueryBuilder(rankDocResults, null, false, Float.MIN_VALUE);
         }
         rankQuery.queryName(retrieverName());
         // ignore prefilters of this level, they were already propagated to children

server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQuery.java

@@ -164,11 +164,7 @@ public float getMaxScore(int docId) {
                         }
 
                         @Override
-                        public float score() {
-                            // We could still end up with a valid 0 score for a RankDoc
-                            // so here we want to differentiate between this and all the tailQuery matches
-                            // that would also produce a 0 score due to filtering, by setting the score to `Float.MIN_VALUE` instead for
-                            // RankDoc matches.
+                        public float score() throws IOException {
                             return Math.max(docs[upTo].score, Float.MIN_VALUE);
                         }
 
@@ -234,6 +230,7 @@ public int hashCode() {
     // RankDocs provided. This query does not contribute to scoring, as it is set as filter when creating the weight
     private final Query tailQuery;
     private final boolean onlyRankDocs;
+    private final float minScore;
 
     /**
      * Creates a {@code RankDocsQuery} based on the provided docs.
@@ -242,8 +239,9 @@ public int hashCode() {
      * @param sources      The original queries that were used to compute the top documents
      * @param queryNames   The names (if present) of the original retrievers
      * @param onlyRankDocs Whether the query should only match the provided rank docs
+     * @param minScore     The minimum score threshold for documents to be included in total hits
      */
-    public RankDocsQuery(IndexReader reader, RankDoc[] rankDocs, Query[] sources, String[] queryNames, boolean onlyRankDocs) {
+    public RankDocsQuery(IndexReader reader, RankDoc[] rankDocs, Query[] sources, String[] queryNames, boolean onlyRankDocs, float minScore) {
         assert sources.length == queryNames.length;
         // clone to avoid side-effect after sorting
         this.docs = rankDocs.clone();
@@ -260,13 +258,15 @@ public RankDocsQuery(IndexReader reader, RankDoc[] rankDocs, Query[] sources, St
             this.tailQuery = null;
         }
         this.onlyRankDocs = onlyRankDocs;
+        this.minScore = minScore;
     }
 
     private RankDocsQuery(RankDoc[] docs, Query topQuery, Query tailQuery, boolean onlyRankDocs) {
         this.docs = docs;
         this.topQuery = topQuery;
         this.tailQuery = tailQuery;
         this.onlyRankDocs = onlyRankDocs;
+        this.minScore = Float.MIN_VALUE;
     }
 
     private static int binarySearch(RankDoc[] docs, int fromIndex, int toIndex, int key) {
@@ -346,7 +346,41 @@ public Matches matches(LeafReaderContext context, int doc) throws IOException {
 
             @Override
             public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
-                return combinedWeight.scorerSupplier(context);
+                return new ScorerSupplier() {
+                    private final ScorerSupplier supplier = combinedWeight.scorerSupplier(context);
+
+                    @Override
+                    public Scorer get(long leadCost) throws IOException {
+                        Scorer scorer = supplier.get(leadCost);
+                        return new Scorer() {
+                            @Override
+                            public DocIdSetIterator iterator() {
+                                return scorer.iterator();
+                            }
+
+                            @Override
+                            public float getMaxScore(int docId) throws IOException {
+                                return scorer.getMaxScore(docId);
+                            }
+
+                            @Override
+                            public float score() throws IOException {
+                                float score = scorer.score();
+                                return score >= minScore ? score : 0f;
+                            }
+
+                            @Override
+                            public int docID() {
+                                return scorer.docID();
+                            }
+                        };
+                    }
+
+                    @Override
+                    public long cost() {
+                        return supplier.cost();
+                    }
+                };
             }
         };
     }

server/src/test/java/org/elasticsearch/index/query/RankDocsQueryBuilderTests.java

@@ -20,6 +20,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.TopScoreDocCollectorManager;
+import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.elasticsearch.search.rank.RankDoc;
@@ -31,6 +32,7 @@
 import java.util.Random;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class RankDocsQueryBuilderTests extends AbstractQueryTestCase<RankDocsQueryBuilder> {
@@ -50,14 +52,30 @@ private RankDoc[] generateRandomRankDocs() {
     @Override
     protected RankDocsQueryBuilder doCreateTestQueryBuilder() {
         RankDoc[] rankDocs = generateRandomRankDocs();
-        return new RankDocsQueryBuilder(rankDocs, null, false);
+        return new RankDocsQueryBuilder(rankDocs, null, false, Float.MIN_VALUE);
     }
 
     @Override
     protected void doAssertLuceneQuery(RankDocsQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
-        assertTrue(query instanceof RankDocsQuery);
+        assertThat(query, instanceOf(RankDocsQuery.class));
         RankDocsQuery rankDocsQuery = (RankDocsQuery) query;
-        assertArrayEquals(queryBuilder.rankDocs(), rankDocsQuery.rankDocs());
+        assertThat(rankDocsQuery.rankDocs(), equalTo(queryBuilder.rankDocs()));
+    }
+
+    protected Query createTestQuery() throws IOException {
+        return createRandomQuery().toQuery(createSearchExecutionContext());
+    }
+
+    private RankDocsQueryBuilder createQueryBuilder() {
+        return createRandomQuery();
+    }
+
+    private RankDocsQueryBuilder createRandomQuery() {
+        RankDoc[] rankDocs = new RankDoc[randomIntBetween(1, 5)];
+        for (int i = 0; i < rankDocs.length; i++) {
+            rankDocs[i] = new RankDoc(randomInt(), randomFloat(), randomIntBetween(0, 2));
+        }
+        return new RankDocsQueryBuilder(rankDocs, null, randomBoolean(), Float.MIN_VALUE);
     }
 
     /**
@@ -151,7 +169,8 @@ public void testRankDocsQueryEarlyTerminate() throws IOException {
                         rankDocs,
                         new Query[] { NumericDocValuesField.newSlowExactQuery("active", 1) },
                         new String[1],
-                        false
+                        false,
+                        Float.MIN_VALUE
                     );
                     var topDocsManager = new TopScoreDocCollectorManager(topSize, null, totalHitsThreshold);
                     var col = searcher.search(q, topDocsManager);
@@ -172,7 +191,8 @@ public void testRankDocsQueryEarlyTerminate() throws IOException {
                         rankDocs,
                         new Query[] { NumericDocValuesField.newSlowExactQuery("active", 1) },
                         new String[1],
-                        false
+                        false,
+                        Float.MIN_VALUE
                     );
                     var topDocsManager = new TopScoreDocCollectorManager(topSize, null, Integer.MAX_VALUE);
                     var col = searcher.search(q, topDocsManager);
@@ -187,7 +207,8 @@ public void testRankDocsQueryEarlyTerminate() throws IOException {
                         rankDocs,
                         new Query[] { NumericDocValuesField.newSlowExactQuery("active", 1) },
                         new String[1],
-                        true
+                        true,
+                        Float.MIN_VALUE
                     );
                     var topDocsManager = new TopScoreDocCollectorManager(topSize, null, Integer.MAX_VALUE);
                     var col = searcher.search(q, topDocsManager);
@@ -204,7 +225,8 @@ public void testRankDocsQueryEarlyTerminate() throws IOException {
                         singleRankDoc,
                         new Query[] { NumericDocValuesField.newSlowExactQuery("active", 1) },
                         new String[1],
-                        false
+                        false,
+                        Float.MIN_VALUE
                     );
                     var topDocsManager = new TopScoreDocCollectorManager(1, null, 0);
                     var col = searcher.search(q, topDocsManager);
@@ -257,10 +279,29 @@ public void shouldThrowForNegativeScores() throws IOException {
             iw.addDocument(new Document());
             try (IndexReader reader = iw.getReader()) {
                 SearchExecutionContext context = createSearchExecutionContext(newSearcher(reader));
-                RankDocsQueryBuilder queryBuilder = new RankDocsQueryBuilder(new RankDoc[] { new RankDoc(0, -1.0f, 0) }, null, false);
+                RankDocsQueryBuilder queryBuilder = new RankDocsQueryBuilder(new RankDoc[] { new RankDoc(0, -1.0f, 0) }, null, false, Float.MIN_VALUE);
                 IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> queryBuilder.doToQuery(context));
                 assertEquals("RankDoc scores must be positive values. Missing a normalization step?", ex.getMessage());
             }
         }
     }
+
+    public void testCreateQuery() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(new Document());
+            try (IndexReader reader = iw.getReader()) {
+                RankDoc[] rankDocs = new RankDoc[] { new RankDoc(0, randomFloat(), 0) };
+                RankDocsQuery q = new RankDocsQuery(
+                    reader,
+                    rankDocs,
+                    new Query[] { new MatchAllDocsQuery() },
+                    new String[] { "test" },
+                    false,
+                    Float.MIN_VALUE
+                );
+                assertNotNull(q);
+                assertArrayEquals(rankDocs, q.rankDocs());
+            }
+        }
+    }
 }

server/src/test/java/org/elasticsearch/search/rank/AbstractRankDocWireSerializingTestCase.java

@@ -50,7 +50,7 @@ public void testRankDocSerialization() throws IOException {
         for (int i = 0; i < totalDocs; i++) {
             docs.add(createTestRankDoc());
         }
-        RankDocsQueryBuilder rankDocsQueryBuilder = new RankDocsQueryBuilder(docs.toArray((T[]) new RankDoc[0]), null, randomBoolean());
+        RankDocsQueryBuilder rankDocsQueryBuilder = new RankDocsQueryBuilder(docs.toArray((T[]) new RankDoc[0]), null, randomBoolean(), Float.MIN_VALUE);
         RankDocsQueryBuilder copy = (RankDocsQueryBuilder) copyNamedWriteable(rankDocsQueryBuilder, writableRegistry(), QueryBuilder.class);
         assertThat(rankDocsQueryBuilder, equalTo(copy));
     }