Adding RankDocsRetrieverBuilder and RankDocsQuery (elastic#111709)

Author: pmpailis

server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RankDocRetrieverBuilderIT.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.plugins.Plugin;
@@ -141,6 +142,11 @@ private AssertingRetrieverBuilder(RetrieverBuilder innerRetriever) {
             this.innerRetriever = innerRetriever;
         }
 
+        @Override
+        public QueryBuilder topDocsQuery() {
+            return null;
+        }
+
         @Override
         public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
             assertNull(ctx.getPointInTimeBuilder());
@@ -200,6 +206,11 @@ public boolean isCompound() {
             return true;
         }
 
+        @Override
+        public QueryBuilder topDocsQuery() {
+            return null;
+        }
+
         @Override
         public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
             assertNotNull(ctx.getPointInTimeBuilder());

server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java

@@ -364,6 +364,7 @@
     exports org.elasticsearch.search.rank.rerank;
     exports org.elasticsearch.search.rescore;
     exports org.elasticsearch.search.retriever;
+    exports org.elasticsearch.search.retriever.rankdoc;
     exports org.elasticsearch.search.runtime;
     exports org.elasticsearch.search.searchafter;
     exports org.elasticsearch.search.slice;

server/src/main/java/module-info.java

@@ -196,6 +196,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ZDT_NANOS_SUPPORT = def(8_726_00_0);
     public static final TransportVersion LTR_SERVERLESS_RELEASE = def(8_727_00_0);
     public static final TransportVersion ALLOW_PARTIAL_SEARCH_RESULTS_IN_PIT = def(8_728_00_0);
+    public static final TransportVersion RANK_DOCS_RETRIEVER = def(8_729_00_0);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -73,6 +73,7 @@
 import org.elasticsearch.index.analysis.NamedAnalyzer;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.lucene.grouping.TopFieldGroups;
+import org.elasticsearch.search.retriever.rankdoc.RankDocsSortField;
 import org.elasticsearch.search.sort.ShardDocSortField;
 
 import java.io.IOException;
@@ -551,6 +552,8 @@ private static SortField rewriteMergeSortField(SortField sortField) {
             return newSortField;
         } else if (sortField.getClass() == ShardDocSortField.class) {
             return new SortField(sortField.getField(), SortField.Type.LONG, sortField.getReverse());
+        } else if (sortField.getClass() == RankDocsSortField.class) {
+            return new SortField(sortField.getField(), SortField.Type.INT, sortField.getReverse());
         } else {
             return sortField;
         }

server/src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -32,6 +32,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.search.ESToParentBlockJoinQuery;
+import org.elasticsearch.search.retriever.rankdoc.RankDocsQuery;
 import org.elasticsearch.search.runtime.AbstractScriptFieldQuery;
 import org.elasticsearch.search.vectors.KnnScoreDocQuery;
 
@@ -255,10 +256,10 @@ public void visitLeaf(Query leafQuery) {
                     hasUnknownLeaf[0] = true;
                 }
                 /**
-                 * KnnScoreDocQuery requires the same reader that built the docs
+                 * KnnScoreDocQuery and RankDocsQuery requires the same reader that built the docs
                  * When using {@link HighlightFlag#WEIGHT_MATCHES} different readers are used and isn't supported by this query
                  */
-                if (leafQuery instanceof KnnScoreDocQuery) {
+                if (leafQuery instanceof KnnScoreDocQuery || leafQuery instanceof RankDocsQuery) {
                     hasUnknownLeaf[0] = true;
                 }
                 super.visitLeaf(query);

server/src/main/java/org/elasticsearch/lucene/search/uhighlight/CustomUnifiedHighlighter.java

@@ -237,6 +237,8 @@
 import org.elasticsearch.search.retriever.RetrieverBuilder;
 import org.elasticsearch.search.retriever.RetrieverParserContext;
 import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
+import org.elasticsearch.search.retriever.rankdoc.RankDocsQueryBuilder;
+import org.elasticsearch.search.retriever.rankdoc.RankDocsSortBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.sort.GeoDistanceSortBuilder;
 import org.elasticsearch.search.sort.ScoreSortBuilder;
@@ -851,17 +853,20 @@ private void registerRescorer(RescorerSpec<?> spec) {
     }
 
     private void registerRankers() {
+        namedWriteables.add(new NamedWriteableRegistry.Entry(RankDoc.class, RankDoc.NAME, RankDoc::new));
         namedWriteables.add(new NamedWriteableRegistry.Entry(RankDoc.class, RankFeatureDoc.NAME, RankFeatureDoc::new));
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(RankShardResult.class, RankFeatureShardResult.NAME, RankFeatureShardResult::new)
         );
+        namedWriteables.add(new NamedWriteableRegistry.Entry(QueryBuilder.class, RankDocsQueryBuilder.NAME, RankDocsQueryBuilder::new));
     }
 
     private void registerSorts() {
         namedWriteables.add(new NamedWriteableRegistry.Entry(SortBuilder.class, GeoDistanceSortBuilder.NAME, GeoDistanceSortBuilder::new));
         namedWriteables.add(new NamedWriteableRegistry.Entry(SortBuilder.class, ScoreSortBuilder.NAME, ScoreSortBuilder::new));
         namedWriteables.add(new NamedWriteableRegistry.Entry(SortBuilder.class, ScriptSortBuilder.NAME, ScriptSortBuilder::new));
         namedWriteables.add(new NamedWriteableRegistry.Entry(SortBuilder.class, FieldSortBuilder.NAME, FieldSortBuilder::new));
+        namedWriteables.add(new NamedWriteableRegistry.Entry(SortBuilder.class, RankDocsSortBuilder.NAME, RankDocsSortBuilder::new));
     }
 
     private static <T> void registerFromPlugin(List<SearchPlugin> plugins, Function<SearchPlugin, List<T>> producer, Consumer<T> consumer) {

server/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -2228,15 +2228,15 @@ public ActionRequestValidationException validate(
             if (sorts() != null) {
                 specified.add(SORT_FIELD.getPreferredName());
             }
-            if (rescores() != null) {
-                specified.add(RESCORE_FIELD.getPreferredName());
-            }
             if (minScore() != null) {
                 specified.add(MIN_SCORE_FIELD.getPreferredName());
             }
             if (rankBuilder() != null) {
                 specified.add(RANK_FIELD.getPreferredName());
             }
+            if (rescores() != null) {
+                specified.add(RESCORE_FIELD.getPreferredName());
+            }
             if (specified.isEmpty() == false) {
                 validationException = addValidationError(
                     "cannot specify [" + RETRIEVER.getPreferredName() + "] and " + specified,

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -8,29 +8,37 @@
 
 package org.elasticsearch.search.rank;
 
+import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.ScoreDoc;
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xcontent.ToXContentFragment;
+import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.Objects;
 
 /**
  * {@code RankDoc} is the base class for all ranked results.
- * Subclasses should extend this with additional information
- * required for their global ranking method.
+ * Subclasses should extend this with additional information required for their global ranking method.
  */
-public abstract class RankDoc extends ScoreDoc implements NamedWriteable {
+public class RankDoc extends ScoreDoc implements NamedWriteable, ToXContentFragment {
+
+    public static final String NAME = "rank_doc";
 
     public static final int NO_RANK = -1;
 
     /**
-     * If this document has been ranked, this is its final
-     * rrf ranking from all the result sets.
+     * If this document has been ranked, this is its final rrf ranking from all the result sets.
      */
     public int rank = NO_RANK;
 
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
     public record RankKey(int doc, int shardIndex) {}
 
     public RankDoc(int doc, float score, int shardIndex) {
@@ -51,7 +59,26 @@ public final void writeTo(StreamOutput out) throws IOException {
         doWriteTo(out);
     }
 
-    protected abstract void doWriteTo(StreamOutput out) throws IOException;
+    protected void doWriteTo(StreamOutput out) throws IOException {};
+
+    /**
+     * Explain the ranking of this document.
+     */
+    public Explanation explain() {
+        return Explanation.match(rank, "doc [" + doc + "] with an original score of [" + score + "] is at rank [" + rank + "].");
+    }
+
+    @Override
+    public final XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field("_rank", rank);
+        builder.field("_doc", doc);
+        builder.field("_shard", shardIndex);
+        builder.field("_score", score);
+        doToXContent(builder, params);
+        return builder;
+    }
+
+    protected void doToXContent(XContentBuilder builder, Params params) throws IOException {}
 
     @Override
     public final boolean equals(Object o) {
@@ -61,17 +88,21 @@ public final boolean equals(Object o) {
         return doc == rd.doc && score == rd.score && shardIndex == rd.shardIndex && rank == rd.rank && doEquals(rd);
     }
 
-    protected abstract boolean doEquals(RankDoc rd);
+    protected boolean doEquals(RankDoc rd) {
+        return true;
+    }
 
     @Override
     public final int hashCode() {
         return Objects.hash(doc, score, shardIndex, doHashCode());
     }
 
-    protected abstract int doHashCode();
+    protected int doHashCode() {
+        return 0;
+    }
 
     @Override
     public String toString() {
-        return "RankDoc{" + "score=" + score + ", doc=" + doc + ", shardIndex=" + shardIndex + '}';
+        return "RankDoc{" + "_rank=" + rank + ", _doc=" + doc + ", _shard=" + shardIndex + ", _score=" + score + '}';
     }
 }

server/src/main/java/org/elasticsearch/search/rank/RankDoc.java

@@ -8,9 +8,11 @@
 
 package org.elasticsearch.search.rank.feature;
 
+import org.apache.lucene.search.Explanation;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.search.rank.RankDoc;
+import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.Objects;
@@ -22,7 +24,7 @@ public class RankFeatureDoc extends RankDoc {
 
     public static final String NAME = "rank_feature_doc";
 
-    // todo: update to support more than 1 fields; and not restrict to string data
+    // TODO: update to support more than 1 fields; and not restrict to string data
     public String featureData;
 
     public RankFeatureDoc(int doc, float score, int shardIndex) {
@@ -34,6 +36,11 @@ public RankFeatureDoc(StreamInput in) throws IOException {
         featureData = in.readOptionalString();
     }
 
+    @Override
+    public Explanation explain() {
+        throw new UnsupportedOperationException("explain is not supported for {" + getClass() + "}");
+    }
+
     public void featureData(String featureData) {
         this.featureData = featureData;
     }
@@ -58,4 +65,9 @@ protected int doHashCode() {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field("featureData", featureData);
+    }
 }